fluxfem 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

fluxfem/mesh/supermesh.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+from typing import Iterable
+
+import numpy as np
+
+from .surface import SurfaceMesh
+
+_SUPERMESH_CACHE: dict[tuple, "SurfaceSupermesh"] = {}
+
+
+@dataclass(eq=False)
+class SurfaceSupermesh:
+    """Intersection supermesh for two surface meshes."""
+    coords: np.ndarray
+    conn: np.ndarray
+    source_facets_a: np.ndarray
+    source_facets_b: np.ndarray
+
+
+def _polygon_area_2d(pts: np.ndarray) -> float:
+    x = pts[:, 0]
+    y = pts[:, 1]
+    return 0.5 * float(np.sum(x * np.roll(y, -1) - y * np.roll(x, -1)))
+
+
+def _cross2(a: np.ndarray, b: np.ndarray) -> float:
+    return float(a[0] * b[1] - a[1] * b[0])
+
+
+def _line_intersection(p1, p2, p3, p4, *, tol: float):
+    d1 = p2 - p1
+    d2 = p4 - p3
+    denom = _cross2(d1, d2)
+    if abs(denom) < tol:
+        return p2
+    t = _cross2(p3 - p1, d2) / denom
+    return p1 + t * d1
+
+
+def _sutherland_hodgman(subject: list[np.ndarray], clip: list[np.ndarray], *, tol: float):
+    if not subject:
+        return []
+    orient = np.sign(_polygon_area_2d(np.array(clip)))
+    if orient == 0:
+        return []
+
+    def inside(pt, a, b):
+        return orient * _cross2(b - a, pt - a) >= -tol
+
+    output = subject
+    for i in range(len(clip)):
+        input_list = output
+        if not input_list:
+            break
+        output = []
+        cp1 = clip[i]
+        cp2 = clip[(i + 1) % len(clip)]
+        s = input_list[-1]
+        for e in input_list:
+            if inside(e, cp1, cp2):
+                if not inside(s, cp1, cp2):
+                    output.append(_line_intersection(s, e, cp1, cp2, tol=tol))
+                output.append(e)
+            elif inside(s, cp1, cp2):
+                output.append(_line_intersection(s, e, cp1, cp2, tol=tol))
+            s = e
+    return output
+
+
+def _plane_basis(normal: np.ndarray):
+    n = normal / np.linalg.norm(normal)
+    ref = np.array([1.0, 0.0, 0.0], dtype=float)
+    if abs(np.dot(n, ref)) > 0.9:
+        ref = np.array([0.0, 1.0, 0.0], dtype=float)
+    t1 = np.cross(n, ref)
+    t1 = t1 / np.linalg.norm(t1)
+    t2 = np.cross(n, t1)
+    return t1, t2, n
+
+
+def _facet_plane(pts: np.ndarray, *, tol: float):
+    n = None
+    for i in range(len(pts) - 2):
+        v1 = pts[i + 1] - pts[i]
+        v2 = pts[i + 2] - pts[i]
+        n_candidate = np.cross(v1, v2)
+        n_norm = np.linalg.norm(n_candidate)
+        if n_norm > tol:
+            n = n_candidate / n_norm
+            d = -float(np.dot(n, pts[i]))
+            return n, d
+    return None, None
+
+
+def _coplanar(pts_a: np.ndarray, pts_b: np.ndarray, *, tol: float) -> bool:
+    n, d = _facet_plane(pts_a, tol=tol)
+    if n is None:
+        return False
+    n2, d2 = _facet_plane(pts_b, tol=tol)
+    if n2 is None:
+        return False
+    if abs(abs(np.dot(n, n2)) - 1.0) > 1e-4:
+        return False
+    dist_a = np.abs(pts_a @ n + d)
+    dist_b = np.abs(pts_b @ n + d)
+    return np.max(dist_a) <= tol and np.max(dist_b) <= tol
+
+
+def _project(points: np.ndarray, origin: np.ndarray, t1: np.ndarray, t2: np.ndarray):
+    rel = points - origin[None, :]
+    x = rel @ t1
+    y = rel @ t2
+    return np.stack([x, y], axis=1)
+
+
+def _unique_points(points: Iterable[np.ndarray], *, tol: float):
+    scale = 1.0 / tol
+    mapping: dict[tuple[int, int, int], int] = {}
+    coords: list[np.ndarray] = []
+    indices: list[int] = []
+    for p in points:
+        key = tuple(np.round(p * scale).astype(int))
+        idx = mapping.get(key)
+        if idx is None:
+            idx = len(coords)
+            mapping[key] = idx
+            coords.append(p)
+        indices.append(idx)
+    return np.asarray(coords, dtype=float), indices
+
+
+def _facet_polygon_coords(coords: np.ndarray, facet: np.ndarray) -> np.ndarray:
+    n = int(len(facet))
+    if n == 9:
+        corner = [0, 2, 8, 6]
+        return coords[facet][corner]
+    return coords[facet]
+
+
+def _triangle_min_angle(p0: np.ndarray, p1: np.ndarray, p2: np.ndarray) -> float:
+    def angle(a, b, c):
+        v1 = a - b
+        v2 = c - b
+        n1 = np.linalg.norm(v1)
+        n2 = np.linalg.norm(v2)
+        if n1 == 0.0 or n2 == 0.0:
+            return 0.0
+        cosang = np.clip(np.dot(v1, v2) / (n1 * n2), -1.0, 1.0)
+        return float(np.arccos(cosang))
+
+    return min(angle(p1, p0, p2), angle(p0, p1, p2), angle(p0, p2, p1))
+
+
+def _triangulate_polygon(indices: list[int], poly2d: np.ndarray) -> list[tuple[int, int, int]]:
+    n = len(indices)
+    if n < 3:
+        return []
+    if n == 3:
+        return [(indices[0], indices[1], indices[2])]
+    if n == 4:
+        p = poly2d
+        diag_pref = os.getenv("FLUXFEM_SUPERMESH_QUAD_DIAG", "alt").lower()
+        if diag_pref == "alt":
+            return [(indices[0], indices[1], indices[3]), (indices[1], indices[2], indices[3])]
+        if diag_pref == "fan":
+            return [(indices[0], indices[1], indices[2]), (indices[0], indices[2], indices[3])]
+        min_a = min(
+            _triangle_min_angle(p[0], p[1], p[2]),
+            _triangle_min_angle(p[0], p[2], p[3]),
+        )
+        min_b = min(
+            _triangle_min_angle(p[0], p[1], p[3]),
+            _triangle_min_angle(p[1], p[2], p[3]),
+        )
+        if min_b > min_a:
+            return [(indices[0], indices[1], indices[3]), (indices[1], indices[2], indices[3])]
+        return [(indices[0], indices[1], indices[2]), (indices[0], indices[2], indices[3])]
+    tris = []
+    for i in range(1, n - 1):
+        tris.append((indices[0], indices[i], indices[i + 1]))
+    return tris
+
+
+def build_surface_supermesh(
+    surface_a: SurfaceMesh,
+    surface_b: SurfaceMesh,
+    *,
+    tol: float = 1e-8,
+    cache_enabled: bool | None = None,
+    cache_trace: bool | None = None,
+) -> SurfaceSupermesh:
+    from ..solver.bc import facet_normals
+    import hashlib
+
+    if cache_enabled is None:
+        cache_enabled = os.getenv("FLUXFEM_SUPERMESH_CACHE", "0") not in ("0", "", "false", "False")
+    if cache_trace is None:
+        cache_trace = os.getenv("FLUXFEM_SUPERMESH_CACHE_TRACE", "0") not in ("0", "", "false", "False")
+
+    def _array_sig(arr: np.ndarray) -> tuple:
+        arr_c = np.ascontiguousarray(arr)
+        h = hashlib.blake2b(arr_c.view(np.uint8), digest_size=8).hexdigest()
+        return (arr_c.shape, str(arr_c.dtype), h)
+    if cache_enabled:
+        global _SUPERMESH_CACHE
+        try:
+            _SUPERMESH_CACHE
+        except NameError:
+            _SUPERMESH_CACHE = {}
+        key = (
+            _array_sig(np.asarray(surface_a.coords)),
+            _array_sig(np.asarray(surface_a.conn)),
+            _array_sig(np.asarray(surface_b.coords)),
+            _array_sig(np.asarray(surface_b.conn)),
+            float(tol),
+        )
+        cached = _SUPERMESH_CACHE.get(key)
+        if cached is not None:
+            if cache_trace:
+                print(f"[supermesh] cache hit n_tris={int(cached.conn.shape[0])}", flush=True)
+            return cached
+
+    coords_a = np.asarray(surface_a.coords, dtype=float)
+    coords_b = np.asarray(surface_b.coords, dtype=float)
+    facets_a = np.asarray(surface_a.conn, dtype=int)
+    facets_b = np.asarray(surface_b.conn, dtype=int)
+    normals_a = facet_normals(surface_a, outward_from=np.mean(coords_a, axis=0), normalize=True)
+
+    all_coords: list[np.ndarray] = []
+    all_conn: list[tuple[int, int, int]] = []
+    src_a: list[int] = []
+    src_b: list[int] = []
+
+    for ia, fa in enumerate(facets_a):
+        pts_a = _facet_polygon_coords(coords_a, fa)
+        min_a = pts_a.min(axis=0)
+        max_a = pts_a.max(axis=0)
+        for ib, fb in enumerate(facets_b):
+            pts_b = _facet_polygon_coords(coords_b, fb)
+            if np.any(pts_b.max(axis=0) < min_a - tol) or np.any(pts_b.min(axis=0) > max_a + tol):
+                continue
+            if not _coplanar(pts_a, pts_b, tol=tol):
+                continue
+
+            n, _d = _facet_plane(pts_a, tol=tol)
+            if n is not None:
+                n_ref = normals_a[int(ia)]
+                if np.dot(n, n_ref) < 0.0:
+                    n = -n
+            t1, t2, _ = _plane_basis(n)
+            origin = pts_a[0]
+
+            poly_a = _project(pts_a, origin, t1, t2)
+            poly_b = _project(pts_b, origin, t1, t2)
+
+            inter = _sutherland_hodgman(
+                [p.copy() for p in poly_a],
+                [p.copy() for p in poly_b],
+                tol=tol,
+            )
+            if len(inter) < 3:
+                continue
+            inter_np = np.asarray(inter)
+            if abs(_polygon_area_2d(inter_np)) <= tol:
+                continue
+            center = np.mean(inter_np, axis=0)
+            angles = np.arctan2(inter_np[:, 1] - center[1], inter_np[:, 0] - center[0])
+            order = np.argsort(angles)
+            inter_np = inter_np[order]
+
+            inter_3d = origin[None, :] + inter_np[:, 0:1] * t1 + inter_np[:, 1:2] * t2
+            coords_local, idx = _unique_points(inter_3d, tol=tol)
+            base = len(all_coords)
+            for p in coords_local:
+                all_coords.append(p)
+            tris = _triangulate_polygon(idx, inter_np)
+            for a_idx, b_idx, c_idx in tris:
+                a_id = base + a_idx
+                b_id = base + b_idx
+                c_id = base + c_idx
+                if n is not None:
+                    pa = all_coords[a_id]
+                    pb = all_coords[b_id]
+                    pc = all_coords[c_id]
+                    n_tri = np.cross(pb - pa, pc - pa)
+                    if np.dot(n_tri, n) < 0.0:
+                        b_id, c_id = c_id, b_id
+                all_conn.append((a_id, b_id, c_id))
+                src_a.append(ia)
+                src_b.append(ib)
+
+    if not all_conn:
+        sm = SurfaceSupermesh(
+            coords=np.zeros((0, 3), dtype=float),
+            conn=np.zeros((0, 3), dtype=int),
+            source_facets_a=np.zeros((0,), dtype=int),
+            source_facets_b=np.zeros((0,), dtype=int),
+        )
+        if cache_enabled:
+            _SUPERMESH_CACHE[key] = sm
+        return sm
+
+    coords = np.asarray(all_coords, dtype=float)
+    conn = np.asarray(all_conn, dtype=int)
+    sm = SurfaceSupermesh(
+        coords=coords,
+        conn=conn,
+        source_facets_a=np.asarray(src_a, dtype=int),
+        source_facets_b=np.asarray(src_b, dtype=int),
+    )
+    if cache_enabled:
+        _SUPERMESH_CACHE[key] = sm
+        if cache_trace:
+            print(f"[supermesh] cache store n_tris={int(sm.conn.shape[0])}", flush=True)
+    return sm

fluxfem/mesh/surface.py

@@ -1,16 +1,32 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Callable, Optional, Protocol, Sequence, TYPE_CHECKING, TypeVar, cast
 import jax
 import jax.numpy as jnp
+
+from .dtypes import INDEX_DTYPE
 import numpy as np
+import numpy.typing as npt
 
 DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
 
 from .base import BaseMesh, BaseMeshPytree
 from .hex import HexMesh, HexMeshPytree
 
+P = TypeVar("P")
+
+if TYPE_CHECKING:
+    from ..solver.bc import SurfaceFormContext
+
+
+class SurfaceSpaceLike(Protocol):
+    value_dim: int
+    mesh: BaseMesh
+
+
+SurfaceLinearForm = Callable[["SurfaceFormContext", P], npt.ArrayLike]
+
 
 def _polygon_area(pts: np.ndarray) -> float:
     """
@@ -24,7 +40,7 @@ def _polygon_area(pts: np.ndarray) -> float:
     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))
+        area += float(0.5 * np.linalg.norm(np.cross(v1, v2)))
     return float(area)
 
 
@@ -53,8 +69,8 @@ class SurfaceMesh(BaseMesh):
         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)
+        facets_j = jnp.asarray(facets, dtype=INDEX_DTYPE)
+        tags_j = None if facet_tags is None else jnp.asarray(facet_tags, dtype=INDEX_DTYPE)
         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)
 
@@ -101,31 +117,70 @@ class SurfaceMesh(BaseMesh):
         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):
+    def assemble_load(
+        self,
+        load: npt.ArrayLike,
+        *,
+        dim: int,
+        n_total_nodes: int | None = None,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         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):
+    def assemble_linear_form(
+        self,
+        form: SurfaceLinearForm[P],
+        params: P,
+        *,
+        dim: int,
+        n_total_nodes: int | None = None,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         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):
+    def assemble_linear_form_on_space(
+        self,
+        space: SurfaceSpaceLike,
+        form: SurfaceLinearForm[P],
+        params: P,
+        *,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         """
         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)
+        mesh = cast(BaseMesh, getattr(space, "mesh", self))
+        n_total_nodes = int(mesh.n_nodes)
         return self.assemble_linear_form(form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
 
+
+@dataclass(frozen=True)
+class SurfaceWithElemConn:
+    surface: SurfaceMesh
+    elem_conn: np.ndarray
+
+
+def surface_with_elem_conn(mesh: BaseMesh, facets, *, mode: str = "touching") -> SurfaceWithElemConn:
+    """
+    Build a SurfaceMesh from facets and return it with a matching elem_conn array.
+    """
+    surface = SurfaceMesh.from_facets(mesh.coords, facets, node_tags=mesh.node_tags)
+    elems = mesh.elements_from_facets(facets, mode=mode)
+    elem_conn = np.asarray(mesh.conn, dtype=int)[elems]
+    return SurfaceWithElemConn(surface=surface, elem_conn=elem_conn)
+
     def assemble_traction(
         self,
-        traction,
+        traction: float | Sequence[float],
         *,
         dim: int = 3,
         n_total_nodes: int | None = None,
-        F0=None,
-        outward_from=None,
-    ):
+        F0: npt.ArrayLike | None = None,
+        outward_from: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         from ..solver.bc import assemble_surface_traction
         return assemble_surface_traction(
             self,
@@ -162,8 +217,8 @@ class SurfaceMeshPytree(BaseMeshPytree):
         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)
+        facets_j = jnp.asarray(facets, dtype=INDEX_DTYPE)
+        tags_j = None if facet_tags is None else jnp.asarray(facet_tags, dtype=INDEX_DTYPE)
         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)
 
@@ -204,31 +259,54 @@ class SurfaceMeshPytree(BaseMeshPytree):
         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):
+    def assemble_load(
+        self,
+        load: npt.ArrayLike,
+        *,
+        dim: int,
+        n_total_nodes: int | None = None,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         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):
+    def assemble_linear_form(
+        self,
+        form: SurfaceLinearForm[P],
+        params: P,
+        *,
+        dim: int,
+        n_total_nodes: int | None = None,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         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):
+    def assemble_linear_form_on_space(
+        self,
+        space: SurfaceSpaceLike,
+        form: SurfaceLinearForm[P],
+        params: P,
+        *,
+        F0: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         """
         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)
+        mesh = cast(BaseMesh, getattr(space, "mesh", self))
+        n_total_nodes = int(mesh.n_nodes)
         return self.assemble_linear_form(form, params, dim=dim, n_total_nodes=n_total_nodes, F0=F0)
 
     def assemble_traction(
         self,
-        traction,
+        traction: float | Sequence[float],
         *,
         dim: int = 3,
         n_total_nodes: int | None = None,
-        F0=None,
-        outward_from=None,
-    ):
+        F0: npt.ArrayLike | None = None,
+        outward_from: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         from ..solver.bc import assemble_surface_traction
         return assemble_surface_traction(
             self,
@@ -241,12 +319,12 @@ class SurfaceMeshPytree(BaseMeshPytree):
 
     def assemble_flux(
         self,
-        flux,
+        flux: npt.ArrayLike,
         *,
         n_total_nodes: int | None = None,
-        F0=None,
-        outward_from=None,
-    ):
+        F0: npt.ArrayLike | None = None,
+        outward_from: npt.ArrayLike | None = None,
+    ) -> np.ndarray:
         from ..solver.bc import assemble_surface_flux
         return assemble_surface_flux(
             self,

fluxfem/mesh/tet.py

@@ -5,6 +5,8 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
+from .dtypes import NP_INDEX_DTYPE
+
 DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
 
 
@@ -58,6 +60,10 @@ class StructuredTetBox:
             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
+                if tet.shape[0] == 10:
+                    # keep edge-node ordering consistent with corner swap
+                    tet[[4, 6]] = tet[[6, 4]]  # edges (0-1) <-> (0-2)
+                    tet[[8, 9]] = tet[[9, 8]]  # edges (1-3) <-> (2-3)
                 conn[idx] = tet
         return conn
 
@@ -115,10 +121,10 @@ class StructuredTetBox:
                         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])
+                        conn_list.append([n0, n1, n2, n3, n01, n12, n02, n03, n13, n23])
 
         coords = np.asarray(coords_list, dtype=DTYPE)
-        conn = np.asarray(conn_list, dtype=np.int32)
+        conn = np.asarray(conn_list, dtype=NP_INDEX_DTYPE)
         conn = self._fix_orientation(coords, conn)
         return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))
 
@@ -169,7 +175,7 @@ class StructuredTetBox:
                             [v010, v001, v011, v111],
                         ]
                     )
-        conn = np.asarray(conn_list, dtype=np.int32)
+        conn = np.asarray(conn_list, dtype=NP_INDEX_DTYPE)
         conn = self._fix_orientation(coords, conn)
         return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))
 
@@ -217,11 +223,14 @@ class StructuredTetTensorBox:
         npx = len(xs)
         npy = len(ys)
         npz = len(zs)
+        X: np.ndarray
+        Y: np.ndarray
+        Z: np.ndarray
         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)
+        ix: np.ndarray = np.arange(npx * npy * npz)
         ne = (npx - 1) * (npy - 1) * (npz - 1)
-        t = np.zeros((8, ne), dtype=np.int64)
+        t: np.ndarray = 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()
@@ -232,7 +241,7 @@ class StructuredTetTensorBox:
         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: np.ndarray = 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]]
@@ -241,6 +250,6 @@ class StructuredTetTensorBox:
         T[:, (5 * ne):] = t[[0, 3, 6, 7]]
 
         coords = p.T.astype(DTYPE, copy=False)
-        conn = T.T.astype(np.int32, copy=False)
+        conn: np.ndarray = T.T.astype(NP_INDEX_DTYPE, copy=False)
         conn = self._fix_orientation(coords, conn)
         return TetMesh(coords=jnp.array(coords), conn=jnp.array(conn))

fluxfem/physics/diffusion.py

@@ -15,4 +15,7 @@ def diffusion_form(ctx: FormContext, kappa: float) -> jnp.ndarray:
     return kappa * G
 
 
+diffusion_form._ff_kind = "bilinear"
+diffusion_form._ff_domain = "volume"
+
 __all__ = ["diffusion_form"]

fluxfem/physics/elasticity/hyperelastic.py

@@ -1,11 +1,23 @@
+from typing import Mapping, TYPE_CHECKING, TypeAlias
+
 import jax
 import jax.numpy as jnp
 import numpy as np
 
 from ...core.forms import FormContext
+from ...core.space import FESpace
+from ...mesh import BaseMesh
 from ...core.basis import build_B_matrices_finite
 from ..postprocess import make_point_data_displacement, write_point_data_vtu
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+ParamsLike: TypeAlias = Mapping[str, float] | tuple[float, float]
+
 
 def right_cauchy_green(F: jnp.ndarray) -> jnp.ndarray:
     """C = F^T F (right Cauchy-Green)."""
@@ -46,7 +58,9 @@ def pk2_neo_hookean(F: jnp.ndarray, mu: float, lam: float) -> jnp.ndarray:
     return mu * (I - C_inv) + lam * jnp.log(J)[..., None, None] * C_inv
 
 
-def neo_hookean_residual_form(ctx: FormContext, u_elem: jnp.ndarray, params) -> jnp.ndarray:
+def neo_hookean_residual_form(
+    ctx: FormContext, u_elem: jnp.ndarray, params: ParamsLike
+) -> jnp.ndarray:
     """
     Compressible Neo-Hookean residual (Total Lagrangian, PK2).
     params: dict-like with keys \"mu\", \"lam\" or tuple (mu, lam)
@@ -78,6 +92,9 @@ def neo_hookean_residual_form(ctx: FormContext, u_elem: jnp.ndarray, params) ->
     return jnp.einsum("qik,qk->qi", BT, S_voigt)   # (n_q, n_ldofs)
 
 
+neo_hookean_residual_form._ff_kind = "residual"
+neo_hookean_residual_form._ff_domain = "volume"
+
 __all__ = [
     "right_cauchy_green",
     "green_lagrange_strain",
@@ -89,11 +106,26 @@ __all__ = [
 ]
 
 
-def make_elastic_point_data(mesh, space, u, *, compute_j: bool = True, deformed_scale: float = 1.0):
+def make_elastic_point_data(
+    mesh: BaseMesh,
+    space: FESpace,
+    u: ArrayLike,
+    *,
+    compute_j: bool = True,
+    deformed_scale: float = 1.0,
+) -> dict[str, np.ndarray]:
     """Alias to postprocess.make_point_data_displacement for backward compatibility."""
     return make_point_data_displacement(mesh, space, u, compute_j=compute_j, deformed_scale=deformed_scale)
 
 
-def write_elastic_vtu(mesh, space, u, filepath: str, *, compute_j: bool = True, deformed_scale: float = 1.0):
+def write_elastic_vtu(
+    mesh: BaseMesh,
+    space: FESpace,
+    u: ArrayLike,
+    filepath: str,
+    *,
+    compute_j: bool = True,
+    deformed_scale: float = 1.0,
+) -> None:
     """Alias to postprocess.write_point_data_vtu for backward compatibility."""
     return write_point_data_vtu(mesh, space, u, filepath, compute_j=compute_j, deformed_scale=deformed_scale)