jax_fem.fe module

jax_fem.fe module#

class jax_fem.fe.FiniteElement(mesh, vec, dim, ele_type, gauss_order, dirichlet_bc_info)[source]#

Finite element class for one variable. This variable can be:

  • Scalar-valued (when vec = 1)

  • Vector-valued (when vec > 1)

Parameters:

  • mesh (Mesh)

  • vec (int)

  • dim (int)

  • ele_type (str)

  • gauss_order (int)

  • dirichlet_bc_info (list)

mesh#

Stores points (coordinates) and cells (connectivity).

Type:

Mesh

vec#

Number of vector components in solution (e.g., 3 for displacement in 3D, 1 for temperature).

Type:

int

dim#

Spatial dimension of the problem. Currently supports:

  • 2 (2D problems)

  • 3 (3D problems)

Type:

int

ele_type#

Element type. Currently supports:

  • ‘QUAD4’

  • ‘QUAD8’

  • ‘TRI3’

  • ‘TRI6’

  • ‘HEX8’

  • ‘HEX20’

  • ‘HEX27’

  • ‘TET4’

  • ‘TET10’

Type:

str

gauss_order#

Order of Gaussian quadrature.

Type:

int

dirichlet_bc_info#

A list for Dirichlet boundary condition information, whose elements are structured as:

  • location_fns: list of callables Each callable takes a point (NumpyArray) and returns a boolean indicating if the point satisfies the location condition

  • vecs: list of integers Each integer must be in the range of 0 to vec - 1, specifying which component of the (vector) variable to apply Dirichlet condition to

  • value_fns: list of callables Each callable takes a point and returns the Dirichlet value

Type:

list

get_shape_grads()[source]#

Compute shape function gradient value.

The gradient is w.r.t physical coordinates. Refer to 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 (NumpyArray) – Shape is (num_cells, num_quads, num_nodes, dim).

  • JxW (NumpyArray) – Shape is (num_cells, num_quads).

get_face_shape_grads(boundary_inds)[source]#

Face shape function gradients and JxW (for surface integral). Nanson’s formula is used to map physical surface ingetral to reference domain. Refer to wikiversity.

Parameters:

boundary_inds (list[NumpyArray]) – Shape is (num_selected_faces, 2).

Returns:

  • face_shape_grads_physical (NumpyArray) – Shape is (num_selected_faces, num_face_quads, num_nodes, dim).

  • nanson_scale (NumpyArray) – Shape is (num_selected_faces, num_face_quads).

get_physical_quad_points()[source]#

Compute physical quadrature points

Returns:

physical_quad_points – Shape is (num_cells, num_quads, dim).

Return type:

NumpyArray

get_physical_surface_quad_points(boundary_inds)[source]#

Compute physical quadrature points on the surface

Parameters:

boundary_inds (list[NumpyArray]) – Shape for NumpyArray is (num_selected_faces, 2).

Returns:

physical_surface_quad_points – Shape for NumpyArray is (num_selected_faces, num_face_quads, dim).

Return type:

NumpyArray

Dirichlet_boundary_conditions(dirichlet_bc_info)[source]#

Indices and values for Dirichlet B.C.

Parameters:

dirichlet_bc_info (list) – [location_fns, vecs, value_fns]

Returns:

  • node_inds_list (list[NumpyArray]) – The value of the NumpyArray ranges from 0 to num_total_nodes - 1.

  • vec_inds_list (list[NumpyArray]) – The value of the NumpyArray ranges from 0 to to vec - 1.

  • vals_list (list[NumpyArray]) – Dirichlet values to be assigned.

update_Dirichlet_boundary_conditions(dirichlet_bc_info)[source]#

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 (list) – [location_fns, vecs, value_fns]

get_boundary_conditions_inds(location_fns)[source]#

Given location functions, compute which faces satisfy the condition.

Parameters:

location_fns (list[callable]) –

The callable is a location function that inputs a point (NumpyArray) and returns if the point satisfies the location condition. For example,

lambda x: np.isclose(x[0], 0.)

If this location function takes 2 arguments, then the first is point and the second is index. For example,

lambda x, ind: np.isclose(x[0], 0.) & np.isin(ind, np.array([1, 3, 10]))

Returns:

boundary_inds_list – Shape of NumpyArray is (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[NumpyArray]

convert_from_dof_to_quad(sol)[source]#

Obtain quad values from nodal solution.

Parameters:

sol (JaxArray) – Shape is (num_total_nodes, vec).

Returns:

u – Shape is (num_cells, num_quads, vec).

Return type:

JaxArray

convert_from_dof_to_face_quad(sol, boundary_inds)[source]#

Obtain surface solution from nodal solution

Parameters:

  • sol (JaxArray) – Shape is (num_total_nodes, vec).

  • boundary_inds (int)

Returns:

u – Shape is (num_selected_faces, num_face_quads, vec).

Return type:

JaxArray

sol_to_grad(sol)[source]#

Obtain solution gradient from nodal solution.

Parameters:

sol (JaxArray) – Shape is (num_total_nodes, vec).

Returns:

u_grads – Shape is (num_cells, num_quads, vec, dim).

Return type:

JaxArray

Contents