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

fluxfem/mesh/surface.py

@@ -1,18 +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:
     """
@@ -26,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)
 
 
@@ -103,20 +117,43 @@ 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)
 
 
@@ -137,13 +174,13 @@ def surface_with_elem_conn(mesh: BaseMesh, facets, *, mode: str = "touching") ->
 
     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,
@@ -222,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,
@@ -259,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

@@ -223,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()
@@ -238,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]]
@@ -247,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_INDEX_DTYPE, 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/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)
@@ -92,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)

fluxfem/physics/elasticity/linear.py

@@ -1,9 +1,14 @@
+from typing import TypeAlias
+
 import jax.numpy as jnp
 
+from ...core.assembly import LinearReturn
 from ...core.forms import FormContext, vector_load_form
-from ...core.basis import build_B_matrices
+from ...core.space import FESpace
 from ...physics.operators import sym_grad
 
+ArrayLike: TypeAlias = jnp.ndarray
+
 # from ...mechanics.kinematics import build_B_matrices
 
 
@@ -41,7 +46,13 @@ def vector_body_force_form(ctx: FormContext, load_vec: jnp.ndarray) -> jnp.ndarr
 vector_body_force_form._ff_kind = "linear"
 vector_body_force_form._ff_domain = "volume"
 
-def assemble_constant_body_force(space, gravity_vec, density: float, *, sparse: bool = False):
+def assemble_constant_body_force(
+    space: FESpace,
+    gravity_vec: ArrayLike,
+    density: float,
+    *,
+    sparse: bool = False,
+) -> LinearReturn:
     """
     Convenience: assemble body force from density * gravity vector.
     gravity_vec: length-3 array-like (direction and magnitude of g)

fluxfem/physics/elasticity/stress.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+from typing import TypeAlias
+
 import jax.numpy as jnp
 
+ArrayLike: TypeAlias = jnp.ndarray
+
 
-def _sym(A: jnp.ndarray) -> jnp.ndarray:
+def _sym(A: ArrayLike) -> ArrayLike:
     return 0.5 * (A + jnp.swapaxes(A, -1, -2))
 
 
-def principal_stresses(S: jnp.ndarray) -> jnp.ndarray:
+def principal_stresses(S: ArrayLike) -> ArrayLike:
     """
     Return principal stresses (eigvals) for symmetric 3x3 stress tensor.
     Supports batching over leading dimensions.
@@ -16,12 +20,12 @@ def principal_stresses(S: jnp.ndarray) -> jnp.ndarray:
     return jnp.linalg.eigvalsh(S_sym)
 
 
-def principal_sum(S: jnp.ndarray) -> jnp.ndarray:
+def principal_sum(S: ArrayLike) -> ArrayLike:
     """Sum of principal stresses (trace)."""
     return jnp.trace(S, axis1=-2, axis2=-1)
 
 
-def max_shear_stress(S: jnp.ndarray) -> jnp.ndarray:
+def max_shear_stress(S: ArrayLike) -> ArrayLike:
     """
     Maximum shear stress = (sigma_max - sigma_min) / 2.
     """
@@ -29,7 +33,7 @@ def max_shear_stress(S: jnp.ndarray) -> jnp.ndarray:
     return 0.5 * (vals[..., -1] - vals[..., 0])
 
 
-def von_mises_stress(S: jnp.ndarray) -> jnp.ndarray:
+def von_mises_stress(S: ArrayLike) -> ArrayLike:
     """
     von Mises equivalent stress: sqrt(3/2 * dev(S):dev(S)).
     """

fluxfem/physics/operators.py

@@ -2,11 +2,18 @@
 
 # fluxfem/mechanics/operators.py
 from __future__ import annotations
+
+from typing import Any, TypeAlias
+
 import jax
 import jax.numpy as jnp
 
+from ..core.context_types import FormFieldLike
+
+ArrayLike: TypeAlias = jnp.ndarray
+
 
-def dot(a, b):
+def dot(a: FormFieldLike | ArrayLike, b: ArrayLike) -> ArrayLike:
     """
     Batched matrix product on the last two axes.
 
@@ -19,7 +26,7 @@ def dot(a, b):
     return jnp.matmul(a, b)
 
 
-def ddot(a, b, c=None):
+def ddot(a: ArrayLike, b: ArrayLike, c: ArrayLike | None = None) -> ArrayLike:
     """
     Double contraction on the last two axes.
 
@@ -32,12 +39,12 @@ def ddot(a, b, c=None):
     return jnp.einsum("...ik,kl,...lm->...im", a_t, b, c)
 
 
-def transpose_last2(a):
+def transpose_last2(a: ArrayLike) -> ArrayLike:
     """Swap the last two axes (batched transpose)."""
     return jnp.swapaxes(a, -1, -2)
 
 
-def sym_grad(field) -> jnp.ndarray:
+def sym_grad(field: FormFieldLike) -> jnp.ndarray:
     """
     Symmetric gradient operator for vector mechanics (small strain).
 
@@ -89,7 +96,7 @@ def sym_grad(field) -> jnp.ndarray:
     return jax.vmap(B_single)(gradN)
 
 
-def sym_grad_u(field, u_elem: jnp.ndarray) -> jnp.ndarray:
+def sym_grad_u(field: FormFieldLike, u_elem: jnp.ndarray) -> jnp.ndarray:
     """
     Apply sym_grad(field) to a local displacement vector.
 

fluxfem/physics/postprocess.py

@@ -1,15 +1,33 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, TypeAlias
+
 import numpy as np
 import jax
 import jax.numpy as jnp
 
 # from ..core.assembly import build_form_contexts
 from ..tools.visualizer import write_vtu
+from ..mesh import BaseMesh
+from ..core.space import FESpace
 from ..core.interp import interpolate_field_at_element_points
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+
 
-def make_point_data_displacement(mesh, space, u, *, compute_j: bool = True, deformed_scale: float = 1.0):
+def make_point_data_displacement(
+    mesh: BaseMesh,
+    space: FESpace,
+    u: ArrayLike,
+    *,
+    compute_j: bool = True,
+    deformed_scale: float = 1.0,
+) -> dict[str, np.ndarray]:
     """
     Common postprocess helper to build point data dictionaries:
       - displacement
@@ -58,14 +76,22 @@ def make_point_data_displacement(mesh, space, u, *, compute_j: bool = True, defo
     return point_data
 
 
-def write_point_data_vtu(mesh, space, u, filepath: str, *, compute_j: bool = True, deformed_scale: float = 1.0):
+def write_point_data_vtu(
+    mesh: BaseMesh,
+    space: FESpace,
+    u: ArrayLike,
+    filepath: str,
+    *,
+    compute_j: bool = True,
+    deformed_scale: float = 1.0,
+) -> None:
     """Write VTU with displacement/deformed_coords and optional J."""
     pdata = make_point_data_displacement(mesh, space, u, compute_j=compute_j, deformed_scale=deformed_scale)
     write_vtu(mesh, filepath, point_data=pdata)
 
 
 __all__ = ["make_point_data_displacement", "write_point_data_vtu", "interpolate_at_points"]
-def interpolate_at_points(space, u, points: np.ndarray):
+def interpolate_at_points(space: FESpace, u: ArrayLike, points: np.ndarray) -> np.ndarray:
     """
     Interpolate displacement field at given physical points (Hex8 only, structured search).
     - points: (m,3) array of physical coordinates.

fluxfem/solver/__init__.py

@@ -24,8 +24,10 @@ from .dirichlet import (
 from .cg import cg_solve, cg_solve_jax, build_cg_operator, CGOperator
 from .preconditioner import make_block_jacobi_preconditioner
 from .block_system import build_block_system, split_block_matrix, BlockSystem
-from .block_matrix import diag as block_diag, make as make_block_matrix
+from .block_matrix import FluxBlockMatrix, diag as block_diag, make as make_block_matrix
 from .newton import newton_solve
+from .result import SolverResult
+from .history import NewtonIterRecord
 from .solve_runner import (
     NonlinearAnalysis,
     NewtonLoopConfig,
@@ -68,9 +70,12 @@ __all__ = [
     "build_block_system",
     "split_block_matrix",
     "BlockSystem",
+    "FluxBlockMatrix",
     "block_diag",
     "make_block_matrix",
     "newton_solve",
+    "SolverResult",
+    "NewtonIterRecord",
     "LinearAnalysis",
     "LinearSolveConfig",
     "LinearStepResult",