zoomy_jax.mesh.mesh module

zoomy_jax.mesh.mesh module#

JAX mesh container and conversion utilities.

MeshJAX is a frozen dataclass holding all mesh arrays as JAX arrays, registered as a JAX pytree for use inside jit/vmap/lax.while_loop.

zoomy_jax.mesh.mesh.compute_derivatives(u, mesh, derivatives_multi_index=None)#

Compute cell-wise LSQ derivatives using JAX vmap.

class zoomy_jax.mesh.mesh.MeshJAX(dimension, type, n_cells, n_inner_cells, n_faces, n_vertices, n_boundary_faces, n_faces_per_cell, vertex_coordinates, cell_vertices, cell_faces, cell_volumes, cell_centers, cell_inradius, cell_neighbors, boundary_face_cells, boundary_face_ghosts, boundary_face_function_numbers, boundary_face_physical_tags, boundary_face_face_indices, face_cells, face_normals, face_volumes, face_centers, face_subvolumes, face_neighbors, boundary_conditions_sorted_physical_tags, boundary_conditions_sorted_names, lsq_gradQ, lsq_neighbors, lsq_monomial_multi_index, lsq_scale_factors, z_ordering)#

Bases: object

Immutable JAX mesh container. All numeric fields are jnp.ndarray.

Parameters:

  • dimension (int) –

  • type (str) –

  • n_cells (int) –

  • n_inner_cells (int) –

  • n_faces (int) –

  • n_vertices (int) –

  • n_boundary_faces (int) –

  • n_faces_per_cell (int) –

  • vertex_coordinates (Array) –

  • cell_vertices (Array) –

  • cell_faces (Array) –

  • cell_volumes (Array) –

  • cell_centers (Array) –

  • cell_inradius (Array) –

  • cell_neighbors (Array) –

  • boundary_face_cells (Array) –

  • boundary_face_ghosts (Array) –

  • boundary_face_function_numbers (Array) –

  • boundary_face_physical_tags (Array) –

  • boundary_face_face_indices (Array) –

  • face_cells (Array) –

  • face_normals (Array) –

  • face_volumes (Array) –

  • face_centers (Array) –

  • face_subvolumes (Array) –

  • face_neighbors (Array) –

  • boundary_conditions_sorted_physical_tags (Array) –

  • boundary_conditions_sorted_names (Any) –

  • lsq_gradQ (Array) –

  • lsq_neighbors (Array) –

  • lsq_monomial_multi_index (Any) –

  • lsq_scale_factors (Any) –

  • z_ordering (Array) –

dimension: int#
type: str#
n_cells: int#
n_inner_cells: int#
n_faces: int#
n_vertices: int#
n_boundary_faces: int#
n_faces_per_cell: int#
vertex_coordinates: Array#
cell_vertices: Array#
cell_faces: Array#
cell_volumes: Array#
cell_centers: Array#
cell_inradius: Array#
cell_neighbors: Array#
boundary_face_cells: Array#
boundary_face_ghosts: Array#
boundary_face_function_numbers: Array#
boundary_face_physical_tags: Array#
boundary_face_face_indices: Array#
face_cells: Array#
face_normals: Array#
face_volumes: Array#
face_centers: Array#
face_subvolumes: Array#
face_neighbors: Array#
boundary_conditions_sorted_physical_tags: Array#
boundary_conditions_sorted_names: Any#
lsq_gradQ: Array#
lsq_neighbors: Array#
lsq_monomial_multi_index: Any#
lsq_scale_factors: Any#
z_ordering: Array#
zoomy_jax.mesh.mesh.convert_mesh_to_jax(mesh)#

Convert a mesh (old Mesh or new LSQMesh/FVMMesh) to MeshJAX.

Accepts both the old monolithic Mesh class and the new hierarchy (LSQMesh, which exposes cached geometry via properties).

Return type:

MeshJAX