jax_fem.fe module

class jax_fem.fe.FiniteElement(mesh, vec, dim, ele_type, gauss_order, dirichlet_bc_info, periodic_bc_info=None)

Bases: object

Defines finite element related to one variable (can be vector valued)

Parameters:

• mesh (Mesh)

• vec (int)

• dim (int)

• ele_type (str)

• gauss_order (int)

• dirichlet_bc_info (List[List[Callable] | List[int]] | None)

• periodic_bc_info (List[List[Callable] | List[int]] | None)

mesh

The mesh object stores points (coordinates) and cells (connectivity).

Type:

Mesh object

vec

The number of vector variable components of the solution. E.g., a 3D displacement field has u_x, u_y and u_z components, so vec=3

Type:

int

dim

The dimension of the problem.

Type:

int

ele_type

Element type

Type:

str

dirichlet_bc_info

location_fnsList[Callable]

Callable : a function that inputs a point and returns if the point satisfies the location condition

vecs: List[int]

integer value must be in the range of 0 to vec - 1, specifying which component of the (vector) variable to apply Dirichlet condition to

value_fnsList[Callable]

Callable : a function that inputs a point and returns the Dirichlet value

Type:

[location_fns, vecs, value_fns]

periodic_bc_info

location_fns_AList[Callable]

Callable : location function for boundary A

location_fns_BList[Callable]

Callable : location function for boundary B

mappingsList[Callable]

Callable: function mapping a point from boundary A to boundary B

vecs: List[int]

which component of the (vector) variable to apply periodic condition to

Type:

[location_fns_A, location_fns_B, mappings, vecs]

Dirichlet_boundary_conditions(dirichlet_bc_info)

Indices and values for Dirichlet B.C.

Parameters:

dirichlet_bc_info ([location_fns, vecs, value_fns])

Returns:

• node_inds_List (List[onp.ndarray]) – The ndarray ranges from 0 to num_total_nodes - 1

• vec_inds_List (List[onp.ndarray]) – The ndarray ranges from 0 to to vec - 1

• vals_List (List[ndarray]) – Dirichlet values to be assigned

convert_from_dof_to_face_quad(sol, boundary_inds)

Obtain surface solution from nodal solution

Parameters:

• sol (np.DeviceArray) – (num_total_nodes, vec)

• boundary_inds (int)

Returns:

u – (num_selected_faces, num_face_quads, vec)

Return type:

np.DeviceArray

convert_from_dof_to_quad(sol)

Obtain quad values from nodal solution

Parameters:

sol (np.DeviceArray) – (num_total_nodes, vec)

Returns:

u – (num_cells, num_quads, vec)

Return type:

np.DeviceArray

dim: int

dirichlet_bc_info: List[List[Callable] | List[int]] | None

ele_type: str

gauss_order: int

get_boundary_conditions_inds(location_fns)

Given location functions, compute which faces satisfy the condition.

Parameters:

location_fns (List[Callable]) –

Callable: a location function that inputs a point (ndarray) and returns if the point satisfies the location condition

e.g., lambda x: np.isclose(x[0], 0.) If this location function takes 2 arguments, then the first is point and the second is index. e.g., lambda x, ind: np.isclose(x[0], 0.) & np.isin(ind, np.array([1, 3, 10]))

Returns:

boundary_inds_list – (num_selected_faces, 2) boundary_inds_list[k][i, 0] returns the global cell index of the ith selected face of boundary subset k boundary_inds_list[k][i, 1] returns the local face index of the ith selected face of boundary subset k

Return type:

List[onp.ndarray]

get_face_shape_grads(boundary_inds)

Face shape function gradients and JxW (for surface integral) Nanson’s formula is used to map physical surface ingetral to reference domain Reference: https://en.wikiversity.org/wiki/Continuum_mechanics/Volume_change_and_area_change

Parameters:

boundary_inds (List[onp.ndarray]) – (num_selected_faces, 2)

Returns:

• face_shape_grads_physical (onp.ndarray) – (num_selected_faces, num_face_quads, num_nodes, dim)

• nanson_scale (onp.ndarray) – (num_selected_faces, num_face_quads)

get_physical_quad_points()

Compute physical quadrature points

Returns:

physical_quad_points – (num_cells, num_quads, dim)

Return type:

onp.ndarray

get_physical_surface_quad_points(boundary_inds)

Compute physical quadrature points on the surface

Parameters:

boundary_inds (List[onp.ndarray]) – ndarray shape: (num_selected_faces, 2)

Returns:

physical_surface_quad_points – (num_selected_faces, num_face_quads, dim)

Return type:

ndarray

get_shape_grads()

Compute shape function gradient value The gradient is w.r.t physical coordinates. See Hughes, Thomas JR. The finite element method: linear static and dynamic finite element analysis. Courier Corporation, 2012. Page 147, Eq. (3.9.3)

Returns:

• shape_grads_physical (onp.ndarray) – (num_cells, num_quads, num_nodes, dim)

• JxW (onp.ndarray) – (num_cells, num_quads)

mesh: Mesh

periodic_bc_info: List[List[Callable] | List[int]] | None = None

print_BC_info()

Print boundary condition information for debugging purposes.

TODO: Not working

sol_to_grad(sol)

Obtain solution gradient from nodal solution

Parameters:

sol (np.DeviceArray) – (num_total_nodes, vec)

Returns:

u_grads – (num_cells, num_quads, vec, dim)

Return type:

np.DeviceArray

update_Dirichlet_boundary_conditions(dirichlet_bc_info)

Reset Dirichlet boundary conditions. Useful when a time-dependent problem is solved, and at each iteration the boundary condition needs to be updated.

Parameters:

dirichlet_bc_info ([location_fns, vecs, value_fns])

vec: int

class jax_fem.fe.Mesh(points, cells, ele_type='TET4')

Bases: object

A custom mesh manager might be better using a third-party library like meshio.

count_selected_faces(location_fn)

Given location functions, compute the count of faces that satisfy the location function. Useful for setting up distributed load conditions.

Parameters:

location_fns (List[Callable]) – Callable: a function that inputs a point and returns a boolean value describing whether the boundary condition should be applied.

Returns:

face_count

Return type:

int

jax_fem.fe.dataclass(cls=None, /, *, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False)

Returns the same class as was passed in, with dunder methods added based on the fields defined in the class.

Examines PEP 526 __annotations__ to determine fields.

If init is true, an __init__() method is added to the class. If repr is true, a __repr__() method is added. If order is true, rich comparison dunder methods are added. If unsafe_hash is true, a __hash__() method function is added. If frozen is true, fields may not be assigned to after instance creation. If match_args is true, the __match_args__ tuple is added. If kw_only is true, then by default all fields are keyword-only. If slots is true, an __slots__ attribute is added.

jax_fem.fe.get_face_shape_vals_and_grads(ele_type, gauss_order=None)

TODO: Add comments

Returns:

• face_shape_vals (ndarray) – (6, 4, 8) = (num_faces, num_face_quads, num_nodes)

• face_shape_grads_ref (ndarray) – (6, 4, 3) = (num_faces, num_face_quads, num_nodes, dim)

• face_weights (ndarray) – (6, 4) = (num_faces, num_face_quads)

• face_normals (ndarray) – (6, 3) = (num_faces, dim)

• face_inds (ndarray) – (6, 4) = (num_faces, num_face_vertices)

jax_fem.fe.get_shape_vals_and_grads(ele_type, gauss_order=None)

TODO: Add comments

Returns:

• shape_values (ndarray) – (8, 8) = (num_quads, num_nodes)

• shape_grads_ref (ndarray) – (8, 8, 3) = (num_quads, num_nodes, dim)

• weights (ndarray) – (8,) = (num_quads,)