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

fluxfem/solver/sparse.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any, Iterable, Sequence, TYPE_CHECKING, TypeAlias
 
 import numpy as np
 import jax
@@ -12,6 +13,90 @@ except Exception:  # pragma: no cover
     sp = None
 
 
+if TYPE_CHECKING:
+    from jax import Array as JaxArray
+
+    ArrayLike: TypeAlias = np.ndarray | JaxArray
+else:
+    ArrayLike: TypeAlias = np.ndarray
+COOTuple: TypeAlias = tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, int]
+
+
+def coalesce_coo(
+    rows: ArrayLike, cols: ArrayLike, data: ArrayLike
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Sum duplicate COO entries by sorting (CPU-friendly).
+    Returns (rows_u, cols_u, data_u) as NumPy arrays.
+    """
+    r = np.asarray(rows, dtype=np.int64)
+    c = np.asarray(cols, dtype=np.int64)
+    d = np.asarray(data)
+    if r.size == 0:
+        return r, c, d
+    order = np.lexsort((c, r))
+    r_s = r[order]
+    c_s = c[order]
+    d_s = d[order]
+    new_group = np.ones(r_s.size, dtype=bool)
+    new_group[1:] = (r_s[1:] != r_s[:-1]) | (c_s[1:] != c_s[:-1])
+    starts = np.nonzero(new_group)[0]
+    r_u = r_s[starts]
+    c_u = c_s[starts]
+    d_u = np.add.reduceat(d_s, starts)
+    return r_u, c_u, d_u
+
+
+def _normalize_flux_mats(mats: Sequence["FluxSparseMatrix"]) -> tuple["FluxSparseMatrix", ...]:
+    if len(mats) == 1 and isinstance(mats[0], (list, tuple)):
+        mats = tuple(mats[0])
+    if not mats:
+        raise ValueError("At least one FluxSparseMatrix is required.")
+    return mats
+
+
+def concat_flux(*mats: "FluxSparseMatrix", n_dofs: int | None = None) -> "FluxSparseMatrix":
+    """
+    Concatenate COO entries from multiple FluxSparseMatrix objects.
+    All matrices must share the same n_dofs unless n_dofs is provided.
+    """
+    mats = _normalize_flux_mats(mats)
+    if n_dofs is None:
+        n_dofs = int(mats[0].n_dofs)
+        for mat in mats[1:]:
+            if int(mat.n_dofs) != n_dofs:
+                raise ValueError("All matrices must share n_dofs for concat_flux.")
+    rows_list = [np.asarray(mat.pattern.rows, dtype=np.int32) for mat in mats]
+    cols_list = [np.asarray(mat.pattern.cols, dtype=np.int32) for mat in mats]
+    data_list = [np.asarray(mat.data) for mat in mats]
+    rows = np.concatenate(rows_list) if rows_list else np.asarray([], dtype=np.int32)
+    cols = np.concatenate(cols_list) if cols_list else np.asarray([], dtype=np.int32)
+    data = np.concatenate(data_list) if data_list else np.asarray([], dtype=float)
+    return FluxSparseMatrix(rows, cols, data, int(n_dofs))
+
+
+def block_diag_flux(*mats: "FluxSparseMatrix") -> "FluxSparseMatrix":
+    """Block-diagonal concatenation for FluxSparseMatrix objects."""
+    mats = _normalize_flux_mats(mats)
+    rows_out = []
+    cols_out = []
+    data_out = []
+    offset = 0
+    for mat in mats:
+        rows = np.asarray(mat.pattern.rows, dtype=np.int32)
+        cols = np.asarray(mat.pattern.cols, dtype=np.int32)
+        data = np.asarray(mat.data)
+        if rows.size:
+            rows_out.append(rows + offset)
+            cols_out.append(cols + offset)
+            data_out.append(data)
+        offset += int(mat.n_dofs)
+    rows = np.concatenate(rows_out) if rows_out else np.asarray([], dtype=np.int32)
+    cols = np.concatenate(cols_out) if cols_out else np.asarray([], dtype=np.int32)
+    data = np.concatenate(data_out) if data_out else np.asarray([], dtype=float)
+    return FluxSparseMatrix(rows, cols, data, int(offset))
+
+
 @jax.tree_util.register_pytree_node_class
 @dataclass(frozen=True)
 class SparsityPattern:
@@ -80,7 +165,14 @@ class FluxSparseMatrix:
     - data stores the numeric values for the current nonlinear iterate
     """
 
-    def __init__(self, rows_or_pattern, cols=None, data=None, n_dofs: int | None = None):
+    def __init__(
+        self,
+        rows_or_pattern: SparsityPattern | ArrayLike,
+        cols: ArrayLike | None = None,
+        data: ArrayLike | None = None,
+        n_dofs: int | None = None,
+        meta: dict | None = None,
+    ):
         # New signature: FluxSparseMatrix(pattern, data)
         if isinstance(rows_or_pattern, SparsityPattern):
             pattern = rows_or_pattern
@@ -88,15 +180,22 @@ class FluxSparseMatrix:
             values = jnp.asarray(values)
         else:
             # Legacy signature: FluxSparseMatrix(rows, cols, data, n_dofs)
-            r_np = np.asarray(rows_or_pattern, dtype=np.int32)
-            c_np = np.asarray(cols, dtype=np.int32)
-            diag_idx_np = np.nonzero(r_np == c_np)[0].astype(np.int32)
+            r_j = jnp.asarray(rows_or_pattern, dtype=jnp.int32)
+            c_j = jnp.asarray(cols, dtype=jnp.int32)
+            is_tracer = isinstance(rows_or_pattern, jax.core.Tracer) or isinstance(cols, jax.core.Tracer)
+            diag_idx_j = None
+            if not is_tracer:
+                diag_idx_j = jnp.nonzero(r_j == c_j)[0].astype(jnp.int32)
+            if n_dofs is None:
+                if is_tracer:
+                    raise ValueError("n_dofs must be provided when constructing FluxSparseMatrix under JIT.")
+                n_dofs = int(np.asarray(cols).max()) + 1
             pattern = SparsityPattern(
-                rows=jnp.asarray(r_np),
-                cols=jnp.asarray(c_np),
-                n_dofs=int(n_dofs) if n_dofs is not None else int(c_np.max()) + 1,
+                rows=r_j,
+                cols=c_j,
+                n_dofs=int(n_dofs) if n_dofs is not None else int(np.asarray(cols).max()) + 1,
                 idx=None,
-                diag_idx=jnp.asarray(diag_idx_np),
+                diag_idx=diag_idx_j,
             )
             values = jnp.asarray(data)
 
@@ -105,27 +204,42 @@ class FluxSparseMatrix:
         self.cols = pattern.cols
         self.n_dofs = int(pattern.n_dofs)
         self.data = values
+        self.meta = dict(meta) if meta is not None else None
 
     @classmethod
-    def from_bilinear(cls, coo_tuple):
+    def from_bilinear(cls, coo_tuple: COOTuple) -> "FluxSparseMatrix":
         """Construct from assemble_bilinear_dense(..., sparse=True)."""
         rows, cols, data, n_dofs = coo_tuple
         return cls(rows, cols, data, n_dofs)
 
     @classmethod
-    def from_linear(cls, coo_tuple):
+    def from_linear(cls, coo_tuple: tuple[jnp.ndarray, jnp.ndarray, int]) -> "FluxSparseMatrix":
         """Construct from assemble_linear_form(..., sparse=True) (matrix interpretation only)."""
         rows, data, n_dofs = coo_tuple
         cols = jnp.zeros_like(rows)
         return cls(rows, cols, data, n_dofs)
 
-    def with_data(self, data):
+    def with_data(self, data: ArrayLike) -> "FluxSparseMatrix":
         """Return a new FluxSparseMatrix sharing the same pattern with updated data."""
-        return FluxSparseMatrix(self.pattern, data)
+        return FluxSparseMatrix(self.pattern, data, meta=self.meta)
+
+    def add_dense(self, dense: ArrayLike) -> "FluxSparseMatrix":
+        """Return a new FluxSparseMatrix with dense entries added on the pattern."""
+        dense_vals = jnp.asarray(dense)[self.pattern.rows, self.pattern.cols]
+        return FluxSparseMatrix(self.pattern, self.data + dense_vals)
 
-    def to_coo(self):
+    def to_coo(self) -> COOTuple:
         return self.pattern.rows, self.pattern.cols, self.data, self.pattern.n_dofs
 
+    @property
+    def nnz(self) -> int:
+        return int(self.data.shape[0])
+
+    def coalesce(self) -> "FluxSparseMatrix":
+        """Return a new FluxSparseMatrix with duplicate entries summed."""
+        rows_u, cols_u, data_u = coalesce_coo(self.pattern.rows, self.pattern.cols, self.data)
+        return FluxSparseMatrix(rows_u, cols_u, data_u, self.pattern.n_dofs)
+
     def to_csr(self):
         if sp is None:
             raise ImportError("scipy is required for to_csr()")
@@ -143,7 +257,7 @@ class FluxSparseMatrix:
         d = np.array(self.data, copy=True)
         return sp.csr_matrix((d, (r, c)), shape=(self.pattern.n_dofs, self.pattern.n_dofs))
 
-    def to_dense(self):
+    def to_dense(self) -> jnp.ndarray:
         # small debug helper
         dense = jnp.zeros((self.pattern.n_dofs, self.pattern.n_dofs), dtype=self.data.dtype)
         dense = dense.at[self.pattern.rows, self.pattern.cols].add(self.data)
@@ -158,7 +272,7 @@ class FluxSparseMatrix:
         idx = jnp.stack([self.pattern.rows, self.pattern.cols], axis=-1)
         return jsparse.BCOO((self.data, idx), shape=(self.pattern.n_dofs, self.pattern.n_dofs))
 
-    def matvec(self, x):
+    def matvec(self, x: ArrayLike) -> jnp.ndarray:
         """Compute y = A x in JAX (iterative solvers)."""
         xj = jnp.asarray(x)
         contrib = self.data * xj[self.pattern.cols]
@@ -167,6 +281,26 @@ class FluxSparseMatrix:
         out = jnp.zeros(self.pattern.n_dofs, dtype=contrib.dtype)
         return out.at[self.pattern.rows].add(contrib)
 
+    def as_cg_operator(
+        self,
+        *,
+        matvec: str = "flux",
+        preconditioner=None,
+        solver: str = "cg",
+        dof_per_node: int | None = None,
+        block_sizes=None,
+    ):
+        from .cg import build_cg_operator
+
+        return build_cg_operator(
+            self,
+            matvec=matvec,
+            preconditioner=preconditioner,
+            solver=solver,
+            dof_per_node=dof_per_node,
+            block_sizes=block_sizes,
+        )
+
     def diag(self):
         """Diagonal entries aggregated for Jacobi preconditioning."""
         if self.pattern.diag_idx is not None:

fluxfem/tools/jit.py

@@ -1,10 +1,22 @@
+from __future__ import annotations
+
 import jax
 
-from ..core.assembly import assemble_residual, assemble_jacobian
+from typing import Callable, TypeVar
+
+from ..core.assembly import JacobianReturn, LinearReturn, ResidualForm, assemble_jacobian, assemble_residual
 from ..core.space import FESpace
 
+P = TypeVar("P")
+
 
-def make_jitted_residual(space: FESpace, res_form, params, *, sparse: bool = False):
+def make_jitted_residual(
+    space: FESpace,
+    res_form: ResidualForm[P],
+    params: P,
+    *,
+    sparse: bool = False,
+) -> Callable[[jax.Array], LinearReturn]:
     """
     Create a jitted residual assembler: u -> R(u).
     params and space are closed over.
@@ -13,7 +25,7 @@ def make_jitted_residual(space: FESpace, res_form, params, *, sparse: bool = Fal
     params_jax = params
 
     @jax.jit
-    def residual(u):
+    def residual(u: jax.Array) -> LinearReturn:
         return assemble_residual(space_jax, res_form, u, params_jax, sparse=sparse)
 
     return residual
@@ -21,12 +33,12 @@ def make_jitted_residual(space: FESpace, res_form, params, *, sparse: bool = Fal
 
 def make_jitted_jacobian(
     space: FESpace,
-    res_form,
-    params,
+    res_form: ResidualForm[P],
+    params: P,
     *,
     sparse: bool = False,
     return_flux_matrix: bool = False,
-):
+) -> Callable[[jax.Array], JacobianReturn]:
     """
     Create a jitted Jacobian assembler: u -> J(u).
     params and space are closed over.
@@ -35,7 +47,7 @@ def make_jitted_jacobian(
     params_jax = params
 
     @jax.jit
-    def jacobian(u):
+    def jacobian(u: jax.Array) -> JacobianReturn:
         return assemble_jacobian(
             space_jax,
             res_form,

fluxfem/tools/timer.py

@@ -5,13 +5,15 @@ from contextlib import AbstractContextManager
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Callable, DefaultDict, Dict, Iterator, List, Optional
+from typing import Any, Callable, DefaultDict, Dict, Iterator, List, Optional, TypeAlias
 
 import logging
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+PlotResult: TypeAlias = tuple[Any, Any]
+
 
 @dataclass
 class SectionStats:
@@ -30,7 +32,7 @@ class BaseTimer(ABC):
 
 class NullTimer(BaseTimer):
     @contextmanager
-    def section(self, name: str):
+    def section(self, name: str) -> Iterator[None]:
         yield
 
 
@@ -85,7 +87,7 @@ class SectionTimer(BaseTimer):
             self._records[full_name].append(duration)
             self._stack.pop()
 
-    def wrap(self, name: str):
+    def wrap(self, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
         """
         Decorator form of :meth:`section`.
 
@@ -97,8 +99,8 @@ class SectionTimer(BaseTimer):
                 ...
         """
 
-        def _decorator(func):
-            def _wrapper(*args, **kwargs):
+        def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+            def _wrapper(*args: Any, **kwargs: Any) -> Any:
                 with self.section(name):
                     return func(*args, **kwargs)
             return _wrapper
@@ -211,7 +213,7 @@ class SectionTimer(BaseTimer):
         self,
         sort_by: str = "total",
         descending: bool = True,
-        logger_instance=None,
+        logger_instance: logging.Logger | None = None,
     ) -> str:
         stats = self.summary(sort_by=sort_by, descending=descending)
         if not stats:
@@ -231,7 +233,7 @@ class SectionTimer(BaseTimer):
 
     def plot_bar(
         self,
-        ax=None,
+        ax: Any | None = None,
         sort_by: str = "total",
         value: str = "total",
         descending: bool = True,
@@ -240,7 +242,7 @@ class SectionTimer(BaseTimer):
         stacked_nested: bool = False,
         moving_average: bool = False,
         use_self_time: bool = False,
-    ):
+    ) -> PlotResult:
         """
         Plot timing results as a horizontal bar chart without relying on pyplot state.
 
@@ -417,7 +419,7 @@ class SectionTimer(BaseTimer):
 
     def plot_pie(
         self,
-        ax=None,
+        ax: Any | None = None,
         sort_by: str = "total",
         value: str = "total",
         descending: bool = True,
@@ -430,7 +432,7 @@ class SectionTimer(BaseTimer):
         show_total: bool = True,
         moving_average: bool = False,
         use_self_time: bool = False,
-    ):
+    ) -> PlotResult:
         """
         Plot timing results as a pie chart to show relative time share.
 
@@ -552,7 +554,7 @@ class SectionTimer(BaseTimer):
 
     def plot(
         self,
-        ax=None,
+        ax: Any | None = None,
         sort_by: str = "total",
         value: str = "total",
         descending: bool = True,
@@ -563,7 +565,7 @@ class SectionTimer(BaseTimer):
         moving_average: bool = False,
         use_self_time: bool = False,
         **kwargs,
-    ):
+    ) -> PlotResult:
         """
         Plot timing results choosing between pie (default) or bar chart.
 

fluxfem/tools/visualizer.py

@@ -3,14 +3,14 @@ from typing import Mapping
 
 import numpy as np
 
-from ..mesh import HexMesh, TetMesh
+from ..mesh import BaseMesh, HexMesh, TetMesh
 
 
 VTK_HEXAHEDRON = 12
 VTK_TETRA = 10
 
 
-def _cell_type_for_mesh(mesh):
+def _cell_type_for_mesh(mesh: BaseMesh) -> int:
     if isinstance(mesh, HexMesh):
         return VTK_HEXAHEDRON
     if isinstance(mesh, TetMesh):
@@ -25,7 +25,13 @@ def _write_dataarray(name: str, data: np.ndarray, ncomp: int = 1) -> str:
     return f'<DataArray type="Float32" Name="{name}" format="ascii"{comp_attr}>{values}</DataArray>'
 
 
-def write_vtu(mesh, filepath: str, *, point_data: Mapping[str, np.ndarray] | None = None, cell_data: Mapping[str, np.ndarray] | None = None):
+def write_vtu(
+    mesh: BaseMesh,
+    filepath: str,
+    *,
+    point_data: Mapping[str, np.ndarray] | None = None,
+    cell_data: Mapping[str, np.ndarray] | None = None,
+) -> None:
     """
     Write an UnstructuredGrid VTU for HexMesh or TetMesh.
     point_data/cell_data: dict name -> ndarray. Point data length must match n_points;
@@ -87,7 +93,13 @@ def write_vtu(mesh, filepath: str, *, point_data: Mapping[str, np.ndarray] | Non
         f.write("\n".join(lines))
 
 
-def write_displacement_vtu(mesh, u, filepath: str, *, name: str = "displacement"):
+def write_displacement_vtu(
+    mesh: BaseMesh,
+    u: np.ndarray,
+    filepath: str,
+    *,
+    name: str = "displacement",
+) -> None:
     """
     Convenience wrapper: reshape displacement vector to point data and write VTU.
     Assumes 3 dof/node ordering [u0,v0,w0, u1,v1,w1, ...].