fluxfem 0.1.3a0__py3-none-any.whl

fluxfem/mesh/io.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+try:
+    import meshio
+except Exception as e:  # pragma: no cover
+    meshio = None
+    meshio_import_error = e
+else:
+    meshio_import_error = None
+
+from .hex import HexMesh
+from .tet import TetMesh
+from .surface import SurfaceMesh
+
+
+def load_gmsh_mesh(path: str):
+    """
+    Load a Gmsh .msh (v2/v4) file containing hex or tet elements (and optional boundary facets).
+
+    Returns:
+      mesh: HexMesh or TetMesh
+      facets: (n_facets, 3 or 4) or None
+      facet_tags: (n_facets,) or None  (gmsh physical tags if present)
+    """
+    if meshio is None:
+        raise ImportError(f"meshio is required to load gmsh meshes: {meshio_import_error}")
+
+    msh = meshio.read(path)
+    coords = np.asarray(msh.points[:, :3], dtype=DTYPE)
+
+    mesh = None
+    if "hexahedron" in msh.cells_dict:
+        conn = np.asarray(msh.cells_dict["hexahedron"], dtype=np.int32)
+        mesh = HexMesh(jnp.asarray(coords), jnp.asarray(conn))
+    elif "tetra" in msh.cells_dict:
+        conn = np.asarray(msh.cells_dict["tetra"], dtype=np.int32)
+        mesh = TetMesh(jnp.asarray(coords), jnp.asarray(conn))
+    else:
+        raise ValueError("gmsh mesh does not contain hexahedron or tetra cells")
+
+    facets = None
+    facet_tags = None
+    # surface facets (quad/tri)
+    cell_dict = msh.cells_dict
+    for fkey in ("quad", "triangle", "tri"):
+        if fkey in cell_dict:
+            facets = np.asarray(cell_dict[fkey], dtype=np.int32)
+            break
+    tags_raw = None
+    if facets is not None:
+        if "gmsh:physical" in msh.cell_data_dict:
+            tags_raw = msh.cell_data_dict["gmsh:physical"].get(fkey, None)
+        else:
+            if "gmsh:physical" in msh.cell_data:
+                for block, data in zip(msh.cells, msh.cell_data["gmsh:physical"]):
+                    if block.type == fkey:
+                        tags_raw = data
+                        break
+        if tags_raw is not None:
+            facet_tags = np.asarray(tags_raw, dtype=np.int32)
+
+    return mesh, facets, facet_tags
+
+
+def load_gmsh_hex_mesh(path: str):
+    mesh, facets, tags = load_gmsh_mesh(path)
+    if not isinstance(mesh, HexMesh):
+        raise ValueError("gmsh mesh is not hexahedral")
+    return mesh, facets, tags
+
+
+def load_gmsh_tet_mesh(path: str):
+    mesh, facets, tags = load_gmsh_mesh(path)
+    if not isinstance(mesh, TetMesh):
+        raise ValueError("gmsh mesh is not tetrahedral")
+    return mesh, facets, tags
+
+
+def make_surface_from_facets(coords: np.ndarray, facets: np.ndarray, tags=None) -> SurfaceMesh:
+    """Helper to build SurfaceMesh from raw coords/facets (optional tags)."""
+    return SurfaceMesh.from_facets(coords, facets, facet_tags=tags)

fluxfem/mesh/predicate.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def plane_predicate(axis: int, value: float, tol: float = 1e-8):
+    """Return predicate True when all nodes lie on plane x[axis]=value (within tol)."""
+    def pred(pts: np.ndarray) -> bool:
+        return bool(np.allclose(pts[:, axis], value, atol=tol))
+    return pred
+
+
+def axis_plane_predicate(axis: int, value: float, tol: float = 1e-8):
+    """Alias of plane_predicate for readability."""
+    return plane_predicate(axis, value, tol=tol)
+
+
+def slab_predicate(axis: int, min_val: float, max_val: float, tol: float = 1e-8):
+    """Return predicate True when nodes lie in a slab min<=x[axis]<=max (within tol)."""
+    def pred(pts: np.ndarray) -> np.ndarray:
+        pts_np = np.asarray(pts, dtype=float)
+        return (pts_np[:, axis] >= min_val - tol) & (pts_np[:, axis] <= max_val + tol)
+    return pred
+
+
+def bbox_predicate(mins: np.ndarray, maxs: np.ndarray, tol: float = 1e-8):
+    """Return predicate True for nodes on the axis-aligned bounding box."""
+    mins_np = np.asarray(mins, dtype=float)
+    maxs_np = np.asarray(maxs, dtype=float)
+
+    def pred(pts: np.ndarray) -> np.ndarray:
+        pts_np = np.asarray(pts, dtype=float)
+        on_min = np.isclose(pts_np, mins_np[None, :], atol=tol)
+        on_max = np.isclose(pts_np, maxs_np[None, :], atol=tol)
+        return np.any(on_min | on_max, axis=1)
+
+    return pred
+
+
+__all__ = [
+    "plane_predicate",
+    "axis_plane_predicate",
+    "slab_predicate",
+    "bbox_predicate",
+]

fluxfem/mesh/surface.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+from .base import BaseMesh, BaseMeshPytree
+from .hex import HexMesh, HexMeshPytree
+
+
+def _polygon_area(pts: np.ndarray) -> float:
+    """
+    Polygon area in 3D by fan triangulation (works for tri/quad faces).
+    Assumes points are planar and ordered.
+    """
+    if pts.shape[0] < 3:
+        return 0.0
+    area = 0.0
+    p0 = pts[0]
+    for i in range(1, pts.shape[0] - 1):
+        v1 = pts[i] - p0
+        v2 = pts[i + 1] - p0
+        area += 0.5 * np.linalg.norm(np.cross(v1, v2))
+    return float(area)
+
+
+@dataclass(eq=False)
+class SurfaceMesh(BaseMesh):
+    """
+    Simple boundary mesh made of facets (tri/quad) that live in the volume mesh nodes.
+    Uses BaseMesh.conn to store facets.
+    """
+
+    facet_tags: Optional[jnp.ndarray] = None
+
+    def __post_init__(self):
+        # Keep facet_tags mirrored in cell_tags for BaseMesh compat.
+        if self.cell_tags is None and self.facet_tags is not None:
+            self.cell_tags = self.facet_tags
+        if self.facet_tags is None and self.cell_tags is not None:
+            self.facet_tags = self.cell_tags
+
+    @classmethod
+    def from_facets(
+        cls,
+        coords: jnp.ndarray,
+        facets: jnp.ndarray,
+        facet_tags: Optional[jnp.ndarray] = None,
+        node_tags: Optional[jnp.ndarray] = None,
+    ) -> "SurfaceMesh":
+        coords_j = jnp.asarray(coords, dtype=DTYPE)
+        facets_j = jnp.asarray(facets, dtype=jnp.int32)
+        tags_j = None if facet_tags is None else jnp.asarray(facet_tags, dtype=jnp.int32)
+        node_tags_j = None if node_tags is None else jnp.asarray(node_tags)
+        return cls(coords=coords_j, conn=facets_j, cell_tags=tags_j, node_tags=node_tags_j, facet_tags=tags_j)
+
+    @classmethod
+    def from_hex_mesh(
+        cls,
+        mesh: HexMesh,
+        facets: jnp.ndarray,
+        facet_tags: Optional[jnp.ndarray] = None,
+    ) -> "SurfaceMesh":
+        """
+        Build a surface mesh that reuses the volume mesh coordinates.
+        Facets must reference the volume mesh node numbering.
+        """
+        return cls.from_facets(mesh.coords, facets, facet_tags=facet_tags, node_tags=mesh.node_tags)
+
+    @property
+    def n_facets(self) -> int:
+        return self.n_elems
+
+    def facet_areas(self) -> np.ndarray:
+        """Return per-facet area (uses NumPy for simplicity)."""
+        coords = np.asarray(self.coords)
+        facets = np.asarray(self.conn, dtype=int)
+        areas = np.zeros(facets.shape[0], dtype=float)
+        for i, nodes in enumerate(facets):
+            pts = coords[nodes]
+            areas[i] = _polygon_area(pts)
+        return areas
+
+    def select_by_tag(self, tag: int) -> "SurfaceMesh":
+        """Return a new SurfaceMesh containing only facets with given tag."""
+        if self.facet_tags is None:
+            raise ValueError("facet_tags not set on this SurfaceMesh")
+        mask = np.asarray(self.facet_tags) == tag
+        return SurfaceMesh.from_facets(
+            self.coords,
+            self.conn[mask],
+            facet_tags=self.facet_tags[mask],
+            node_tags=self.node_tags,
+        )
+
+    def facet_normals(self, *, outward_from=None, normalize: bool = True) -> np.ndarray:
+        from ..solver.bc import facet_normals
+        return facet_normals(self, outward_from=outward_from, normalize=normalize)
+
+    def assemble_load(self, load, *, dim: int, n_total_nodes: int | None = None, F0=None):
+        from ..solver.bc import assemble_surface_load
+        return assemble_surface_load(self, load, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_linear_form(self, form, params, *, dim: int, n_total_nodes: int | None = None, F0=None):
+        from ..solver.bc import assemble_surface_linear_form
+        return assemble_surface_linear_form(self, form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_linear_form_on_space(self, space, form, params, *, F0=None):
+        """
+        Assemble surface linear form using global size inferred from a volume space.
+        """
+        dim = int(getattr(space, "value_dim", 1))
+        n_total_nodes = int(getattr(space, "mesh", self).n_nodes)
+        return self.assemble_linear_form(form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_traction(
+        self,
+        traction,
+        *,
+        dim: int = 3,
+        n_total_nodes: int | None = None,
+        F0=None,
+        outward_from=None,
+    ):
+        from ..solver.bc import assemble_surface_traction
+        return assemble_surface_traction(
+            self,
+            traction,
+            dim=dim,
+            n_total_nodes=n_total_nodes,
+            F0=F0,
+            outward_from=outward_from,
+        )
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class SurfaceMeshPytree(BaseMeshPytree):
+    """
+    Simple boundary mesh made of facets (tri/quad) that live in the volume mesh nodes.
+    Uses BaseMesh.conn to store facets.
+    """
+
+    facet_tags: Optional[jnp.ndarray] = None
+
+    def __post_init__(self):
+        if self.cell_tags is None and self.facet_tags is not None:
+            self.cell_tags = self.facet_tags
+        if self.facet_tags is None and self.cell_tags is not None:
+            self.facet_tags = self.cell_tags
+
+    @classmethod
+    def from_facets(
+        cls,
+        coords: jnp.ndarray,
+        facets: jnp.ndarray,
+        facet_tags: Optional[jnp.ndarray] = None,
+        node_tags: Optional[jnp.ndarray] = None,
+    ) -> "SurfaceMeshPytree":
+        coords_j = jnp.asarray(coords, dtype=DTYPE)
+        facets_j = jnp.asarray(facets, dtype=jnp.int32)
+        tags_j = None if facet_tags is None else jnp.asarray(facet_tags, dtype=jnp.int32)
+        node_tags_j = None if node_tags is None else jnp.asarray(node_tags)
+        return cls(coords=coords_j, conn=facets_j, cell_tags=tags_j, node_tags=node_tags_j, facet_tags=tags_j)
+
+    @classmethod
+    def from_hex_mesh(
+        cls,
+        mesh: HexMesh | HexMeshPytree,
+        facets: jnp.ndarray,
+        facet_tags: Optional[jnp.ndarray] = None,
+    ) -> "SurfaceMeshPytree":
+        return cls.from_facets(mesh.coords, facets, facet_tags=facet_tags, node_tags=mesh.node_tags)
+
+    @property
+    def n_facets(self) -> int:
+        return self.n_elems
+
+    def facet_areas(self) -> np.ndarray:
+        coords = np.asarray(self.coords)
+        facets = np.asarray(self.conn, dtype=int)
+        areas = np.zeros(facets.shape[0], dtype=float)
+        for i, nodes in enumerate(facets):
+            pts = coords[nodes]
+            areas[i] = _polygon_area(pts)
+        return areas
+
+    def select_by_tag(self, tag: int) -> "SurfaceMeshPytree":
+        if self.facet_tags is None:
+            raise ValueError("facet_tags not set on this SurfaceMesh")
+        mask = np.asarray(self.facet_tags) == tag
+        return SurfaceMeshPytree.from_facets(
+            self.coords,
+            self.conn[mask],
+            facet_tags=self.facet_tags[mask],
+            node_tags=self.node_tags,
+        )
+
+    def facet_normals(self, *, outward_from=None, normalize: bool = True) -> np.ndarray:
+        from ..solver.bc import facet_normals
+        return facet_normals(self, outward_from=outward_from, normalize=normalize)
+
+    def assemble_load(self, load, *, dim: int, n_total_nodes: int | None = None, F0=None):
+        from ..solver.bc import assemble_surface_load
+        return assemble_surface_load(self, load, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_linear_form(self, form, params, *, dim: int, n_total_nodes: int | None = None, F0=None):
+        from ..solver.bc import assemble_surface_linear_form
+        return assemble_surface_linear_form(self, form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_linear_form_on_space(self, space, form, params, *, F0=None):
+        """
+        Assemble surface linear form using global size inferred from a volume space.
+        """
+        dim = int(getattr(space, "value_dim", 1))
+        n_total_nodes = int(getattr(space, "mesh", self).n_nodes)
+        return self.assemble_linear_form(form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
+
+    def assemble_traction(
+        self,
+        traction,
+        *,
+        dim: int = 3,
+        n_total_nodes: int | None = None,
+        F0=None,
+        outward_from=None,
+    ):
+        from ..solver.bc import assemble_surface_traction
+        return assemble_surface_traction(
+            self,
+            traction,
+            dim=dim,
+            n_total_nodes=n_total_nodes,
+            F0=F0,
+            outward_from=outward_from,
+        )
+
+    def assemble_flux(
+        self,
+        flux,
+        *,
+        n_total_nodes: int | None = None,
+        F0=None,
+        outward_from=None,
+    ):
+        from ..solver.bc import assemble_surface_flux
+        return assemble_surface_flux(
+            self,
+            flux,
+            n_total_nodes=n_total_nodes,
+            F0=F0,
+            outward_from=outward_from,
+        )

fluxfem/mesh/tet.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Optional
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+
+from .base import BaseMesh, BaseMeshPytree
+
+
+@dataclass
+class TetMesh(BaseMesh):
+    """Unstructured tetra mesh."""
+    def face_node_patterns(self):
+        # 4-node faces of a tet
+        return [
+            (0, 1, 2),
+            (0, 1, 3),
+            (0, 2, 3),
+            (1, 2, 3),
+        ]
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class TetMeshPytree(BaseMeshPytree):
+    """Unstructured tetra mesh (pytree)."""
+    def face_node_patterns(self):
+        return [
+            (0, 1, 2),
+            (0, 1, 3),
+            (0, 2, 3),
+            (1, 2, 3),
+        ]
+
+
+@dataclass
+class StructuredTetBox:
+    """Regular grid subdivided into 5 tets per cube (simplex) in a structured layout."""
+
+    nx: int
+    ny: int
+    nz: int
+    lx: float = 1.0
+    ly: float = 1.0
+    lz: float = 1.0
+    origin: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    order: int = 1  # 1: 4-node, 2: 10-node (edge mids)
+
+    def _fix_orientation(self, coords: np.ndarray, conn: np.ndarray) -> np.ndarray:
+        conn = conn.copy()
+        for idx, tet in enumerate(conn):
+            p = coords[tet]
+            p0, p1, p2, p3 = p[0], p[1], p[2], p[3]  # first 4 are corners
+            J = np.stack([p1 - p0, p2 - p0, p3 - p0], axis=1)
+            if np.linalg.det(J) < 0:
+                tet[[1, 2]] = tet[[2, 1]]  # swap corner1/corner2
+                conn[idx] = tet
+        return conn
+
+    def _midpoint(self, a, b):
+        return 0.5 * (a + b)
+
+    def _build_tet10(self, xs, ys, zs) -> TetMesh:
+        coords_list: list[list[float]] = []
+        node_map: dict[tuple[float, float, float], int] = {}
+
+        def add_node(pt: np.ndarray) -> int:
+            key = tuple(np.round(pt.astype(np.float64), 10))
+            if key not in node_map:
+                node_map[key] = len(coords_list)
+                coords_list.append([float(pt[0]), float(pt[1]), float(pt[2])])
+            return node_map[key]
+
+        # corners
+        for k in range(self.nz + 1):
+            for j in range(self.ny + 1):
+                for i in range(self.nx + 1):
+                    add_node(np.array([xs[i], ys[j], zs[k]], dtype=np.float64))
+
+        conn_list = []
+        mid = self._midpoint
+        for k in range(self.nz):
+            for j in range(self.ny):
+                for i in range(self.nx):
+                    p000 = np.array([xs[i], ys[j], zs[k]], dtype=np.float64)
+                    p100 = np.array([xs[i + 1], ys[j], zs[k]], dtype=np.float64)
+                    p010 = np.array([xs[i], ys[j + 1], zs[k]], dtype=np.float64)
+                    p110 = np.array([xs[i + 1], ys[j + 1], zs[k]], dtype=np.float64)
+                    p001 = np.array([xs[i], ys[j], zs[k + 1]], dtype=np.float64)
+                    p101 = np.array([xs[i + 1], ys[j], zs[k + 1]], dtype=np.float64)
+                    p011 = np.array([xs[i], ys[j + 1], zs[k + 1]], dtype=np.float64)
+                    p111 = np.array([xs[i + 1], ys[j + 1], zs[k + 1]], dtype=np.float64)
+
+                    corners = [
+                        (p000, p100, p010, p001),
+                        (p100, p110, p010, p111),
+                        (p100, p010, p001, p111),
+                        (p100, p001, p101, p111),
+                        (p010, p001, p011, p111),
+                    ]
+
+                    for p0, p1, p2, p3 in corners:
+                        n0 = add_node(p0)
+                        n1 = add_node(p1)
+                        n2 = add_node(p2)
+                        n3 = add_node(p3)
+                        # edge midpoints
+                        n01 = add_node(mid(p0, p1))
+                        n02 = add_node(mid(p0, p2))
+                        n03 = add_node(mid(p0, p3))
+                        n12 = add_node(mid(p1, p2))
+                        n13 = add_node(mid(p1, p3))
+                        n23 = add_node(mid(p2, p3))
+                        conn_list.append([n0, n1, n2, n3, n01, n02, n03, n12, n13, n23])
+
+        coords = np.asarray(coords_list, dtype=DTYPE)
+        conn = np.asarray(conn_list, dtype=np.int32)
+        conn = self._fix_orientation(coords, conn)
+        return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))
+
+    def build(self) -> TetMesh:
+        if self.nx <= 0 or self.ny <= 0 or self.nz <= 0:
+            raise ValueError("nx, ny, nz must be positive")
+        if self.order not in (1, 2):
+            raise ValueError("order must be 1 or 2")
+
+        ox, oy, oz = self.origin
+        xs = np.linspace(ox, ox + self.lx, self.nx + 1, dtype=np.float64)
+        ys = np.linspace(oy, oy + self.ly, self.ny + 1, dtype=np.float64)
+        zs = np.linspace(oz, oz + self.lz, self.nz + 1, dtype=np.float64)
+
+        if self.order == 1:
+            return self._build_linear(xs, ys, zs)
+        return self._build_tet10(xs, ys, zs)
+
+    # keep linear build for order=1
+    def _build_linear(self, xs, ys, zs) -> TetMesh:
+        ox, oy, oz = 0.0, 0.0, 0.0  # unused but keep signature consistent
+        coords_list = []
+        for k in range(self.nz + 1):
+            for j in range(self.ny + 1):
+                for i in range(self.nx + 1):
+                    coords_list.append([xs[i], ys[j], zs[k]])
+        coords = np.asarray(coords_list, dtype=DTYPE)
+        def node_id(i: int, j: int, k: int) -> int:
+            return k * (self.ny + 1) * (self.nx + 1) + j * (self.nx + 1) + i
+        conn_list = []
+        for k in range(self.nz):
+            for j in range(self.ny):
+                for i in range(self.nx):
+                    v000 = node_id(i, j, k)
+                    v100 = node_id(i + 1, j, k)
+                    v010 = node_id(i, j + 1, k)
+                    v110 = node_id(i + 1, j + 1, k)
+                    v001 = node_id(i, j, k + 1)
+                    v101 = node_id(i + 1, j, k + 1)
+                    v011 = node_id(i, j + 1, k + 1)
+                    v111 = node_id(i + 1, j + 1, k + 1)
+                    conn_list.extend(
+                        [
+                            [v000, v100, v010, v001],
+                            [v100, v110, v010, v111],
+                            [v100, v010, v001, v111],
+                            [v100, v001, v101, v111],
+                            [v010, v001, v011, v111],
+                        ]
+                    )
+        conn = np.asarray(conn_list, dtype=np.int32)
+        conn = self._fix_orientation(coords, conn)
+        return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))
+
+
+@dataclass
+class StructuredTetTensorBox:
+    """
+    Regular grid subdivided into 6 tets per cube (matches skfem MeshTet.init_tensor).
+    """
+
+    nx: int
+    ny: int
+    nz: int
+    lx: float = 1.0
+    ly: float = 1.0
+    lz: float = 1.0
+    origin: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    order: int = 1  # only linear supported
+
+    def _fix_orientation(self, coords: np.ndarray, conn: np.ndarray) -> np.ndarray:
+        conn = conn.copy()
+        for idx, tet in enumerate(conn):
+            p = coords[tet]
+            p0, p1, p2, p3 = p[0], p[1], p[2], p[3]
+            J = np.stack([p1 - p0, p2 - p0, p3 - p0], axis=1)
+            if np.linalg.det(J) < 0:
+                tet[[1, 2]] = tet[[2, 1]]
+                conn[idx] = tet
+        return conn
+
+    def build(self) -> TetMesh:
+        if self.nx <= 0 or self.ny <= 0 or self.nz <= 0:
+            raise ValueError("nx, ny, nz must be positive")
+        if self.order != 1:
+            raise ValueError("StructuredTetTensorBox only supports order=1")
+
+        ox, oy, oz = self.origin
+        xs = np.linspace(ox, ox + self.lx, self.nx + 1, dtype=np.float64)
+        ys = np.linspace(oy, oy + self.ly, self.ny + 1, dtype=np.float64)
+        zs = np.linspace(oz, oz + self.lz, self.nz + 1, dtype=np.float64)
+        return self._build_linear(xs, ys, zs)
+
+    def _build_linear(self, xs, ys, zs) -> TetMesh:
+        # Mirror scikit-fem MeshTet.init_tensor.
+        npx = len(xs)
+        npy = len(ys)
+        npz = len(zs)
+        X, Y, Z = np.meshgrid(np.sort(xs), np.sort(ys), np.sort(zs))
+        p = np.vstack((X.flatten("F"), Y.flatten("F"), Z.flatten("F")))
+        ix = np.arange(npx * npy * npz)
+        ne = (npx - 1) * (npy - 1) * (npz - 1)
+        t = np.zeros((8, ne), dtype=np.int64)
+        ix = ix.reshape(npy, npx, npz, order="F").copy()
+        t[0] = ix[0:(npy - 1), 0:(npx - 1), 0:(npz - 1)].reshape(ne, 1, order="F").copy().flatten()
+        t[1] = ix[1:npy, 0:(npx - 1), 0:(npz - 1)].reshape(ne, 1, order="F").copy().flatten()
+        t[2] = ix[0:(npy - 1), 1:npx, 0:(npz - 1)].reshape(ne, 1, order="F").copy().flatten()
+        t[3] = ix[0:(npy - 1), 0:(npx - 1), 1:npz].reshape(ne, 1, order="F").copy().flatten()
+        t[4] = ix[1:npy, 1:npx, 0:(npz - 1)].reshape(ne, 1, order="F").copy().flatten()
+        t[5] = ix[1:npy, 0:(npx - 1), 1:npz].reshape(ne, 1, order="F").copy().flatten()
+        t[6] = ix[0:(npy - 1), 1:npx, 1:npz].reshape(ne, 1, order="F").copy().flatten()
+        t[7] = ix[1:npy, 1:npx, 1:npz].reshape(ne, 1, order="F").copy().flatten()
+
+        T = np.zeros((4, 6 * ne), dtype=np.int64)
+        T[:, :ne] = t[[0, 1, 5, 7]]
+        T[:, ne:(2 * ne)] = t[[0, 1, 4, 7]]
+        T[:, (2 * ne):(3 * ne)] = t[[0, 2, 4, 7]]
+        T[:, (3 * ne):(4 * ne)] = t[[0, 3, 5, 7]]
+        T[:, (4 * ne):(5 * ne)] = t[[0, 2, 6, 7]]
+        T[:, (5 * ne):] = t[[0, 3, 6, 7]]
+
+        coords = p.T.astype(DTYPE, copy=False)
+        conn = T.T.astype(np.int32, copy=False)
+        conn = self._fix_orientation(coords, conn)
+        return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))

fluxfem/physics/__init__.py

@@ -0,0 +1,53 @@
+"""Physics-level helpers (constitutive models, material laws, etc.)."""
+
+from .elasticity import (
+    lame_parameters,
+    isotropic_3d_D,
+    linear_elasticity_form,
+    vector_body_force_form,
+    constant_body_force_vector_form,
+    assemble_constant_body_force,
+    right_cauchy_green,
+    green_lagrange_strain,
+    deformation_gradient,
+    pk2_neo_hookean,
+    neo_hookean_residual_form,
+    make_elastic_point_data,
+    write_elastic_vtu,
+    principal_stresses,
+    principal_sum,
+    max_shear_stress,
+    von_mises_stress,
+)
+from .diffusion import diffusion_form
+from .operators import dot, ddot, transpose_last2, sym_grad, sym_grad_u
+from .postprocess import make_point_data_displacement, write_point_data_vtu, interpolate_at_points
+
+__all__ = [
+    "lame_parameters",
+    "isotropic_3d_D",
+    "linear_elasticity_form",
+    "vector_body_force_form",
+    "constant_body_force_vector_form",
+    "assemble_constant_body_force",
+    "diffusion_form",
+    "dot",
+    "ddot",
+    "transpose_last2",
+    "sym_grad",
+    "sym_grad_u",
+    "right_cauchy_green",
+    "green_lagrange_strain",
+    "deformation_gradient",
+    "pk2_neo_hookean",
+    "neo_hookean_residual_form",
+    "make_elastic_point_data",
+    "write_elastic_vtu",
+    "make_point_data_displacement",
+    "write_point_data_vtu",
+    "interpolate_at_points",
+    "principal_stresses",
+    "principal_sum",
+    "max_shear_stress",
+    "von_mises_stress",
+]