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

fluxfem/__init__.py

@@ -31,9 +31,6 @@ __all__ = [
     "compile_surface_bilinear",
     "compile_mixed_surface_residual",
     "compile_mixed_residual",
-    "outer",
-    "sdot",
-    "dOmega",
     "FormContext",
     "MixedFormContext",
     "VolumeContext",
@@ -85,11 +82,6 @@ __all__ = [
     "constant_body_force_form",
     "constant_body_force_vector_form",
     "diffusion_form",
-    "dot",
-    "ddot",
-    "transpose_last2",
-    "sym_grad",
-    "sym_grad_u",
     "right_cauchy_green",
     "green_lagrange_strain",
     "deformation_gradient",
@@ -172,6 +164,7 @@ __all__ = [
     "coo_to_csr",
     "SparsityPattern",
     "FluxSparseMatrix",
+    "FluxBlockMatrix",
     "coalesce_coo",
     "concat_flux",
     "block_diag_flux",
@@ -233,11 +226,6 @@ _PHYSICS_EXPORTS = {
     "vector_body_force_form",
     "constant_body_force_vector_form",
     "diffusion_form",
-    "dot",
-    "ddot",
-    "transpose_last2",
-    "sym_grad",
-    "sym_grad_u",
     "right_cauchy_green",
     "green_lagrange_strain",
     "deformation_gradient",

fluxfem/core/__init__.py

@@ -80,76 +80,6 @@ from .weakform import (
     compile_linear,
     compile_residual,
     compile_mixed_residual,
-    grad,
-    sym_grad,
-    outer,
-    dot,
-    sdot,
-    ddot,
-    inner,
-    action,
-    gaction,
-    normal,
-    ds,
-    dOmega,
-    I,
-    det,
-    inv,
-    transpose,
-    log,
-    transpose_last2,
-    matmul,
-    matmul_std,
-    einsum,
-)
-from ..mesh import (
-    BaseMeshPytree,
-    SurfaceWithFacetMap,
-    bbox_predicate,
-    plane_predicate,
-    axis_plane_predicate,
-    slab_predicate,
-    HexMesh,
-    HexMeshPytree,
-    StructuredHexBox,
-    SurfaceMesh,
-    SurfaceMeshPytree,
-    SurfaceWithElemConn,
-    surface_with_elem_conn,
-    SurfaceSupermesh,
-    build_surface_supermesh,
-    tag_axis_minmax_facets,
-    TetMesh,
-    TetMeshPytree,
-    StructuredTetBox,
-    StructuredTetTensorBox,
-    load_gmsh_mesh,
-    load_gmsh_hex_mesh,
-    load_gmsh_tet_mesh,
-    make_surface_from_facets,
-    SurfaceSupermesh,
-    build_surface_supermesh,
-    MortarMatrix,
-    assemble_mortar_matrices,
-    assemble_contact_onesided_floor,
-    assemble_mixed_surface_residual,
-    assemble_mixed_surface_jacobian,
-    map_surface_facets_to_tet_elements,
-    map_surface_facets_to_hex_elements,
-    tri_area,
-    tri_quadrature,
-    facet_triangles,
-    facet_shape_values,
-    volume_shape_values_at_points,
-    quad_shape_and_local,
-    quad9_shape_values,
-    hex27_gradN,
-    ContactSurfaceSpace,
-    ContactSide,
-    OneSidedContact,
-    OneSidedContactSurfaceSpace,
-    facet_gap_values,
-    active_contact_facets,
 )
 from .basis import (
     HexTriLinearBasis,
@@ -167,9 +97,58 @@ import importlib
 
 from .solver import spdirect_solve_cpu, spdirect_solve_gpu, spdirect_solve_jax, coo_to_csr
 
+_MESH_EXPORTS = {
+    "BaseMeshPytree",
+    "SurfaceWithFacetMap",
+    "bbox_predicate",
+    "plane_predicate",
+    "axis_plane_predicate",
+    "slab_predicate",
+    "HexMesh",
+    "HexMeshPytree",
+    "StructuredHexBox",
+    "SurfaceMesh",
+    "SurfaceMeshPytree",
+    "SurfaceWithElemConn",
+    "surface_with_elem_conn",
+    "SurfaceSupermesh",
+    "build_surface_supermesh",
+    "tag_axis_minmax_facets",
+    "TetMesh",
+    "TetMeshPytree",
+    "StructuredTetBox",
+    "StructuredTetTensorBox",
+    "load_gmsh_mesh",
+    "load_gmsh_hex_mesh",
+    "load_gmsh_tet_mesh",
+    "make_surface_from_facets",
+    "MortarMatrix",
+    "assemble_mortar_matrices",
+    "assemble_contact_onesided_floor",
+    "assemble_mixed_surface_residual",
+    "assemble_mixed_surface_jacobian",
+    "map_surface_facets_to_tet_elements",
+    "map_surface_facets_to_hex_elements",
+    "tri_area",
+    "tri_quadrature",
+    "facet_triangles",
+    "facet_shape_values",
+    "volume_shape_values_at_points",
+    "quad_shape_and_local",
+    "quad9_shape_values",
+    "hex27_gradN",
+    "ContactSurfaceSpace",
+    "ContactSide",
+    "OneSidedContact",
+    "OneSidedContactSurfaceSpace",
+    "facet_gap_values",
+    "active_contact_facets",
+}
+
 _SOLVER_EXPORTS = {
     "SparsityPattern",
     "FluxSparseMatrix",
+    "FluxBlockMatrix",
     "coalesce_coo",
     "concat_flux",
     "block_diag_flux",
@@ -227,7 +206,9 @@ _SOLVER_BC_EXPORTS = {
 
 
 def __getattr__(name: str):
-    if name in _SOLVER_EXPORTS:
+    if name in _MESH_EXPORTS:
+        module = importlib.import_module("..mesh", __name__)
+    elif name in _SOLVER_EXPORTS:
         module = importlib.import_module("..solver", __name__)
     elif name in _SOLVER_BC_EXPORTS:
         module = importlib.import_module("..solver.bc", __name__)
@@ -410,6 +391,7 @@ __all__ = [
     "coo_to_csr",
     "SparsityPattern",
     "FluxSparseMatrix",
+    "FluxBlockMatrix",
     "coalesce_coo",
     "concat_flux",
     "block_diag_flux",

fluxfem/core/assembly.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Callable, Literal, Optional, Protocol, TYPE_CHECKING, TypeAlias, TypeVar, Union
+from typing import Any, Callable, Literal, Mapping, Optional, Protocol, TYPE_CHECKING, TypeAlias, TypeVar, Union, cast
 import numpy as np
 import jax
 import jax.numpy as jnp
@@ -10,11 +10,16 @@ from .forms import FormContext
 from .space import FESpaceBase
 
 # Shared call signatures for kernels/forms
-Array = jnp.ndarray
+Array: TypeAlias = jnp.ndarray
 P = TypeVar("P")
 
-Kernel = Callable[[FormContext, P], Array]
+FormKernel: TypeAlias = Callable[[FormContext, P], Array]
+# Form kernels return integrands; element kernels return integrated element arrays.
+Kernel: TypeAlias = Callable[[FormContext, P], Array]
+ResidualInput: TypeAlias = Array | Mapping[str, Array]
+ResidualValue: TypeAlias = Array | Mapping[str, Array]
 ResidualForm = Callable[[FormContext, Array, P], Array]
+ResidualFormLike = Callable[[FormContext, ResidualInput, P], ResidualValue]
 ElementDofMapper = Callable[[Array], Array]
 
 if TYPE_CHECKING:
@@ -107,7 +112,7 @@ class BatchedAssembler:
     def __init__(
         self,
         space: SpaceLike,
-        elem_data: Any,
+        elem_data: FormContext,
         elem_dofs: Array,
         *,
         pattern: SparsityPattern | None = None,
@@ -119,8 +124,8 @@ class BatchedAssembler:
         self.n_ldofs = int(space.n_ldofs)
         self.n_dofs = int(space.n_dofs)
         self.pattern = pattern
-        self._rows = None
-        self._cols = None
+        self._rows: Array | None = None
+        self._cols: Array | None = None
 
     @classmethod
     def from_space(
@@ -135,7 +140,7 @@ class BatchedAssembler:
 
     def make_mask(self, n_active: int) -> Array:
         n_active = max(0, min(int(n_active), self.n_elems))
-        mask = np.zeros((self.n_elems,), dtype=float)
+        mask: np.ndarray = np.zeros((self.n_elems,), dtype=float)
         if n_active:
             mask[:n_active] = 1.0
         return jnp.asarray(mask)
@@ -177,7 +182,7 @@ class BatchedAssembler:
 
     def assemble_bilinear(
         self,
-        form: Kernel[P],
+        form: FormKernel[P],
         params: P,
         *,
         mask: Array | None = None,
@@ -208,7 +213,7 @@ class BatchedAssembler:
 
     def assemble_linear(
         self,
-        form: Kernel[P],
+        form: FormKernel[P],
         params: P,
         *,
         mask: Array | None = None,
@@ -359,7 +364,7 @@ class SpaceLike(FESpaceBase, Protocol):
 
 def assemble_bilinear_dense(
     space: SpaceLike,
-    kernel: Kernel[P],
+    kernel: FormKernel[P],
     params: P,
     *,
     sparse: bool = False,
@@ -412,7 +417,7 @@ def assemble_bilinear_dense(
 
 def assemble_bilinear_form(
     space: SpaceLike,
-    form: Kernel[P],
+    form: FormKernel[P],
     params: P,
     *,
     pattern: SparsityPattern | None = None,
@@ -589,7 +594,7 @@ def assemble_mass_matrix(
 
 def assemble_linear_form(
     space: SpaceLike,
-    form: Kernel[P],
+    form: FormKernel[P],
     params: P,
     *,
     kernel: ElementLinearKernel | None = None,
@@ -670,7 +675,7 @@ def assemble_linear_form(
     return F
 
 
-def assemble_functional(space: SpaceLike, form: Kernel[P], params: P) -> jnp.ndarray:
+def assemble_functional(space: SpaceLike, form: FormKernel[P], params: P) -> jnp.ndarray:
     """
     Assemble scalar functional J = ∫ form(ctx, params) dΩ.
     Expects form(ctx, params) -> (n_q,) or (n_q, 1).
@@ -798,7 +803,7 @@ def assemble_jacobian_elementwise(
     )
     K_flat = jnp.zeros(n_entries, dtype=data.dtype)
     K_flat = jax.lax.scatter_add(K_flat, idx[:, None], data, sdn)
-    return K_flat.reshape(pat.n_dofs, pat.n_dofs)
+    return K_flat.reshape(n_dofs, n_dofs)
 
 
 def assemble_residual_global(
@@ -885,7 +890,7 @@ assemble_residual_elementwise_xla = assemble_residual_elementwise
 
 
 def make_element_bilinear_kernel(
-    form: Kernel[P], params: P, *, jit: bool = True
+    form: FormKernel[P], params: P, *, jit: bool = True
 ) -> ElementBilinearKernel:
     """Element kernel: (ctx) -> Ke."""
 
@@ -900,7 +905,7 @@ def make_element_bilinear_kernel(
 
 
 def make_element_linear_kernel(
-    form: Kernel[P], params: P, *, jit: bool = True
+    form: FormKernel[P], params: P, *, jit: bool = True
 ) -> ElementLinearKernel:
     """Element kernel: (ctx) -> fe."""
 
@@ -945,8 +950,8 @@ def make_element_jacobian_kernel(
 
 
 def element_residual(
-    res_form: ResidualForm[P], ctx: FormContext, u_elem: jnp.ndarray, params: P
-) -> Any:
+    res_form: ResidualFormLike[P], ctx: FormContext, u_elem: ResidualInput, params: P
+) -> ResidualValue:
     """
     Element residual vector r_e(u_e) = sum_q w_q * detJ_q * res_form(ctx, u_e, params).
     Returns shape (n_ldofs,).
@@ -972,8 +977,8 @@ def element_residual(
 
 
 def element_jacobian(
-    res_form: ResidualForm[P], ctx: FormContext, u_elem: jnp.ndarray, params: P
-) -> Any:
+    res_form: ResidualFormLike[P], ctx: FormContext, u_elem: ResidualInput, params: P
+) -> ResidualValue:
     """
     Element Jacobian K_e = d r_e / d u_e (AD via jacfwd), shape (n_ldofs, n_ldofs).
     """
@@ -984,7 +989,7 @@ def element_jacobian(
 
 
 def make_element_kernel(
-    form: Kernel[P] | ResidualForm[P],
+    form: FormKernel[P] | ResidualForm[P],
     params: P,
     *,
     kind: Literal["bilinear", "linear", "residual", "jacobian"],
@@ -999,22 +1004,26 @@ def make_element_kernel(
       - "residual": kernel(ctx, u_elem) -> (n_ldofs,)
       - "jacobian": kernel(u_elem, ctx) -> (n_ldofs, n_ldofs)
     """
-    kind = kind.lower()
+    kind = cast(Literal["bilinear", "linear", "residual", "jacobian"], kind.lower())
     if kind == "bilinear":
-        return make_element_bilinear_kernel(form, params, jit=jit)
+        form_bilinear = cast(FormKernel[P], form)
+        return make_element_bilinear_kernel(form_bilinear, params, jit=jit)
     if kind == "linear":
+        form_linear = cast(FormKernel[P], form)
         def per_element(ctx: FormContext):
-            integrand = form(ctx, params)
-            if getattr(form, "_includes_measure", False):
+            integrand = form_linear(ctx, params)
+            if getattr(form_linear, "_includes_measure", False):
                 return integrand.sum(axis=0)
             wJ = ctx.w * ctx.test.detJ
             return (integrand * wJ[:, None]).sum(axis=0)
 
         return jax.jit(per_element) if jit else per_element
     if kind == "residual":
-        return make_element_residual_kernel(form, params)
+        form_residual = cast(ResidualForm[P], form)
+        return make_element_residual_kernel(form_residual, params)
     if kind == "jacobian":
-        return make_element_jacobian_kernel(form, params)
+        form_residual = cast(ResidualForm[P], form)
+        return make_element_jacobian_kernel(form_residual, params)
     raise ValueError(f"Unknown kernel kind: {kind}")
 
 
@@ -1327,19 +1336,19 @@ def scalar_body_force_form(ctx: FormContext, load: float) -> jnp.ndarray:
     return load * ctx.test.N  # (n_q, n_ldofs)
 
 
-scalar_body_force_form._ff_kind = "linear"
-scalar_body_force_form._ff_domain = "volume"
+scalar_body_force_form._ff_kind = "linear"  # type: ignore[attr-defined]
+scalar_body_force_form._ff_domain = "volume"  # type: ignore[attr-defined]
 
 
-def make_scalar_body_force_form(body_force: Callable[[Array], Array]) -> Kernel[Any]:
+def make_scalar_body_force_form(body_force: Callable[[Array], Array]) -> FormKernel[Any]:
     """
     Build a scalar linear form from a callable f(x_q) -> (n_q,).
     """
     def _form(ctx: FormContext, _params):
         f_q = body_force(ctx.x_q)
         return f_q[..., None] * ctx.test.N
-    _form._ff_kind = "linear"
-    _form._ff_domain = "volume"
+    _form._ff_kind = "linear"  # type: ignore[attr-defined]
+    _form._ff_domain = "volume"  # type: ignore[attr-defined]
     return _form
 
 

fluxfem/core/basis.py

@@ -421,7 +421,7 @@ class HexTriLinearBasis(SmallStrainBMixin, TotalLagrangeBMixin):
 
     def tree_flatten(self):
         children = (self.quad_points, self.quad_weights)
-        aux_data = {}
+        aux_data: dict[str, object] = {}
         return children, aux_data
 
     @classmethod
@@ -569,7 +569,7 @@ class HexSerendipityBasis20(SmallStrainBMixin, TotalLagrangeBMixin):
 
     def tree_flatten(self):
         children = (self.quad_points, self.quad_weights)
-        aux_data = {}
+        aux_data: dict[str, object] = {}
         return children, aux_data
 
     @classmethod

fluxfem/core/context_types.py

@@ -1,36 +1,60 @@
 from __future__ import annotations
 
-from typing import Any, Protocol, TypeAlias, runtime_checkable
+from typing import Any, Mapping, Protocol, TYPE_CHECKING, TypeAlias, runtime_checkable
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
 
 
 @runtime_checkable
 class VolumeContext(Protocol):
     """Minimum interface for volume weak-form evaluation."""
 
-    test: Any
-    trial: Any
-    w: Any
+    test: "FormFieldLike"
+    trial: "FormFieldLike"
+    w: ArrayLike
 
 
 @runtime_checkable
 class SurfaceContext(Protocol):
     """Minimum interface for surface weak-form evaluation."""
 
-    v: Any
-    w: Any
-    detJ: Any
-    normal: Any
+    v: "FormFieldLike"
+    w: ArrayLike
+    detJ: ArrayLike
+    normal: ArrayLike
 
 
 @runtime_checkable
 class FormFieldLike(Protocol):
     """Minimum interface for form fields used in weak-form evaluation."""
 
-    N: Any
-    gradN: Any
-    detJ: Any
+    N: ArrayLike
+    gradN: ArrayLike
+    detJ: ArrayLike
     value_dim: int
+    basis: Any
+
+
+@runtime_checkable
+class WeakFormContext(Protocol):
+    """Context interface used when resolving field references."""
+
+    test: FormFieldLike
+    trial: FormFieldLike
+    v: FormFieldLike
+    unknown: FormFieldLike | None
+    fields: Mapping[str, Any] | None
+    test_fields: Mapping[str, FormFieldLike] | None
+    trial_fields: Mapping[str, FormFieldLike] | None
+    unknown_fields: Mapping[str, FormFieldLike] | None
 
 
-UElement: TypeAlias = Any
+UElement: TypeAlias = ArrayLike | Mapping[str, ArrayLike]
 ParamsLike: TypeAlias = Any

fluxfem/core/mixed_space.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Mapping, Sequence, Callable
+from typing import Any, Callable, Mapping, Sequence, TYPE_CHECKING, TypeAlias, TypeVar, cast
 
 import numpy as np
 import jax.numpy as jnp
@@ -13,6 +13,16 @@ from ..solver.dirichlet import DirichletBC, free_dofs
 from ..solver.sparse import FluxSparseMatrix
 from .space import FESpaceClosure
 
+P = TypeVar("P")
+
+if TYPE_CHECKING:
+    from .assembly import JacobianReturn, LinearReturn
+
+MixedResidualForm: TypeAlias = Callable[
+    [MixedFormContext, Mapping[str, jnp.ndarray], P],
+    Mapping[str, jnp.ndarray],
+]
+
 
 @dataclass(eq=False)
 class MixedFESpace:
@@ -131,13 +141,25 @@ class MixedFESpace:
 
     def get_sparsity_pattern(self, *, with_idx: bool = True):
         from .assembly import make_sparsity_pattern
-        return make_sparsity_pattern(self, with_idx=with_idx)
+        return make_sparsity_pattern(cast(Any, self), with_idx=with_idx)
 
-    def assemble_residual(self, res_form, u, params, **kwargs):
+    def assemble_residual(
+        self,
+        res_form: MixedResidualForm[P],
+        u: Mapping[str, jnp.ndarray] | Sequence[jnp.ndarray] | jnp.ndarray,
+        params: P,
+        **kwargs,
+    ) -> "LinearReturn":
         from .mixed_assembly import assemble_mixed_residual
         return assemble_mixed_residual(self, res_form, u, params, **kwargs)
 
-    def assemble_jacobian(self, res_form, u, params, **kwargs):
+    def assemble_jacobian(
+        self,
+        res_form: MixedResidualForm[P],
+        u: Mapping[str, jnp.ndarray] | Sequence[jnp.ndarray] | jnp.ndarray,
+        params: P,
+        **kwargs,
+    ) -> "JacobianReturn":
         from .mixed_assembly import assemble_mixed_jacobian
         return assemble_mixed_jacobian(self, res_form, u, params, **kwargs)
 
@@ -233,7 +255,7 @@ class MixedProblem:
     pattern: object | None = None
     n_chunks: int | None = None
     pad_trace: bool = False
-    _compiled: object = field(init=False, repr=False)
+    _compiled: Callable[..., Any] = field(init=False, repr=False)
 
     def __post_init__(self):
         if isinstance(self.residuals, MixedWeakForm):
@@ -257,18 +279,30 @@ class MixedProblem:
             def _wrapped(ctx, u_elem, _params):
                 return self._compiled(ctx, u_elem, params(ctx))
 
-            _wrapped._includes_measure = getattr(self._compiled, "_includes_measure", False)
+            _wrapped._includes_measure = getattr(self._compiled, "_includes_measure", False)  # type: ignore[attr-defined]
             return _wrapped, None
         return self._compiled, params
 
-    def assemble_residual(self, u, *, params=None, **kwargs):
+    def assemble_residual(
+        self,
+        u: Mapping[str, jnp.ndarray] | Sequence[jnp.ndarray] | jnp.ndarray,
+        *,
+        params: P | None = None,
+        **kwargs,
+    ) -> "LinearReturn":
         use_params = self.params if params is None else params
         res_form, use_params = self._wrap_params(use_params)
         return self.space.assemble_residual(
             res_form, u, use_params, **self._merge_kwargs(kwargs)
         )
 
-    def assemble_jacobian(self, u, *, params=None, **kwargs):
+    def assemble_jacobian(
+        self,
+        u: Mapping[str, jnp.ndarray] | Sequence[jnp.ndarray] | jnp.ndarray,
+        *,
+        params: P | None = None,
+        **kwargs,
+    ) -> "JacobianReturn":
         use_params = self.params if params is None else params
         res_form, use_params = self._wrap_params(use_params)
         return self.space.assemble_jacobian(

fluxfem/core/mixed_weakform.py

@@ -40,7 +40,7 @@ def _wrap_params(res_form, params):
         def _wrapped(ctx, u_elem, _params):
             return res_form(ctx, u_elem, params(ctx))
 
-        _wrapped._includes_measure = getattr(res_form, "_includes_measure", False)
+        _wrapped._includes_measure = getattr(res_form, "_includes_measure", False)  # type: ignore[attr-defined]
         return _wrapped, None
     return res_form, params