fluxfem 0.1.1a0__py3-none-any.whl

fluxfem/helpers_num.py

@@ -0,0 +1,11 @@
+"""Numeric helpers (array operators)."""
+from __future__ import annotations
+
+from .physics.operators import dot, ddot, sym_grad, transpose_last2
+
+__all__ = [
+    "dot",
+    "ddot",
+    "sym_grad",
+    "transpose_last2",
+]

fluxfem/helpers_wf.py

@@ -0,0 +1,42 @@
+"""WeakForm/Expr helpers (symbolic operators)."""
+from __future__ import annotations
+
+from .core.weakform import (
+    grad,
+    sym_grad,
+    dot,
+    sdot,
+    ddot,
+    inner,
+    action,
+    gaction,
+    I,
+    det,
+    inv,
+    transpose,
+    transpose_last2,
+    log,
+    normal,
+    ds,
+    dOmega,
+)
+
+__all__ = [
+    "grad",
+    "sym_grad",
+    "dot",
+    "sdot",
+    "ddot",
+    "inner",
+    "action",
+    "gaction",
+    "I",
+    "det",
+    "inv",
+    "transpose",
+    "transpose_last2",
+    "log",
+    "normal",
+    "ds",
+    "dOmega",
+]

fluxfem/mesh/__init__.py

@@ -0,0 +1,29 @@
+from .hex import HexMesh, HexMeshPytree, StructuredHexBox, tag_axis_minmax_facets
+from .tet import TetMesh, TetMeshPytree, StructuredTetBox, StructuredTetTensorBox
+from .base import BaseMesh, BaseMeshPytree
+from .predicate import bbox_predicate, plane_predicate, axis_plane_predicate, slab_predicate
+from .surface import SurfaceMesh, SurfaceMeshPytree
+from .io import load_gmsh_mesh, load_gmsh_hex_mesh, load_gmsh_tet_mesh, make_surface_from_facets
+
+__all__ = [
+    "BaseMesh",
+    "BaseMeshPytree",
+    "bbox_predicate",
+    "plane_predicate",
+    "axis_plane_predicate",
+    "slab_predicate",
+    "HexMesh",
+    "HexMeshPytree",
+    "StructuredHexBox",
+    "tag_axis_minmax_facets",
+    "TetMesh",
+    "TetMeshPytree",
+    "StructuredTetBox",
+    "StructuredTetTensorBox",
+    "SurfaceMesh",
+    "SurfaceMeshPytree",
+    "load_gmsh_mesh",
+    "load_gmsh_hex_mesh",
+    "load_gmsh_tet_mesh",
+    "make_surface_from_facets",
+]

fluxfem/mesh/base.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable, Iterable, Optional, Sequence
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+
+@dataclass
+class BaseMeshClosure:
+    coords: jnp.ndarray
+    conn: jnp.ndarray
+    cell_tags: Optional[jnp.ndarray] = None
+    node_tags: Optional[jnp.ndarray] = None
+
+    @property
+    def n_nodes(self) -> int:
+        return self.coords.shape[0]
+
+    @property
+    def n_elems(self) -> int:
+        return self.conn.shape[0]
+
+    def element_coords(self) -> jnp.ndarray:
+        return self.coords[self.conn]
+
+    # ------------------------------------------------------------------
+    # Face patterns must be provided by concrete mesh types.
+    def face_node_patterns(self):
+        """
+        Return a list of tuples, each tuple giving local node indices of a face.
+        Override in concrete mesh classes (HexMesh, TetMesh, etc).
+        """
+        raise NotImplementedError("face_node_patterns must be implemented by mesh subtype")
+
+    # Convenience helpers for boundary tagging / DOF lookup
+    def node_indices_where(self, predicate: Callable[[np.ndarray], np.ndarray]) -> np.ndarray:
+        """
+        Return node indices whose coordinates satisfy the predicate.
+        predicate: callable that takes coords (np.ndarray of shape (n_nodes, dim)) and returns boolean mask.
+        """
+        coords_np = np.asarray(self.coords)
+        mask = predicate(coords_np)
+        return np.nonzero(mask)[0]
+
+    def node_indices_where_point(self, predicate: Callable[[np.ndarray], bool]) -> np.ndarray:
+        """
+        Return node indices for which predicate(coord) is True.
+        predicate: callable accepting a single point (dim,) -> bool
+        """
+        coords_np = np.asarray(self.coords)
+        mask = [bool(predicate(pt)) for pt in coords_np]
+        return np.nonzero(mask)[0]
+
+    def axis_extrema_nodes(self, axis: int = 0, side: str = "min", tol: float = 1e-8) -> np.ndarray:
+        """
+        Nodes lying on min or max of a given axis.
+        side: "min" or "max"
+        """
+        coords_np = np.asarray(self.coords)
+        vals = coords_np[:, axis]
+        target = vals.min() if side == "min" else vals.max()
+        mask = np.isclose(vals, target, atol=tol)
+        return np.nonzero(mask)[0]
+
+    def boundary_nodes_bbox(self, tol: float = 1e-8) -> np.ndarray:
+        """
+        Nodes on the axis-aligned bounding box (min/max in each coordinate).
+        Useful for box-shaped meshes like StructuredHexBox.
+        """
+        coords_np = np.asarray(self.coords)
+        mins = coords_np.min(axis=0)
+        maxs = coords_np.max(axis=0)
+        mask = np.zeros(coords_np.shape[0], dtype=bool)
+        for axis in range(coords_np.shape[1]):
+            mask |= np.isclose(coords_np[:, axis], mins[axis], atol=tol)
+            mask |= np.isclose(coords_np[:, axis], maxs[axis], atol=tol)
+        return np.nonzero(mask)[0]
+
+    def node_dofs(self, nodes: Iterable[int], components: Sequence[int] | str = "xyz", dof_per_node: Optional[int] = None) -> np.ndarray:
+        """
+        Build flattened DOF indices for given node ids.
+
+        components:
+            - sequence of component indices (e.g., [0,1,2])
+            - or string like "x", "xy", "xyz" (case-insensitive; maps x/y/z -> 0/1/2)
+        dof_per_node: optional; inferred from max component index + 1 if not provided.
+        """
+        nodes_arr = np.asarray(list(nodes), dtype=int)
+        if isinstance(components, str):
+            comp_map = {"x": 0, "y": 1, "z": 2}
+            comps = np.asarray([comp_map[c.lower()] for c in components], dtype=int)
+        else:
+            comps = np.asarray(list(components), dtype=int)
+        inferred = int(comps.max()) + 1 if comps.size else 1
+        dofpn = inferred if dof_per_node is None else int(dof_per_node)
+        if dofpn <= comps.max():
+            raise ValueError(f"dof_per_node={dofpn} is inconsistent with requested component {comps.max()}")
+        dofs = [dofpn * int(n) + int(c) for n in nodes_arr for c in comps]
+        return np.asarray(dofs, dtype=int)
+
+    def dofs_where(self, predicate: Callable[[np.ndarray], np.ndarray], components: Sequence[int] | str = "xyz", dof_per_node: Optional[int] = None) -> np.ndarray:
+        """
+        DOF indices for nodes selected by a predicate over all coords.
+        predicate takes coords (np.ndarray, shape (n_nodes, dim)) and returns boolean mask.
+        """
+        nodes = self.node_indices_where(predicate)
+        return self.node_dofs(nodes, components=components, dof_per_node=dof_per_node)
+
+    def dofs_where_point(self, predicate: Callable[[np.ndarray], bool], components: Sequence[int] | str = "xyz", dof_per_node: Optional[int] = None) -> np.ndarray:
+        """
+        DOF indices for nodes selected by a per-point predicate.
+        predicate takes a single coord (dim,) -> bool.
+        """
+        nodes = self.node_indices_where_point(predicate)
+        return self.node_dofs(nodes, components=components, dof_per_node=dof_per_node)
+
+    def boundary_dofs_where(self, predicate: Callable[[np.ndarray], np.ndarray], components: Sequence[int] | str = "xyz", dof_per_node: Optional[int] = None) -> np.ndarray:
+        """
+        Return DOF indices for boundary nodes whose coordinates satisfy predicate.
+        predicate takes coords (np.ndarray, shape (n_nodes, dim)) and returns boolean mask.
+        """
+        coords_np = np.asarray(self.coords)
+        mask = np.asarray(predicate(coords_np), dtype=bool)
+        bmask = self.boundary_node_mask()
+        nodes = np.nonzero(mask & bmask)[0]
+        return self.node_dofs(nodes, components=components, dof_per_node=dof_per_node)
+
+    def boundary_dofs_bbox(
+        self,
+        *,
+        components: Sequence[int] | str = "xyz",
+        dof_per_node: Optional[int] = None,
+        tol: float = 1e-8,
+    ) -> np.ndarray:
+        """
+        DOF indices on the axis-aligned bounding box (min/max in each coordinate).
+        """
+        nodes = self.boundary_nodes_bbox(tol=tol)
+        return self.node_dofs(nodes, components=components, dof_per_node=dof_per_node)
+
+    def boundary_node_indices(self) -> np.ndarray:
+        """
+        Return node indices on the boundary based on element face adjacency.
+        """
+        cached = getattr(self, "_boundary_nodes_cache", None)
+        if cached is not None:
+            return cached
+        conn = np.asarray(self.conn)
+        patterns = self.face_node_patterns()
+        face_counts: dict[tuple[int, ...], int] = {}
+        for elem_conn in conn:
+            for pattern in patterns:
+                nodes = tuple(sorted(int(elem_conn[i]) for i in pattern))
+                face_counts[nodes] = face_counts.get(nodes, 0) + 1
+        bnodes = set()
+        for nodes, count in face_counts.items():
+            if count == 1:
+                bnodes.update(nodes)
+        out = np.asarray(sorted(bnodes), dtype=int)
+        setattr(self, "_boundary_nodes_cache", out)
+        return out
+
+    def boundary_node_mask(self) -> np.ndarray:
+        """
+        Return boolean mask for boundary nodes (shape: n_nodes).
+        """
+        mask = np.zeros(self.n_nodes, dtype=bool)
+        nodes = self.boundary_node_indices()
+        mask[nodes] = True
+        return mask
+
+    def make_node_tags(self, predicate: Callable[[np.ndarray], np.ndarray], tag: int, base: Optional[np.ndarray] = None) -> jnp.ndarray:
+        """
+        Build a node_tags array by applying predicate to coords and setting tag where True.
+        Returns a jnp.ndarray (int32). Does not mutate the mesh.
+        """
+        base_tags = np.zeros(self.n_nodes, dtype=np.int32) if base is None else np.asarray(base, dtype=np.int32).copy()
+        mask = predicate(np.asarray(self.coords))
+        base_tags[mask] = int(tag)
+        return jnp.asarray(base_tags, dtype=jnp.int32)
+
+    def with_node_tags(self, node_tags: np.ndarray | jnp.ndarray):
+        """
+        Return a new mesh instance with provided node_tags.
+        """
+        return self.__class__(coords=self.coords, conn=self.conn, cell_tags=self.cell_tags, node_tags=jnp.asarray(node_tags))
+
+    def boundary_facets_where(self, predicate: Callable[[np.ndarray], bool], tag: int | None = None):
+        """
+        Collect boundary facets whose node coordinates satisfy predicate.
+
+        predicate receives a (n_face_nodes, dim) NumPy array and returns True/False.
+        Returns facets (and optional tags if tag is provided).
+        """
+        coords = np.asarray(self.coords)
+        conn = np.asarray(self.conn)
+        patterns = self.face_node_patterns()
+
+        facet_map: dict[tuple[int, ...], tuple[list[int], Optional[int]]] = {}
+
+        for elem_conn in conn:
+            elem_nodes = coords[elem_conn]
+            for pattern in patterns:
+                nodes = [int(elem_conn[i]) for i in pattern]
+                face_coords = elem_nodes[list(pattern)]
+                if not predicate(face_coords):
+                    continue
+                key = tuple(sorted(nodes))
+                if key not in facet_map:
+                    facet_map[key] = (nodes, tag)
+
+        if not facet_map:
+            if tag is None:
+                return jnp.empty((0, len(patterns[0]) if patterns else 0), dtype=jnp.int32)
+            return jnp.empty((0, len(patterns[0]) if patterns else 0), dtype=jnp.int32), jnp.empty((0,), dtype=jnp.int32)
+
+        facets = []
+        tags = []
+        for nodes, t in facet_map.values():
+            facets.append(nodes)
+            if tag is not None:
+                tags.append(t if t is not None else 0)
+
+        facets_arr = jnp.array(facets, dtype=jnp.int32)
+        if tag is None:
+            return facets_arr
+        return facets_arr, jnp.array(tags, dtype=jnp.int32)
+
+
+@jax.tree_util.register_pytree_node_class
+class BaseMeshPytree(BaseMeshClosure):
+    def tree_flatten(self):
+        children = (self.coords, self.conn, self.cell_tags, self.node_tags)
+        return children, {}
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        coords, conn, cell_tags, node_tags = children
+        return cls(coords, conn, cell_tags, node_tags)
+
+
+BaseMesh = BaseMeshClosure
+

fluxfem/mesh/hex.py

@@ -0,0 +1,327 @@
+
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Optional, Dict, Tuple, List, Callable
+import jax
+import jax.numpy as jnp
+from ..core.dtypes import DEFAULT_DTYPE
+import numpy as np
+
+from .base import BaseMesh, BaseMeshPytree
+
+
+@dataclass
+class HexMesh(BaseMesh):
+    """
+    Structured / unstructured hex mesh (8-node linear hex elements).
+    
+    coords: (n_nodes, 3) float32
+    conn:   (n_elems, 8) int32  # node indices of each element
+    """
+    coords: jnp.ndarray  # shape (n_nodes, 3)
+    conn: jnp.ndarray    # shape (n_elems, 8), int32
+
+    def face_node_patterns(self):
+        return [
+            (0, 1, 2, 3),  # -z
+            (4, 5, 6, 7),  # +z
+            (0, 1, 5, 4),  # -y
+            (3, 2, 6, 7),  # +y
+            (0, 4, 7, 3),  # -x
+            (1, 2, 6, 5),  # +x
+        ]
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(eq=False)
+class HexMeshPytree(BaseMeshPytree):
+    def face_node_patterns(self):
+        return [
+            (0, 1, 2, 3),  # -z
+            (4, 5, 6, 7),  # +z
+            (0, 1, 5, 4),  # -y
+            (3, 2, 6, 7),  # +y
+            (0, 4, 7, 3),  # -x
+            (1, 2, 6, 5),  # +x
+        ]
+
+
+def tag_axis_minmax_facets(
+    mesh: HexMesh,
+    axis: int = 0,
+    dirichlet_tag: int = 1,
+    neumann_tag: int = 2,
+    tol: float = 1e-8,
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Tag boundary facets on min/max of the given axis.
+
+    Returns:
+        facets: (n_facets, 4) int32, quad node ids
+        facet_tags: (n_facets,) int32, dirichlet_tag on min side, neumann_tag on max side
+    """
+    coords = np.asarray(mesh.coords)
+    conn = np.asarray(mesh.conn)
+
+    axis_vals = coords[:, axis]
+    v_min = float(axis_vals.min())
+    v_max = float(axis_vals.max())
+
+    face_patterns: List[Tuple[int, int, int, int]] = [
+        (0, 1, 2, 3),  # -z
+        (4, 5, 6, 7),  # +z
+        (0, 1, 5, 4),  # -y
+        (3, 2, 6, 7),  # +y
+        (0, 4, 7, 3),  # -x
+        (1, 2, 6, 5),  # +x
+    ]
+
+    facet_map: Dict[Tuple[int, ...], Tuple[List[int], int]] = {}
+
+    for elem_conn in conn:
+        elem_nodes = coords[elem_conn]
+        for pattern in face_patterns:
+            nodes = [int(elem_conn[i]) for i in pattern]
+            vals = elem_nodes[list(pattern), axis]
+            tag = None
+            if np.allclose(vals, v_min, atol=tol):
+                tag = dirichlet_tag
+            elif np.allclose(vals, v_max, atol=tol):
+                tag = neumann_tag
+
+            if tag is None:
+                continue
+
+            key = tuple(sorted(nodes))
+            if key not in facet_map:
+                facet_map[key] = (nodes, tag)
+
+    if not facet_map:
+        return jnp.empty((0, 4), dtype=jnp.int32), jnp.empty((0,), dtype=jnp.int32)
+
+    facets = []
+    tags = []
+    for nodes, tag in facet_map.values():
+        facets.append(nodes)
+        tags.append(tag)
+
+    return jnp.array(facets, dtype=jnp.int32), jnp.array(tags, dtype=jnp.int32)
+
+
+@dataclass
+class StructuredHexBox:
+    """
+    Uniform hex mesh generator on a box, returned as an unstructured HexMesh.
+    """
+    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: 8-node Hex, 2: 20-node serendipity Hex, 3: 27-node triquadratic Hex
+
+    def build(self) -> HexMesh:
+        """
+        Build a regular grid of nx×ny×nz elements over [origin, origin + (lx, ly, lz)].
+        order=1 → 8-node Hex, order=2 → 20-node serendipity Hex。
+        """
+        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, 3):
+            raise ValueError("order must be 1, 2, or 3")
+
+        ox, oy, oz = self.origin
+        xs = jnp.linspace(ox, ox + self.lx, self.nx + 1, dtype=DEFAULT_DTYPE)
+        ys = jnp.linspace(oy, oy + self.ly, self.ny + 1, dtype=DEFAULT_DTYPE)
+        zs = jnp.linspace(oz, oz + self.lz, self.nz + 1, dtype=DEFAULT_DTYPE)
+
+        if self.order == 1:
+            return self._build_hex8(xs, ys, zs)
+        if self.order == 2:
+            return self._build_hex20(xs, ys, zs)
+        return self._build_hex27(xs, ys, zs)
+
+    def _build_hex8(self, xs, ys, zs) -> HexMesh:
+        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 = jnp.array(coords_list, dtype=DEFAULT_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):
+                    n000 = node_id(i,     j,     k)
+                    n100 = node_id(i + 1, j,     k)
+                    n110 = node_id(i + 1, j + 1, k)
+                    n010 = node_id(i,     j + 1, k)
+                    n001 = node_id(i,     j,     k + 1)
+                    n101 = node_id(i + 1, j,     k + 1)
+                    n111 = node_id(i + 1, j + 1, k + 1)
+                    n011 = node_id(i,     j + 1, k + 1)
+                    conn_list.append([n000, n100, n110, n010, n001, n101, n111, n011])
+
+        conn = jnp.array(conn_list, dtype=jnp.int32)
+        return HexMesh(coords=coords, conn=conn)
+
+    def _build_hex20(self, xs, ys, zs) -> HexMesh:
+        """
+        Build 20-node serendipity Hex mesh (corner + edge midpoints).
+        """
+        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), 8))
+            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]
+
+        # pre-create corner nodes
+        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 = []
+        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)
+                    p110 = np.array([xs[i + 1], ys[j + 1], zs[k]],     dtype=np.float64)
+                    p010 = np.array([xs[i],     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)
+                    p111 = np.array([xs[i + 1], ys[j + 1], zs[k + 1]], dtype=np.float64)
+                    p011 = np.array([xs[i],     ys[j + 1], zs[k + 1]], dtype=np.float64)
+
+                    c000 = add_node(p000)
+                    c100 = add_node(p100)
+                    c110 = add_node(p110)
+                    c010 = add_node(p010)
+                    c001 = add_node(p001)
+                    c101 = add_node(p101)
+                    c111 = add_node(p111)
+                    c011 = add_node(p011)
+
+                    # edge midpoints
+                    e01 = add_node(0.5 * (p000 + p100))
+                    e12 = add_node(0.5 * (p100 + p110))
+                    e23 = add_node(0.5 * (p110 + p010))
+                    e30 = add_node(0.5 * (p010 + p000))
+
+                    e45 = add_node(0.5 * (p001 + p101))
+                    e56 = add_node(0.5 * (p101 + p111))
+                    e67 = add_node(0.5 * (p111 + p011))
+                    e74 = add_node(0.5 * (p011 + p001))
+
+                    e04 = add_node(0.5 * (p000 + p001))
+                    e15 = add_node(0.5 * (p100 + p101))
+                    e26 = add_node(0.5 * (p110 + p111))
+                    e37 = add_node(0.5 * (p010 + p011))
+
+                    conn_list.append(
+                        [
+                            c000, c100, c110, c010,  # corners
+                            c001, c101, c111, c011,
+                            e01, e12, e23, e30,     # edge mids
+                            e45, e56, e67, e74,
+                            e04, e15, e26, e37,
+                        ]
+                    )
+
+        coords = jnp.array(coords_list, dtype=DEFAULT_DTYPE)
+        conn = jnp.array(conn_list, dtype=jnp.int32)
+        return HexMesh(coords=coords, conn=conn)
+
+    def _build_hex27(self, xs, ys, zs) -> HexMesh:
+        """
+        Build 27-node triquadratic Hex mesh (tensor-product nodes: corners, edge mids, face centers, body center).
+        """
+        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), 8))
+            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]
+
+        def mid(a, b):
+            return 0.5 * (a + b)
+
+        # pre-create grid nodes at original vertices and midpoints on axes
+        xs_mid = [mid(xs[i], xs[i + 1]) for i in range(len(xs) - 1)]
+        ys_mid = [mid(ys[j], ys[j + 1]) for j in range(len(ys) - 1)]
+        zs_mid = [mid(zs[k], zs[k + 1]) for k in range(len(zs) - 1)]
+
+        # create all possible nodes (vertices + edge mids + face centers + cell centers)
+        for k_idx, zk in enumerate(zs):
+            for j_idx, yj in enumerate(ys):
+                for i_idx, xi in enumerate(xs):
+                    add_node(np.array([xi, yj, zk], dtype=np.float64))
+        for k_idx, zk in enumerate(zs):
+            for j_idx, yj in enumerate(ys):
+                for i_mid in xs_mid:
+                    add_node(np.array([i_mid, yj, zk], dtype=np.float64))
+        for k_idx, zk in enumerate(zs):
+            for i_idx, xi in enumerate(xs):
+                for j_mid in ys_mid:
+                    add_node(np.array([xi, j_mid, zk], dtype=np.float64))
+        for j_idx, yj in enumerate(ys):
+            for i_idx, xi in enumerate(xs):
+                for k_mid in zs_mid:
+                    add_node(np.array([xi, yj, k_mid], dtype=np.float64))
+        # face centers
+        for k_idx, zk in enumerate(zs):
+            for j_mid in ys_mid:
+                for i_mid in xs_mid:
+                    add_node(np.array([i_mid, j_mid, zk], dtype=np.float64))
+        for k_mid in zs_mid:
+            for j_idx, yj in enumerate(ys):
+                for i_mid in xs_mid:
+                    add_node(np.array([i_mid, yj, k_mid], dtype=np.float64))
+        for k_mid in zs_mid:
+            for j_mid in ys_mid:
+                for i_idx, xi in enumerate(xs):
+                    add_node(np.array([xi, j_mid, k_mid], dtype=np.float64))
+        # cell centers (unique per cell)
+        for k in range(self.nz):
+            for j in range(self.ny):
+                for i in range(self.nx):
+                    cx = mid(xs[i], xs[i + 1])
+                    cy = mid(ys[j], ys[j + 1])
+                    cz = mid(zs[k], zs[k + 1])
+                    add_node(np.array([cx, cy, cz], dtype=np.float64))
+
+        conn_list = []
+        for k in range(self.nz):
+            for j in range(self.ny):
+                for i in range(self.nx):
+                    x_vals = [xs[i], mid(xs[i], xs[i + 1]), xs[i + 1]]
+                    y_vals = [ys[j], mid(ys[j], ys[j + 1]), ys[j + 1]]
+                    z_vals = [zs[k], mid(zs[k], zs[k + 1]), zs[k + 1]]
+
+                    nodes = []
+                    for kk in range(3):
+                        for jj in range(3):
+                            for ii in range(3):
+                                nodes.append(add_node(np.array([x_vals[ii], y_vals[jj], z_vals[kk]], dtype=np.float64)))
+
+                    # order in lexicographic k,j,i -> length 27
+                    conn_list.append(nodes)
+
+        coords = jnp.array(coords_list, dtype=DEFAULT_DTYPE)
+        conn = jnp.array(conn_list, dtype=jnp.int32)
+        return HexMesh(coords=coords, conn=conn)