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

fluxfem/core/space.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import operator
 from dataclasses import dataclass, field
-from typing import Any, Callable, Protocol, TYPE_CHECKING, TypeVar
+from typing import Any, Callable, Protocol, TYPE_CHECKING, TypeAlias, TypeVar, cast
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -16,15 +16,39 @@ if TYPE_CHECKING:
         ElementLinearKernel,
         ElementResidualKernel,
         Kernel,
+        FormKernel,
+        BilinearReturn,
+        JacobianReturn,
+        LinearReturn,
+        MassReturn,
         ResidualForm,
     )
+    from .weakform import BilinearForm, LinearForm
+    from ..solver import FluxSparseMatrix, SparsityPattern
 else:
     Kernel = Callable[..., Any]
+    FormKernel = Callable[..., Any]
     ResidualForm = Callable[..., Any]
     ElementBilinearKernel = Callable[..., Any]
     ElementLinearKernel = Callable[..., Any]
     ElementResidualKernel = Callable[..., Any]
     ElementJacobianKernel = Callable[..., Any]
+    BilinearReturn = Any
+    JacobianReturn = Any
+    LinearReturn = Any
+    MassReturn = Any
+    BilinearForm = Any
+    LinearForm = Any
+    FluxSparseMatrix = Any
+    SparsityPattern = Any
+
+KernelCacheKey: TypeAlias = tuple[str, int, int, bool]
+KernelCache: TypeAlias = dict[
+    KernelCacheKey,
+    ElementBilinearKernel | ElementLinearKernel | ElementResidualKernel | ElementJacobianKernel,
+]
+PatternCache: TypeAlias = dict[bool, SparsityPattern]
+PatternLike: TypeAlias = str | SparsityPattern | None
 
 _WARNED_UNTAGGED_KERNELS: set[int] = set()
 
@@ -93,8 +117,12 @@ class FESpaceBase(Protocol):
     """
     elem_dofs: jnp.ndarray
     value_dim: int
-    n_dofs: int
-    n_ldofs: int
+
+    @property
+    def n_dofs(self) -> int: ...
+
+    @property
+    def n_ldofs(self) -> int: ...
 
     def build_form_contexts(self, dep: jnp.ndarray | None = None) -> FormContext: ...
 
@@ -120,8 +148,8 @@ class FESpaceClosure:
     _n_dofs: int | None = None
     _n_ldofs: int | None = None
     data: SpaceData | None = None
-    _pattern_cache: dict[bool, object] = field(default_factory=dict, repr=False)
-    _kernel_cache: dict[tuple, object] = field(default_factory=dict, repr=False)
+    _pattern_cache: PatternCache = field(default_factory=dict, repr=False)
+    _kernel_cache: KernelCache = field(default_factory=dict, repr=False)
     _elem_rows_cache: jnp.ndarray | None = field(default=None, repr=False)
 
     def __post_init__(self):
@@ -232,7 +260,7 @@ class FESpaceClosure:
             block_sizes=block_sizes,
         )
 
-    def _kernel_cache_key(self, kind: str, form, params, *, jit: bool):
+    def _kernel_cache_key(self, kind: str, form, params, *, jit: bool) -> KernelCacheKey | None:
         if not jit:
             return None
         form_key = getattr(form, "__wrapped__", form)
@@ -256,15 +284,15 @@ class FESpaceClosure:
 
     def assemble(
         self,
-        form,
-        params=None,
+        form: FormKernel[P] | BilinearForm | LinearForm,
+        params: P | None = None,
         *,
         kind: str | None = None,
         n_chunks: int | None = None,
         dep: jnp.ndarray | None = None,
         jit: bool = True,
-        pattern: str | object | None = "auto",
-        kernel=None,
+        pattern: PatternLike = "auto",
+        kernel: ElementBilinearKernel | ElementLinearKernel | None = None,
         **kwargs,
     ):
         """
@@ -324,13 +352,14 @@ class FESpaceClosure:
                     jit=jit,
                     maker=make_element_bilinear_kernel,
                 )
+            form_kernel = cast(FormKernel, form)
             pattern_use = None
             if pattern == "auto":
                 pattern_use = self.get_sparsity_pattern(with_idx=True)
             else:
                 pattern_use = pattern
             return self.assemble_bilinear_form(
-                form,
+                form_kernel,
                 params,
                 n_chunks=n_chunks,
                 dep=dep,
@@ -347,8 +376,9 @@ class FESpaceClosure:
                     jit=jit,
                     maker=make_element_linear_kernel,
                 )
+            form_kernel = cast(FormKernel, form)
             return self.assemble_linear_form(
-                form,
+                form_kernel,
                 params,
                 n_chunks=n_chunks,
                 dep=dep,
@@ -360,14 +390,14 @@ class FESpaceClosure:
     # --- Thin wrappers over functional assembly APIs (kept functional for JAX friendliness) ---
     def assemble_bilinear_form(
         self,
-        form: Kernel[P],
+        form: FormKernel[P],
         params: P,
         *,
         n_chunks: int | None = None,
         dep: jnp.ndarray | None = None,
         kernel: ElementBilinearKernel | None = None,
         **kwargs,
-    ):
+    ) -> FluxSparseMatrix:
         """Assemble bilinear form; kernel(ctx) -> (n_ldofs, n_ldofs) if provided."""
         from .assembly import assemble_bilinear_form
         if "pattern" not in kwargs or kwargs.get("pattern") is None:
@@ -378,29 +408,39 @@ class FESpaceClosure:
 
     def assemble_linear_form(
         self,
-        form: Kernel[P],
+        form: FormKernel[P],
         params: P,
         *,
         n_chunks: int | None = None,
         dep: jnp.ndarray | None = None,
         kernel: ElementLinearKernel | None = None,
         **kwargs,
-    ):
+    ) -> LinearReturn:
         """Assemble linear form; kernel(ctx) -> (n_ldofs,) if provided."""
         from .assembly import assemble_linear_form
         return assemble_linear_form(
             self, form, params, n_chunks=n_chunks, dep=dep, kernel=kernel, **kwargs
         )
 
-    def assemble_functional(self, form, params):
+    def assemble_functional(self, form: FormKernel[P], params: P) -> jnp.ndarray:
         from .assembly import assemble_functional
         return assemble_functional(self, form, params)
 
-    def assemble_mass_matrix(self, *, n_chunks=None, **kwargs):
+    def assemble_mass_matrix(
+        self,
+        *,
+        n_chunks: int | None = None,
+        **kwargs,
+    ) -> MassReturn:
         from .assembly import assemble_mass_matrix
         return assemble_mass_matrix(self, n_chunks=n_chunks, **kwargs)
 
-    def assemble_bilinear_dense(self, kernel, params, **kwargs):
+    def assemble_bilinear_dense(
+        self,
+        kernel: FormKernel[P],
+        params: P,
+        **kwargs,
+    ) -> BilinearReturn:
         from .assembly import assemble_bilinear_dense
         return assemble_bilinear_dense(self, kernel, params, **kwargs)
 
@@ -412,7 +452,7 @@ class FESpaceClosure:
         *,
         kernel: ElementResidualKernel | None = None,
         **kwargs,
-    ):
+    ) -> LinearReturn:
         """Assemble residual; kernel(ctx, u_elem) -> (n_ldofs,) if provided."""
         from .assembly import assemble_residual
         return assemble_residual(self, res_form, u, params, kernel=kernel, **kwargs)
@@ -425,7 +465,7 @@ class FESpaceClosure:
         *,
         kernel: ElementJacobianKernel | None = None,
         **kwargs,
-    ):
+    ) -> JacobianReturn:
         """Assemble Jacobian; kernel(u_elem, ctx) -> (n_ldofs, n_ldofs) if provided."""
         from .assembly import assemble_jacobian
         return assemble_jacobian(self, res_form, u, params, kernel=kernel, **kwargs)
@@ -585,7 +625,7 @@ def make_tet10_space_pytree(
     """Create a pytree quadratic tet space (10-node elements)."""
     basis = make_tet10_basis_pytree(intorder)
     element = None if dim == 1 else ElementVector(dim)
-    return make_space_pytree(mesh, basis, element)
+    return make_space_pytree(cast(BaseMeshPytree, mesh), basis, element)
 
 
 def make_hex_space(mesh: HexMesh, dim: int = 1, intorder: int = 2) -> FESpace:
@@ -601,7 +641,7 @@ def make_hex_space_pytree(
     """Create a pytree trilinear hex space (8-node elements)."""
     basis = make_hex_basis_pytree(intorder)
     element = None if dim == 1 else ElementVector(dim)
-    return make_space_pytree(mesh, basis, element)
+    return make_space_pytree(cast(BaseMeshPytree, mesh), basis, element)
 
 
 def make_hex20_space(
@@ -619,7 +659,7 @@ def make_hex20_space_pytree(
     """Create a pytree serendipity hex space (20-node elements)."""
     basis = make_hex20_basis_pytree(intorder)
     element = None if dim == 1 else ElementVector(dim)
-    return make_space_pytree(mesh, basis, element)
+    return make_space_pytree(cast(BaseMeshPytree, mesh), basis, element)
 
 
 def make_hex27_space(
@@ -637,14 +677,14 @@ def make_hex27_space_pytree(
     """Create a pytree triquadratic hex space (27-node elements)."""
     basis = make_hex27_basis_pytree(intorder)
     element = None if dim == 1 else ElementVector(dim)
-    return make_space_pytree(mesh, basis, element)
+    return make_space_pytree(cast(BaseMeshPytree, mesh), basis, element)
 
 
 def make_tet_space(mesh: TetMesh, dim: int = 1, intorder: int = 2) -> FESpace:
     """Create a linear or quadratic tet space based on mesh nodes."""
     n_nodes = mesh.conn.shape[1]
     if n_nodes == 10:
-        basis = make_tet10_basis(intorder if intorder > 1 else 2)
+        basis: Basis3D = make_tet10_basis(intorder if intorder > 1 else 2)
     else:
         basis = make_tet_basis(intorder)
     element = None if dim == 1 else ElementVector(dim)
@@ -657,8 +697,8 @@ def make_tet_space_pytree(
     """Create a pytree linear or quadratic tet space based on mesh nodes."""
     n_nodes = mesh.conn.shape[1]
     if n_nodes == 10:
-        basis = make_tet10_basis_pytree(intorder if intorder > 1 else 2)
+        basis: Basis3D = make_tet10_basis_pytree(intorder if intorder > 1 else 2)
     else:
         basis = make_tet_basis_pytree(intorder)
     element = None if dim == 1 else ElementVector(dim)
-    return make_space_pytree(mesh, basis, element)
+    return make_space_pytree(cast(BaseMeshPytree, mesh), basis, element)