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

fluxfem/solver/dirichlet.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, TYPE_CHECKING, TypeAlias
 
 import numpy as np
 import jax.numpy as jnp
@@ -13,8 +13,16 @@ except Exception:  # pragma: no cover
 
 from .sparse import FluxSparseMatrix, coalesce_coo
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
 
-def _normalize_dirichlet_values(dofs, vals):
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+DirichletLike: TypeAlias = tuple[np.ndarray, np.ndarray]
+
+
+def _normalize_dirichlet_values(dofs: ArrayLike, vals: ArrayLike | None) -> np.ndarray:
     if vals is None:
         return np.zeros(np.asarray(dofs).shape[0], dtype=float)
     arr = np.asarray(vals)
@@ -23,7 +31,7 @@ def _normalize_dirichlet_values(dofs, vals):
     return arr
 
 
-def _normalize_dirichlet(dofs, vals):
+def _normalize_dirichlet(dofs: ArrayLike, vals: ArrayLike | None) -> DirichletLike:
     dir_arr = np.asarray(dofs, dtype=int)
     return dir_arr, _normalize_dirichlet_values(dir_arr, vals)
 
@@ -37,7 +45,7 @@ class CondensedSystem:
     dir_vals: np.ndarray
     n_dofs: int
 
-    def expand(self, u_free, *, fill_dirichlet: bool = True):
+    def expand(self, u_free: ArrayLike, *, fill_dirichlet: bool = True) -> np.ndarray:
         u_full = np.zeros(self.n_dofs, dtype=np.asarray(u_free).dtype)
         u_full[self.free_dofs] = np.asarray(u_free)
         if fill_dirichlet and self.dir_dofs.size:
@@ -61,7 +69,7 @@ class DirichletBC:
         object.__setattr__(self, "vals", vals)
 
     @classmethod
-    def from_boundary_dofs(cls, mesh, predicate, *, values=None, **kwargs):
+    def from_boundary_dofs(cls, mesh, predicate, *, values: ArrayLike | None = None, **kwargs) -> "DirichletBC":
         """
         Build from mesh.boundary_dofs_where predicate.
 
@@ -72,7 +80,16 @@ class DirichletBC:
         return cls(dofs, vals)
 
     @classmethod
-    def from_bbox(cls, mesh, *, mins=None, maxs=None, tol: float = 1e-8, values=None, **kwargs):
+    def from_bbox(
+        cls,
+        mesh,
+        *,
+        mins: ArrayLike | None = None,
+        maxs: ArrayLike | None = None,
+        tol: float = 1e-8,
+        values: ArrayLike | None = None,
+        **kwargs,
+    ) -> "DirichletBC":
         """
         Build from the mesh axis-aligned bounding box.
 
@@ -94,13 +111,13 @@ class DirichletBC:
     def as_tuple(self) -> tuple[np.ndarray, np.ndarray]:
         return self.dofs, self.vals
 
-    def condense_system(self, A, F, *, check: bool = True) -> CondensedSystem:
+    def condense_system(self, A: Any, F: ArrayLike, *, check: bool = True) -> CondensedSystem:
         return condense_dirichlet_system(A, F, self.dofs, self.vals, check=check)
 
-    def enforce_system(self, A, F):
+    def enforce_system(self, A: Any, F: ArrayLike):
         return enforce_dirichlet_system(A, F, self.dofs, self.vals)
 
-    def condense_flux(self, A: FluxSparseMatrix, F):
+    def condense_flux(self, A: FluxSparseMatrix, F: ArrayLike):
         """
         Condense for FluxSparseMatrix and return (K_free, F_free, free_dofs).
         """
@@ -108,16 +125,22 @@ class DirichletBC:
         free = condensed.free_dofs
         return restrict_flux_to_free(A, free), condensed.F, free
 
-    def enforce_flux(self, A: FluxSparseMatrix, F):
+    def enforce_flux(self, A: FluxSparseMatrix, F: ArrayLike):
         return enforce_dirichlet_fluxsparse(A, F, self.dofs, self.vals)
 
-    def split_matrix(self, A, *, n_total: int | None = None):
+    def split_matrix(self, A: Any, *, n_total: int | None = None):
         return split_dirichlet_matrix(A, self.dofs, n_total=n_total)
 
     def free_dofs(self, n_dofs: int) -> np.ndarray:
         return free_dofs(n_dofs, self.dofs)
 
-    def expand_solution(self, u_free, *, free=None, n_total: int | None = None):
+    def expand_solution(
+        self,
+        u_free: ArrayLike,
+        *,
+        free: np.ndarray | None = None,
+        n_total: int | None = None,
+    ) -> np.ndarray:
         if free is None:
             if n_total is None:
                 raise ValueError("n_total is required when free is not provided.")

fluxfem/solver/history.py

@@ -1,7 +1,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, List, Optional
+from typing import Any, List, Optional, TYPE_CHECKING, TypeAlias
+
+import jax.numpy as jnp
+import numpy as np
+
+from .result import SolverResult
+
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
 
 
 @dataclass
@@ -23,9 +35,9 @@ class NewtonIterRecord:
 @dataclass
 class LoadStepResult:
     load_factor: float
-    info: Any
+    info: SolverResult
     solve_time: float
-    u: Any
+    u: ArrayLike
     iter_history: List[NewtonIterRecord] = field(default_factory=list)
     exception: Optional[str] = None
     meta: dict[str, Any] = field(default_factory=dict)

fluxfem/solver/newton.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 import time
 
+from typing import Any, Callable, Mapping, TYPE_CHECKING, TypeAlias
+
 import numpy as np
 import jax
 import jax.numpy as jnp
@@ -8,6 +10,7 @@ import jax.numpy as jnp
 from ..core.assembly import (
     assemble_residual_scatter,
     assemble_jacobian_scatter,
+    ResidualForm,
     make_element_residual_kernel,
     make_element_jacobian_kernel,
     make_sparsity_pattern,
@@ -17,14 +20,22 @@ from .cg import cg_solve, cg_solve_jax
 from .preconditioner import make_block_jacobi_preconditioner
 from .result import SolverResult
 from .sparse import SparsityPattern, FluxSparseMatrix
-from .dirichlet import _normalize_dirichlet
+from .dirichlet import DirichletBC, _normalize_dirichlet
+
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+ExtraTerm: TypeAlias = Callable[[np.ndarray], tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, Mapping[str, Any]] | None]
 
 
 def newton_solve(
     space,
-    res_form,
-    u0,
-    params,
+    res_form: ResidualForm[Any],
+    u0: ArrayLike,
+    params: Any,
     *,
     tol: float = 1e-8,
     atol: float = 0.0,
@@ -32,18 +43,18 @@ def newton_solve(
     linear_solver: str = "spsolve",  # "spsolve", "spdirect_solve_gpu", "cg", "cg_jax", "cg_custom", or "cg_matfree"
     linear_maxiter: int | None = None,
     linear_tol: float | None = None,
-    linear_preconditioner=None,
+    linear_preconditioner: object | None = None,
     matfree_mode: str = "linearize",
-    matfree_cache: dict | None = None,
-    dirichlet=None,
-    callback=None,
+    matfree_cache: dict[str, Any] | None = None,
+    dirichlet: tuple[np.ndarray, np.ndarray] | None = None,
+    callback: Callable[[np.ndarray, SolverResult], Any] | None = None,
     line_search: bool = False,
     max_ls: int = 10,
     ls_c: float = 1e-4,
-    external_vector=None,
-    jacobian_pattern=None,
-    extra_terms=None,
-):
+    external_vector: np.ndarray | None = None,
+    jacobian_pattern: SparsityPattern | None = None,
+    extra_terms: list[ExtraTerm] | None = None,
+) -> tuple[np.ndarray, SolverResult]:
     """
     Gridap-style Newton–Raphson solver on free DOFs only.
 
@@ -61,6 +72,8 @@ def newton_solve(
     """
 
     if dirichlet is not None:
+        if isinstance(dirichlet, DirichletBC):
+            dirichlet = dirichlet.as_tuple()
         dir_dofs, dir_vals = dirichlet
         dir_dofs, dir_vals = _normalize_dirichlet(dir_dofs, dir_vals)
         if dir_vals.ndim == 0:

fluxfem/solver/petsc.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import numpy as np
 import time
 import warnings
-from typing import Any, Callable
+from typing import Any, Callable, TypeAlias
 
 try:
     import scipy.sparse as sp
@@ -12,6 +12,10 @@ except Exception:  # pragma: no cover
 
 from .sparse import FluxSparseMatrix
 
+ArrayLike: TypeAlias = np.ndarray
+MatVec: TypeAlias = Callable[[np.ndarray], np.ndarray]
+SolveInfo: TypeAlias = dict[str, Any]
+
 
 def petsc_is_available() -> bool:
     try:
@@ -31,7 +35,9 @@ def _require_petsc4py():
         raise ImportError("petsc4py is required for PETSc solves. Install with the petsc extra.") from exc
 
 
-def _coo_to_csr(rows, cols, data, n_dofs: int):
+def _coo_to_csr(
+    rows: ArrayLike, cols: ArrayLike, data: ArrayLike, n_dofs: int
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     r = np.asarray(rows, dtype=np.int64)
     c = np.asarray(cols, dtype=np.int64)
     d = np.asarray(data)
@@ -71,7 +77,7 @@ def _infer_n_dofs(K: Any, F: Any | None, n_dofs: int | None) -> int:
     raise ValueError("n_dofs is required when operator shape is not available.")
 
 
-def _matvec_builder(A: Any):
+def _matvec_builder(A: Any) -> MatVec:
     if isinstance(A, FluxSparseMatrix):
         return lambda x: np.asarray(A.matvec(x))
     if hasattr(A, "matvec"):
@@ -89,7 +95,7 @@ def _matvec_builder(A: Any):
     return mv
 
 
-def _diag_from_coo(rows, cols, data, n_dofs: int) -> np.ndarray:
+def _diag_from_coo(rows: ArrayLike, cols: ArrayLike, data: ArrayLike, n_dofs: int) -> np.ndarray:
     r = np.asarray(rows, dtype=np.int64)
     c = np.asarray(cols, dtype=np.int64)
     d = np.asarray(data)
@@ -119,7 +125,7 @@ def _diag_from_operator(A: Any, n_dofs: int) -> np.ndarray:
     raise ValueError("diag0 preconditioner requires access to the matrix diagonal.")
 
 
-def _as_csr(K: Any):
+def _as_csr(K: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
     if isinstance(K, FluxSparseMatrix):
         rows, cols, data, n_dofs = K.to_coo()
         indptr, indices, data = _coo_to_csr(rows, cols, data, int(n_dofs))
@@ -158,7 +164,7 @@ def petsc_solve(
     atol: float | None = None,
     max_it: int | None = None,
     options: dict[str, Any] | None = None,
-) -> np.ndarray | tuple[np.ndarray, dict[str, Any]]:
+) -> np.ndarray:
     """
     Solve K u = F using PETSc.
 
@@ -235,7 +241,7 @@ def petsc_shell_solve(
     options: dict[str, Any] | None = None,
     options_prefix: str | None = "fluxfem_",
     return_info: bool = False,
-) -> np.ndarray:
+) -> np.ndarray | tuple[np.ndarray, SolveInfo]:
     """
     Solve A x = F using PETSc with a matrix-free Shell Mat.
 

fluxfem/solver/preconditioner.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Iterable, Sequence
+from typing import Callable, Iterable, Sequence, TypeAlias
 
 import numpy as np
 import jax.numpy as jnp
@@ -12,6 +12,9 @@ except Exception:  # pragma: no cover
 
 from .sparse import FluxSparseMatrix
 
+ArrayLike: TypeAlias = jnp.ndarray
+Preconditioner: TypeAlias = Callable[[jnp.ndarray], jnp.ndarray]
+
 
 def _extract_block_sizes(
     n: int,
@@ -51,11 +54,11 @@ def _extract_block_sizes(
 
 
 def make_block_jacobi_preconditioner(
-    A,
+    A: FluxSparseMatrix | "jsparse.BCOO",
     *,
     dof_per_node: int | None = None,
     block_sizes: Sequence[int] | None = None,
-):
+) -> Preconditioner:
     """
     Build block Jacobi preconditioner for blocked DOF layouts.
 
@@ -98,7 +101,7 @@ def make_block_jacobi_preconditioner(
     blocks = blocks + 1e-12 * jnp.eye(block_size)[None, :, :]
     inv_blocks = jnp.linalg.inv(blocks)
 
-    def precon(r):
+    def precon(r: jnp.ndarray) -> jnp.ndarray:
         rb = r.reshape((n_block, block_size))
         zb = jnp.einsum("bij,bj->bi", inv_blocks, rb)
         return zb.reshape((-1,))

fluxfem/solver/solve_runner.py

@@ -3,22 +3,34 @@ from __future__ import annotations
 import time
 from dataclasses import dataclass, field
 import warnings
-from typing import Any, Callable, Iterable, List, Sequence
+from typing import Any, Callable, Iterable, List, Sequence, TYPE_CHECKING, TypeAlias
 
 import numpy as np
 import jax.numpy as jnp
 
-from ..core.assembly import assemble_bilinear_form
+from ..core.assembly import FormKernel, ResidualForm, assemble_bilinear_form
 from ..core.solver import spdirect_solve_cpu, spdirect_solve_gpu
 from .cg import cg_solve, cg_solve_jax
 from .petsc import petsc_shell_solve
 from .sparse import FluxSparseMatrix
-from .dirichlet import expand_dirichlet_solution
+from .dirichlet import DirichletBC, expand_dirichlet_solution
 from .newton import newton_solve
 from .result import SolverResult
 from .history import NewtonIterRecord, LoadStepResult
 from ..tools.timer import SectionTimer, NullTimer
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+DirichletLike: TypeAlias = tuple[np.ndarray, np.ndarray]
+ExtraTerm: TypeAlias = Callable[
+    [np.ndarray],
+    tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, dict[str, Any]] | None,
+]
+
 
 @dataclass
 class NonlinearAnalysis:
@@ -46,15 +58,19 @@ class NonlinearAnalysis:
     """
 
     space: Any
-    residual_form: Any
+    residual_form: ResidualForm[Any]
     params: Any
-    base_external_vector: Any | None = None
-    dirichlet: tuple | None = None
-    extra_terms: list[Callable] | None = None
+    base_external_vector: ArrayLike | None = None
+    dirichlet: DirichletLike | None = None
+    extra_terms: list[ExtraTerm] | None = None
     jacobian_pattern: Any | None = None
     dtype: Any = jnp.float64
 
-    def external_for_load(self, load_factor: float):
+    def __post_init__(self) -> None:
+        if isinstance(self.dirichlet, DirichletBC):
+            self.dirichlet = self.dirichlet.as_tuple()
+
+    def external_for_load(self, load_factor: float) -> ArrayLike | None:
         if self.base_external_vector is None:
             return None
         return jnp.asarray(load_factor * self.base_external_vector, dtype=self.dtype)
@@ -105,14 +121,14 @@ class NewtonSolveRunner:
 
     def run(
         self,
-        u0=None,
+        u0: ArrayLike | None = None,
         *,
         load_sequence: Sequence[float] | None = None,
-        newton_callback: Callable | None = None,
+        newton_callback: Callable[[dict[str, Any]], None] | None = None,
         step_callback: Callable[[LoadStepResult], None] | None = None,
         timer: "SectionTimer | None" = None,
         report_timing: bool = True
-    ):
+    ) -> tuple[np.ndarray, list[LoadStepResult]]:
         """
         Execute Newton solves over the configured load schedule.
 
@@ -276,7 +292,9 @@ class NewtonSolveRunner:
         return u, history
 
 
-def _condense_flux_dirichlet(K: FluxSparseMatrix, F, dirichlet):
+def _condense_flux_dirichlet(
+    K: FluxSparseMatrix, F: ArrayLike, dirichlet: DirichletLike
+) -> tuple[Any, np.ndarray, np.ndarray | None, np.ndarray, np.ndarray, np.ndarray]:
     dir_dofs, dir_vals = dirichlet
     dir_arr = np.asarray(dir_dofs, dtype=int)
     dir_vals_arr = np.asarray(dir_vals, dtype=float)
@@ -294,12 +312,12 @@ def _condense_flux_dirichlet(K: FluxSparseMatrix, F, dirichlet):
 
 def solve_nonlinear(
     space,
-    residual_form,
-    params,
+    residual_form: ResidualForm[Any],
+    params: Any,
     *,
-    dirichlet: tuple | None = None,
-    base_external_vector=None,
-    extra_terms=None,
+    dirichlet: DirichletLike | None = None,
+    base_external_vector: ArrayLike | None = None,
+    extra_terms: list[ExtraTerm] | None = None,
     dtype=jnp.float64,
     maxiter: int = 20,
     tol: float = 1e-8,
@@ -314,8 +332,8 @@ def solve_nonlinear(
     ls_c: float = 1e-4,
     n_steps: int = 1,
     jacobian_pattern=None,
-    u0=None,
-):
+    u0: ArrayLike | None = None,
+) -> tuple[np.ndarray, list[LoadStepResult]]:
     """
     Convenience wrapper: build NonlinearAnalysis and run NewtonSolveRunner.
     """
@@ -359,10 +377,10 @@ class LinearAnalysis:
     """
 
     space: Any
-    bilinear_form: Any
+    bilinear_form: FormKernel[Any]
     params: Any
-    base_rhs_vector: Any
-    dirichlet: tuple | None = None
+    base_rhs_vector: ArrayLike
+    dirichlet: DirichletLike | None = None
     pattern: Any | None = None
     dtype: Any = jnp.float64
 
@@ -373,7 +391,7 @@ class LinearAnalysis:
             pattern=self.pattern,
         )
 
-    def rhs_for_load(self, load_factor: float):
+    def rhs_for_load(self, load_factor: float) -> ArrayLike:
         return jnp.asarray(load_factor * self.base_rhs_vector, dtype=self.dtype)
 
 
@@ -437,7 +455,7 @@ class LinearStepResult:
     """
     info: SolverResult
     solve_time: float
-    u: Any
+    u: ArrayLike
 
 
 class LinearSolveRunner:

fluxfem/solver/solver.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any, TYPE_CHECKING, TypeAlias
+
 import numpy as np
 import jax.numpy as jnp
 
@@ -18,6 +20,16 @@ from .dirichlet import (
 from .sparse import FluxSparseMatrix
 from ..core.space import FESpace
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+DirichletLike: TypeAlias = DirichletBC | tuple[np.ndarray, np.ndarray]
+SolveInfo: TypeAlias = dict[str, Any]
+SolveReturn: TypeAlias = tuple[np.ndarray, SolveInfo]
+
 
 class LinearSolver:
     """
@@ -32,7 +44,7 @@ class LinearSolver:
         self.tol = tol
         self.maxiter = maxiter
 
-    def _solve_free(self, A, b):
+    def _solve_free(self, A: Any, b: Any) -> SolveReturn:
         if self.method == "cg":
             x, info = cg_solve_jax(A, b, tol=self.tol, maxiter=self.maxiter)
             return np.asarray(x), {"iters": info.get("iters"), "converged": info.get("converged", True)}
@@ -59,13 +71,13 @@ class LinearSolver:
 
     def solve(
         self,
-        A,
-        b,
+        A: Any,
+        b: Any,
         *,
-        dirichlet=None,
+        dirichlet: DirichletLike | None = None,
         dirichlet_mode: str = "condense",
         n_total: int | None = None,
-    ):
+    ) -> SolveReturn:
         if dirichlet is None:
             return self._solve_free(A, b)
 
@@ -109,8 +121,8 @@ class NonlinearSolver:
     def __init__(
         self,
         space: FESpace,
-        res_form,
-        params,
+        res_form: Any,
+        params: Any,
         *,
         tol: float = 1e-8,
         maxiter: int = 20,
@@ -119,10 +131,10 @@ class NonlinearSolver:
         max_ls: int = 10,
         ls_c: float = 1e-4,
         linear_tol: float | None = None,
-        dirichlet=None,
-        external_vector=None,
+        dirichlet: DirichletLike | None = None,
+        external_vector: ArrayLike | None = None,
         linear_maxiter: int | None = None,
-        jacobian_pattern=None,
+        jacobian_pattern: Any | None = None,
     ):
         self.space = space
         self.res_form = res_form
@@ -139,7 +151,7 @@ class NonlinearSolver:
         self.linear_maxiter = linear_maxiter
         self.jacobian_pattern = jacobian_pattern
 
-    def solve(self, u0):
+    def solve(self, u0: ArrayLike):
         return newton_solve(
             self.space,
             self.res_form,