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

fluxfem/solver/block_matrix.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Mapping, Sequence
+from collections.abc import Mapping as AbcMapping
+from typing import Any, Iterator, Mapping, Sequence, TypeAlias
 
 import numpy as np
 
@@ -12,12 +13,16 @@ except Exception:  # pragma: no cover
 from .block_system import split_block_matrix
 from .sparse import FluxSparseMatrix
 
+MatrixLike: TypeAlias = Any
+FieldKey: TypeAlias = str | int
+BlockMap: TypeAlias = dict[FieldKey, dict[FieldKey, MatrixLike]]
 
-def diag(**blocks):
+
+def diag(**blocks: MatrixLike) -> dict[str, MatrixLike]:
     return dict(blocks)
 
 
-def _infer_sizes_from_diag(diag_blocks):
+def _infer_sizes_from_diag(diag_blocks: Mapping[FieldKey, MatrixLike]) -> dict[FieldKey, int]:
     sizes = {}
     for name, blk in diag_blocks.items():
         if isinstance(blk, FluxSparseMatrix):
@@ -35,7 +40,19 @@ def _infer_sizes_from_diag(diag_blocks):
     return sizes
 
 
-def _add_blocks(a, b):
+def _infer_format(blocks: AbcMapping[FieldKey, AbcMapping[FieldKey, MatrixLike]], fmt: str) -> str:
+    if fmt != "auto":
+        return fmt
+    for row in blocks.values():
+        for blk in row.values():
+            if isinstance(blk, FluxSparseMatrix):
+                return "flux"
+            if sp is not None and sp.issparse(blk):
+                return "csr"
+    return "dense"
+
+
+def _add_blocks(a: MatrixLike | None, b: MatrixLike | None) -> MatrixLike | None:
     if a is None:
         return b
     if b is None:
@@ -53,7 +70,7 @@ def _add_blocks(a, b):
     return np.asarray(a) + np.asarray(b)
 
 
-def _transpose_block(block, rule: str):
+def _transpose_block(block: MatrixLike, rule: str) -> MatrixLike:
     if isinstance(block, FluxSparseMatrix):
         if sp is None:
             raise ImportError("scipy is required to transpose FluxSparseMatrix blocks.")
@@ -67,17 +84,146 @@ def _transpose_block(block, rule: str):
     return out
 
 
+class FluxBlockMatrix(AbcMapping[FieldKey, Mapping[FieldKey, MatrixLike]]):
+    """
+    Lazy block-matrix container that assembles on demand.
+    """
+
+    def __init__(
+        self,
+        blocks: BlockMap,
+        *,
+        sizes: Mapping[FieldKey, int],
+        order: Sequence[FieldKey] | None = None,
+        symmetric: bool = False,
+        transpose_rule: str = "T",
+    ) -> None:
+        self._blocks = blocks
+        self.sizes = {name: int(size) for name, size in sizes.items()}
+        self.field_order = tuple(order) if order is not None else tuple(self.sizes.keys())
+        self.symmetric = bool(symmetric)
+        self.transpose_rule = transpose_rule
+        for name in self.field_order:
+            if name not in self.sizes:
+                raise KeyError(f"Missing size for field '{name}'")
+
+    def __getitem__(self, key: FieldKey) -> Mapping[FieldKey, MatrixLike]:
+        return self._blocks[key]
+
+    def __iter__(self) -> Iterator[FieldKey]:
+        return iter(self._blocks)
+
+    def __len__(self) -> int:
+        return len(self._blocks)
+
+    @property
+    def blocks(self) -> BlockMap:
+        return self._blocks
+
+    def assemble(self, *, format: str = "flux") -> MatrixLike:
+        if format not in {"auto", "flux", "csr", "dense"}:
+            raise ValueError("format must be one of: auto, flux, csr, dense")
+        use_format = _infer_format(self._blocks, format)
+
+        offsets = {}
+        offset = 0
+        for name in self.field_order:
+            size = int(self.sizes[name])
+            offsets[name] = offset
+            offset += size
+        n_total = offset
+
+        def _block_shape(name_i: FieldKey, name_j: FieldKey) -> tuple[int, int]:
+            return (int(self.sizes[name_i]), int(self.sizes[name_j]))
+
+        if use_format == "flux":
+            rows_list = []
+            cols_list = []
+            data_list = []
+            for name_i in self.field_order:
+                row_blocks = self._blocks.get(name_i, {})
+                for name_j in self.field_order:
+                    blk = row_blocks.get(name_j)
+                    if blk is None:
+                        continue
+                    shape = _block_shape(name_i, name_j)
+                    if isinstance(blk, FluxSparseMatrix):
+                        if shape[0] != shape[1] or int(blk.n_dofs) != shape[0]:
+                            raise ValueError(f"Block {name_i},{name_j} has incompatible FluxSparseMatrix size")
+                        r = np.asarray(blk.pattern.rows, dtype=np.int64)
+                        c = np.asarray(blk.pattern.cols, dtype=np.int64)
+                        d = np.asarray(blk.data)
+                    elif sp is not None and sp.issparse(blk):
+                        coo = blk.tocoo()
+                        r = np.asarray(coo.row, dtype=np.int64)
+                        c = np.asarray(coo.col, dtype=np.int64)
+                        d = np.asarray(coo.data)
+                        if coo.shape != shape:
+                            raise ValueError(f"Block {name_i},{name_j} has shape {coo.shape}, expected {shape}")
+                    else:
+                        arr = np.asarray(blk)
+                        if arr.shape != shape:
+                            raise ValueError(f"Block {name_i},{name_j} has shape {arr.shape}, expected {shape}")
+                        r, c = np.nonzero(arr)
+                        d = arr[r, c]
+                    if r.size:
+                        rows_list.append(r + offsets[name_i])
+                        cols_list.append(c + offsets[name_j])
+                        data_list.append(d)
+            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, n_total)
+
+        if use_format == "csr" and sp is None:
+            raise ImportError("scipy is required for CSR block systems.")
+        block_rows = []
+        for name_i in self.field_order:
+            row = []
+            row_blocks = self._blocks.get(name_i, {})
+            for name_j in self.field_order:
+                blk = row_blocks.get(name_j)
+                shape = _block_shape(name_i, name_j)
+                if blk is None:
+                    if use_format == "csr":
+                        row.append(sp.csr_matrix(shape))
+                    else:
+                        row.append(np.zeros(shape, dtype=float))
+                    continue
+                if isinstance(blk, FluxSparseMatrix):
+                    if sp is None:
+                        raise ImportError("scipy is required to assemble sparse block systems.")
+                    blk = blk.to_csr()
+                if sp is not None and sp.issparse(blk):
+                    blk = blk.tocsr()
+                    if blk.shape != shape:
+                        raise ValueError(f"Block {name_i},{name_j} has shape {blk.shape}, expected {shape}")
+                    row.append(blk)
+                else:
+                    arr = np.asarray(blk)
+                    if arr.shape != shape:
+                        raise ValueError(f"Block {name_i},{name_j} has shape {arr.shape}, expected {shape}")
+                    if use_format == "csr":
+                        row.append(sp.csr_matrix(arr))
+                    else:
+                        row.append(arr)
+            block_rows.append(row)
+        if use_format == "csr":
+            return sp.bmat(block_rows, format="csr")
+        return np.block(block_rows)
+
+
 def make(
     *,
-    diag: Mapping[str, object] | Sequence[object],
-    rel: Mapping[tuple[str, str], object] | None = None,
-    add_contiguous: object | None = None,
-    sizes: Mapping[str, int] | None = None,
+    diag: Mapping[FieldKey, MatrixLike] | Sequence[MatrixLike],
+    rel: Mapping[tuple[FieldKey, FieldKey], MatrixLike] | None = None,
+    add_contiguous: MatrixLike | None = None,
+    sizes: Mapping[FieldKey, int] | None = None,
     symmetric: bool = False,
     transpose_rule: str = "T",
-):
+) -> FluxBlockMatrix:
     """
-    Build a blocks dict from diagonal blocks, optional relations, and a full matrix.
+    Build a lazy FluxBlockMatrix from diagonal blocks, optional relations, and a full matrix.
     """
     if isinstance(diag, Mapping):
         diag_map = dict(diag)
@@ -126,7 +272,13 @@ def make(
                         _transpose_block(blk, transpose_rule),
                     )
 
-    return blocks
+    return FluxBlockMatrix(
+        blocks,
+        sizes=sizes,
+        order=order,
+        symmetric=symmetric,
+        transpose_rule=transpose_rule,
+    )
 
 
-__all__ = ["diag", "make"]
+__all__ = ["FluxBlockMatrix", "diag", "make"]

fluxfem/solver/block_system.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Mapping, Sequence
+from typing import Any, Mapping, Sequence, TypeAlias
 
 import numpy as np
 
@@ -13,23 +13,27 @@ except Exception:  # pragma: no cover
 from .dirichlet import DirichletBC, free_dofs
 from .sparse import FluxSparseMatrix
 
+MatrixLike: TypeAlias = Any
+FieldKey: TypeAlias = str | int
+BlockMap: TypeAlias = Mapping[FieldKey, Mapping[FieldKey, MatrixLike]]
+
 
 @dataclass(frozen=True)
 class BlockSystem:
-    K: object
+    K: MatrixLike
     F: np.ndarray
     free_dofs: np.ndarray
     dirichlet: DirichletBC
-    field_order: tuple[str, ...]
-    field_slices: dict[str, slice]
+    field_order: tuple[FieldKey, ...]
+    field_slices: dict[FieldKey, slice]
 
-    def expand(self, u_free):
+    def expand(self, u_free: np.ndarray) -> np.ndarray:
         return self.dirichlet.expand_solution(u_free, free=self.free_dofs, n_total=self.F.shape[0])
 
-    def split(self, u_full: np.ndarray) -> dict[str, np.ndarray]:
+    def split(self, u_full: np.ndarray) -> dict[FieldKey, np.ndarray]:
         return {name: np.asarray(u_full)[self.field_slices[name]] for name in self.field_order}
 
-    def join(self, fields: Mapping[str, np.ndarray]) -> np.ndarray:
+    def join(self, fields: Mapping[FieldKey, np.ndarray]) -> np.ndarray:
         parts = []
         for name in self.field_order:
             if name not in fields:
@@ -38,7 +42,9 @@ class BlockSystem:
         return np.concatenate(parts, axis=0)
 
 
-def _build_field_slices(order, sizes):
+def _build_field_slices(
+    order: Sequence[FieldKey], sizes: Mapping[FieldKey, int]
+) -> tuple[dict[FieldKey, int], dict[FieldKey, slice], int]:
     offsets = {}
     slices = {}
     offset = 0
@@ -50,7 +56,12 @@ def _build_field_slices(order, sizes):
     return offsets, slices, offset
 
 
-def split_block_matrix(matrix, *, sizes: Mapping[str, int], order: Sequence[str] | None = None):
+def split_block_matrix(
+    matrix: MatrixLike,
+    *,
+    sizes: Mapping[FieldKey, int],
+    order: Sequence[FieldKey] | None = None,
+) -> dict[FieldKey, dict[FieldKey, MatrixLike]]:
     """
     Split a block matrix into a dict-of-dicts by field order and sizes.
     """
@@ -72,7 +83,7 @@ def split_block_matrix(matrix, *, sizes: Mapping[str, int], order: Sequence[str]
     if mat.shape != (n_total, n_total):
         raise ValueError(f"matrix has shape {mat.shape}, expected {(n_total, n_total)}")
 
-    blocks: dict[str, dict[str, object]] = {}
+    blocks: dict[FieldKey, dict[FieldKey, MatrixLike]] = {}
     for name_i in field_order:
         row = {}
         i0 = offsets[name_i]
@@ -85,7 +96,7 @@ def split_block_matrix(matrix, *, sizes: Mapping[str, int], order: Sequence[str]
     return blocks
 
 
-def _infer_format(blocks, fmt):
+def _infer_format(blocks: BlockMap, fmt: str) -> str:
     if fmt != "auto":
         return fmt
     for row in blocks.values():
@@ -97,7 +108,7 @@ def _infer_format(blocks, fmt):
     return "dense"
 
 
-def _infer_sizes_from_diag_seq(diag_seq):
+def _infer_sizes_from_diag_seq(diag_seq: Sequence[MatrixLike]) -> dict[int, int]:
     sizes = {}
     for idx, blk in enumerate(diag_seq):
         if isinstance(blk, FluxSparseMatrix):
@@ -115,7 +126,11 @@ def _infer_sizes_from_diag_seq(diag_seq):
     return sizes
 
 
-def _coerce_rhs(rhs, order, sizes):
+def _coerce_rhs(
+    rhs: MatrixLike | Sequence[MatrixLike] | Mapping[FieldKey, MatrixLike] | None,
+    order: Sequence[FieldKey],
+    sizes: Mapping[FieldKey, int],
+) -> np.ndarray:
     if rhs is None:
         return np.zeros(sum(int(sizes[n]) for n in order), dtype=float)
     if isinstance(rhs, Mapping):
@@ -139,7 +154,9 @@ def _coerce_rhs(rhs, order, sizes):
     return np.concatenate(parts, axis=0)
 
 
-def _build_dirichlet_from_fields(fields, offsets, *, merge: str):
+def _build_dirichlet_from_fields(
+    fields: Mapping[FieldKey, object], offsets: Mapping[FieldKey, int], *, merge: str
+) -> DirichletBC:
     if merge not in {"check_equal", "error", "first", "last"}:
         raise ValueError("merge must be one of: check_equal, error, first, last")
     dof_map: dict[int, float] = {}
@@ -174,7 +191,13 @@ def _build_dirichlet_from_fields(fields, offsets, *, merge: str):
     return DirichletBC(dofs_sorted, vals_sorted)
 
 
-def _build_dirichlet_from_sequence(seq, order, offsets, *, merge: str):
+def _build_dirichlet_from_sequence(
+    seq: Sequence[object | None],
+    order: Sequence[FieldKey],
+    offsets: Mapping[FieldKey, int],
+    *,
+    merge: str,
+) -> DirichletBC:
     if merge not in {"check_equal", "error", "first", "last"}:
         raise ValueError("merge must be one of: check_equal, error, first, last")
     if len(seq) != len(order):
@@ -211,7 +234,7 @@ def _build_dirichlet_from_sequence(seq, order, offsets, *, merge: str):
     return DirichletBC(dofs_sorted, vals_sorted)
 
 
-def _transpose_block(block, rule: str):
+def _transpose_block(block: MatrixLike, rule: str) -> MatrixLike:
     if isinstance(block, FluxSparseMatrix):
         if sp is None:
             raise ImportError("scipy is required to transpose FluxSparseMatrix blocks.")
@@ -225,7 +248,7 @@ def _transpose_block(block, rule: str):
     return out
 
 
-def _add_blocks(a, b):
+def _add_blocks(a: MatrixLike | None, b: MatrixLike | None) -> MatrixLike | None:
     if a is None:
         return b
     if b is None:
@@ -245,14 +268,14 @@ def _add_blocks(a, b):
 
 def _blocks_from_diag_rel(
     *,
-    diag: Mapping[str, object] | Sequence[object],
-    sizes: Mapping[str, int],
-    order: Sequence[str],
-    rel: Mapping[tuple[str, str], object] | None = None,
-    add_contiguous: object | None = None,
+    diag: Mapping[FieldKey, MatrixLike] | Sequence[MatrixLike],
+    sizes: Mapping[FieldKey, int],
+    order: Sequence[FieldKey],
+    rel: Mapping[tuple[FieldKey, FieldKey], MatrixLike] | None = None,
+    add_contiguous: MatrixLike | None = None,
     symmetric: bool = False,
     transpose_rule: str = "T",
-) -> Mapping[str, Mapping[str, object]]:
+) -> BlockMap:
     if isinstance(diag, Mapping):
         diag_map = dict(diag)
     else:
@@ -296,12 +319,12 @@ def _blocks_from_diag_rel(
 
 def build_block_system(
     *,
-    diag: Mapping[str, object] | Sequence[object],
-    sizes: Mapping[str, int] | None = None,
-    rel: Mapping[tuple[str, str], object] | None = None,
-    add_contiguous: object | None = None,
-    rhs: Mapping[str, object] | Sequence[object] | np.ndarray | None = None,
-    constraints=None,
+    diag: Mapping[FieldKey, MatrixLike] | Sequence[MatrixLike],
+    sizes: Mapping[FieldKey, int] | None = None,
+    rel: Mapping[tuple[FieldKey, FieldKey], MatrixLike] | None = None,
+    add_contiguous: MatrixLike | None = None,
+    rhs: Mapping[FieldKey, MatrixLike] | Sequence[MatrixLike] | np.ndarray | None = None,
+    constraints: object | None = None,
     merge: str = "check_equal",
     format: str = "auto",
     symmetric: bool = False,

fluxfem/solver/cg.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any, Callable, TypeAlias
+
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
@@ -14,8 +16,12 @@ from dataclasses import dataclass
 
 from .preconditioner import make_block_jacobi_preconditioner
 
+ArrayLike: TypeAlias = jnp.ndarray
+MatVec: TypeAlias = Callable[[jnp.ndarray], jnp.ndarray]
+CGInfo: TypeAlias = dict[str, Any]
+
 
-def _matvec_builder(A):
+def _matvec_builder(A: Any) -> MatVec:
     if jsparse is not None and isinstance(A, jsparse.BCOO):
         return lambda x: A @ x
     if isinstance(A, FluxSparseMatrix):
@@ -33,7 +39,7 @@ def _matvec_builder(A):
     return mv
 
 
-def _coo_tuple_from_any(A):
+def _coo_tuple_from_any(A: Any):
     if isinstance(A, FluxSparseMatrix):
         return A.to_coo()
     if isinstance(A, tuple) and len(A) == 4:
@@ -53,7 +59,7 @@ def _coo_tuple_from_any(A):
     return None
 
 
-def _to_flux_matrix(A):
+def _to_flux_matrix(A: Any) -> FluxSparseMatrix:
     if isinstance(A, FluxSparseMatrix):
         return A
     coo = _coo_tuple_from_any(A)
@@ -62,7 +68,7 @@ def _to_flux_matrix(A):
     return FluxSparseMatrix.from_bilinear(coo)
 
 
-def _to_bcoo_matrix(A):
+def _to_bcoo_matrix(A: Any):
     if jsparse is None:
         raise ImportError("jax.experimental.sparse is required for BCOO matvec")
     if jsparse is not None and isinstance(A, jsparse.BCOO):
@@ -75,7 +81,7 @@ def _to_bcoo_matrix(A):
     return jsparse.BCOO((data, idx), shape=(n, n))
 
 
-def _normalize_matvec_matrix(A, matvec: str):
+def _normalize_matvec_matrix(A: Any, matvec: str):
     if matvec == "flux":
         return _to_flux_matrix(A)
     if matvec == "bcoo":
@@ -101,7 +107,14 @@ class CGOperator:
     preconditioner: object | None = None
     solver: str = "cg"
 
-    def solve(self, b, *, x0=None, tol: float = 1e-8, maxiter: int | None = None):
+    def solve(
+        self,
+        b: jnp.ndarray,
+        *,
+        x0: jnp.ndarray | None = None,
+        tol: float = 1e-8,
+        maxiter: int | None = None,
+    ):
         if self.solver == "cg":
             return cg_solve(
                 self.A,
@@ -124,13 +137,13 @@ class CGOperator:
 
 
 def build_cg_operator(
-    A,
+    A: Any,
     *,
     matvec: str = "flux",
-    preconditioner=None,
+    preconditioner: object | None = None,
     solver: str = "cg",
     dof_per_node: int | None = None,
-    block_sizes=None,
+    block_sizes: object | None = None,
 ) -> CGOperator:
     """
     Normalize CG inputs into a single operator interface.
@@ -144,7 +157,7 @@ def build_cg_operator(
     return CGOperator(A=A_mat, preconditioner=precon, solver=solver)
 
 
-def _diag_builder(A, n: int):
+def _diag_builder(A: Any, n: int) -> jnp.ndarray:
     """
     Build diagonal for a Jacobi preconditioner when available.
     """
@@ -171,14 +184,14 @@ def _diag_builder(A, n: int):
 
 
 def _cg_solve_single(
-    A,
-    b,
+    A: Any,
+    b: jnp.ndarray,
     *,
-    x0=None,
+    x0: jnp.ndarray | None = None,
     tol: float = 1e-8,
     maxiter: int | None = None,
-    preconditioner=None,
-):
+    preconditioner: object | None = None,
+) -> tuple[jnp.ndarray, CGInfo]:
     """
     Conjugate gradient (Ax=b) in JAX.
     A: FluxSparseMatrix / (rows, cols, data, n) / dense array
@@ -241,14 +254,14 @@ def _cg_solve_single(
 
 
 def cg_solve(
-    A,
-    b,
+    A: Any,
+    b: jnp.ndarray,
     *,
-    x0=None,
+    x0: jnp.ndarray | None = None,
     tol: float = 1e-8,
     maxiter: int | None = None,
-    preconditioner=None,
-):
+    preconditioner: object | None = None,
+) -> tuple[jnp.ndarray, CGInfo]:
     """
     Conjugate gradient (Ax=b) in JAX.
     Supports single RHS (n,) or multiple RHS (n, n_rhs).
@@ -290,14 +303,14 @@ def cg_solve(
 
 
 def _cg_solve_jax_single(
-    A,
-    b,
+    A: Any,
+    b: jnp.ndarray,
     *,
-    x0=None,
+    x0: jnp.ndarray | None = None,
     tol: float = 1e-8,
     maxiter: int | None = None,
-    preconditioner=None,
-):
+    preconditioner: object | None = None,
+) -> tuple[jnp.ndarray, CGInfo]:
     """
     Conjugate gradient via jax.scipy.sparse.linalg.cg.
     A: FluxSparseMatrix / (rows, cols, data, n) / dense array / callable
@@ -359,14 +372,14 @@ def _cg_solve_jax_single(
 
 
 def cg_solve_jax(
-    A,
-    b,
+    A: Any,
+    b: jnp.ndarray,
     *,
-    x0=None,
+    x0: jnp.ndarray | None = None,
     tol: float = 1e-8,
     maxiter: int | None = None,
-    preconditioner=None,
-):
+    preconditioner: object | None = None,
+) -> tuple[jnp.ndarray, CGInfo]:
     """
     Conjugate gradient via jax.scipy.sparse.linalg.cg.
     Supports single RHS (n,) or multiple RHS (n, n_rhs).