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

fluxfem/solver/solve_runner.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import time
 from dataclasses import dataclass, field
+import warnings
 from typing import Any, Callable, Iterable, List, Sequence
 
 import numpy as np
@@ -10,6 +11,7 @@ import jax.numpy as jnp
 from ..core.assembly import 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 .newton import newton_solve
@@ -35,6 +37,8 @@ class NonlinearAnalysis:
         Unscaled external load vector (scaled by load factor in `external_for_load`).
     dirichlet : tuple | None
         (dofs, values) for Dirichlet boundary conditions.
+    extra_terms : list[callable] | None
+        Optional extra term assemblers returning (K, f[, metrics]).
     jacobian_pattern : Any | None
         Optional sparsity pattern to reuse between load steps.
     dtype : Any
@@ -46,6 +50,7 @@ class NonlinearAnalysis:
     params: Any
     base_external_vector: Any | None = None
     dirichlet: tuple | None = None
+    extra_terms: list[Callable] | None = None
     jacobian_pattern: Any | None = None
     dtype: Any = jnp.float64
 
@@ -71,6 +76,7 @@ class NewtonLoopConfig:
     linear_maxiter: int | None = None
     linear_tol: float | None = None
     linear_preconditioner: Any | None = None
+    matfree_mode: str = "linearize"
     load_sequence: Sequence[float] | None = None
     n_steps: int = 1
 
@@ -95,6 +101,7 @@ class NewtonSolveRunner:
     def __init__(self, analysis: NonlinearAnalysis, config: NewtonLoopConfig):
         self.analysis = analysis
         self.config = config
+        self._matfree_cache: dict[str, Any] = {}
 
     def run(
         self,
@@ -154,6 +161,16 @@ class NewtonSolveRunner:
                     schedule.append(lf_clamped)
                     prev = lf_clamped
                 history: List[LoadStepResult] = []
+                matfree_cache = None
+                if self.config.linear_preconditioner == "diag0":
+                    n_free = self.analysis.space.n_dofs
+                    if self.analysis.dirichlet is not None:
+                        n_free -= len(self.analysis.dirichlet[0])
+                    cached_free = self._matfree_cache.get("n_free_dofs")
+                    if cached_free is not None and cached_free != n_free:
+                        self._matfree_cache.clear()
+                    self._matfree_cache["n_free_dofs"] = n_free
+                    matfree_cache = self._matfree_cache
                 for step_i, load_factor in enumerate(schedule, start=1):
                     with timer.section("step"):
                         external = self.analysis.external_for_load(load_factor)
@@ -203,6 +220,8 @@ class NewtonSolveRunner:
                             linear_maxiter=self.config.linear_maxiter,
                             linear_tol=self.config.linear_tol,
                             linear_preconditioner=self.config.linear_preconditioner,
+                            matfree_cache=matfree_cache,
+                            matfree_mode=self.config.matfree_mode,
                             dirichlet=self.analysis.dirichlet,
                             line_search=self.config.line_search,
                             max_ls=self.config.max_ls,
@@ -210,6 +229,7 @@ class NewtonSolveRunner:
                             external_vector=external,
                             callback=cb,
                             jacobian_pattern=self.analysis.jacobian_pattern,
+                            extra_terms=self.analysis.extra_terms,
                         )
                         exception = None
                     except Exception as e:  # pragma: no cover - defensive
@@ -279,6 +299,7 @@ def solve_nonlinear(
     *,
     dirichlet: tuple | None = None,
     base_external_vector=None,
+    extra_terms=None,
     dtype=jnp.float64,
     maxiter: int = 20,
     tol: float = 1e-8,
@@ -287,6 +308,7 @@ def solve_nonlinear(
     linear_maxiter: int | None = None,
     linear_tol: float | None = None,
     linear_preconditioner=None,
+    matfree_mode: str = "linearize",
     line_search: bool = False,
     max_ls: int = 10,
     ls_c: float = 1e-4,
@@ -303,6 +325,7 @@ def solve_nonlinear(
         params=params,
         base_external_vector=base_external_vector,
         dirichlet=dirichlet,
+        extra_terms=extra_terms,
         dtype=dtype,
         jacobian_pattern=jacobian_pattern,
     )
@@ -314,6 +337,7 @@ def solve_nonlinear(
         linear_maxiter=linear_maxiter,
         linear_tol=linear_tol,
         linear_preconditioner=linear_preconditioner,
+        matfree_mode=matfree_mode,
         line_search=line_search,
         max_ls=max_ls,
         ls_c=ls_c,
@@ -359,10 +383,42 @@ class LinearSolveConfig:
     Control parameters for the linear solve with optional load scaling.
     """
 
-    method: str = "spsolve"  # "spsolve" | "spdirect_solve_gpu" | "cg" | "cg_custom"
+    method: str = "spsolve"  # "spsolve" | "spdirect_solve_gpu" | "cg" | "cg_custom" | "petsc_shell"
     tol: float = 1e-8
     maxiter: int | None = None
     preconditioner: Any | None = None
+    ksp_type: str | None = None
+    pc_type: str | None = None
+    ksp_rtol: float | None = None
+    ksp_atol: float | None = None
+    ksp_max_it: int | None = None
+    petsc_ksp_norm_type: str | None = None
+    petsc_ksp_monitor_true_residual: bool = False
+    petsc_ksp_converged_reason: bool = False
+    petsc_ksp_monitor_short: bool = False
+    petsc_shell_pmat: bool = False
+    petsc_shell_pmat_mode: str = "full"
+    petsc_shell_pmat_rebuild_iters: int | None = None
+    petsc_shell_fallback: bool = False
+    petsc_shell_fallback_ksp_types: tuple[str, ...] = ("bcgs", "gmres")
+    petsc_shell_fallback_rebuild_pmat: bool = True
+
+    @classmethod
+    def from_preset(cls, name: str) -> "LinearSolveConfig":
+        preset = name.lower()
+        if preset == "contact":
+            return cls(
+                method="petsc_shell",
+                ksp_type="bcgs",
+                pc_type="ilu",
+                petsc_shell_pmat=True,
+                petsc_shell_pmat_mode="full",
+                petsc_ksp_norm_type="unpreconditioned",
+                petsc_ksp_monitor_true_residual=True,
+                petsc_ksp_converged_reason=True,
+                petsc_shell_fallback=True,
+            )
+        raise ValueError(f"Unknown LinearSolveConfig preset: {name}")
 
 
 @dataclass
@@ -392,6 +448,9 @@ class LinearSolveRunner:
     def __init__(self, analysis: LinearAnalysis, config: LinearSolveConfig):
         self.analysis = analysis
         self.config = config
+        self._petsc_shell_pmat = None
+        self._petsc_shell_last_iters = None
+        self._petsc_shell_pmat_rebuilds = 0
 
     def run(
         self,
@@ -488,6 +547,114 @@ class LinearSolveRunner:
                             stop_reason=("converged" if lin_conv else "linfail"),
                             nan_detected=bool(np.isnan(lin_res)) if lin_res is not None else False,
                         )
+                    elif self.config.method == "petsc_shell":
+                        base_ksp_type = self.config.ksp_type or "gmres"
+                        pc_type = self.config.pc_type if self.config.pc_type is not None else "none"
+                        ksp_rtol = self.config.ksp_rtol if self.config.ksp_rtol is not None else self.config.tol
+                        ksp_atol = self.config.ksp_atol
+                        ksp_max_it = self.config.ksp_max_it if self.config.ksp_max_it is not None else self.config.maxiter
+                        petsc_options = {}
+                        if self.config.petsc_ksp_norm_type:
+                            petsc_options["fluxfem_ksp_norm_type"] = self.config.petsc_ksp_norm_type
+                        if self.config.petsc_ksp_monitor_true_residual:
+                            petsc_options["fluxfem_ksp_monitor_true_residual"] = ""
+                        if self.config.petsc_ksp_converged_reason:
+                            petsc_options["fluxfem_ksp_converged_reason"] = ""
+                        if self.config.petsc_ksp_monitor_short:
+                            petsc_options["fluxfem_ksp_monitor_short"] = ""
+                        if not petsc_options:
+                            petsc_options = None
+                        use_pmat = bool(self.config.petsc_shell_pmat)
+                        rebuild_thresh = self.config.petsc_shell_pmat_rebuild_iters
+                        if use_pmat:
+                            pmat_mode = (self.config.petsc_shell_pmat_mode or "full").lower()
+                            if pmat_mode == "none":
+                                use_pmat = False
+                                pmat = None
+                            elif pmat_mode == "full":
+                                pmat = K_ff
+                            else:
+                                warnings.warn(
+                                    f"petsc_shell_pmat_mode='{pmat_mode}' is not supported in runner; "
+                                    "falling back to 'full'.",
+                                    RuntimeWarning,
+                                )
+                                pmat = K_ff
+                            if use_pmat:
+                                if self._petsc_shell_pmat is None:
+                                    self._petsc_shell_pmat = pmat
+                                    self._petsc_shell_pmat_rebuilds += 1
+                                elif rebuild_thresh is not None and self._petsc_shell_last_iters is not None:
+                                    if self._petsc_shell_last_iters > rebuild_thresh:
+                                        self._petsc_shell_pmat = pmat
+                                        self._petsc_shell_pmat_rebuilds += 1
+                                pmat = self._petsc_shell_pmat
+                        if not use_pmat:
+                            pmat = None
+
+                        def _attempt_solve(ksp_type: str):
+                            return petsc_shell_solve(
+                                K_ff,
+                                F_free,
+                                preconditioner=self.config.preconditioner,
+                                ksp_type=ksp_type,
+                                pc_type=pc_type,
+                                rtol=ksp_rtol,
+                                atol=ksp_atol,
+                                max_it=ksp_max_it,
+                                pmat=pmat,
+                                options=petsc_options,
+                                return_info=True,
+                            )
+
+                        fallback_ksp = [base_ksp_type]
+                        if self.config.petsc_shell_fallback:
+                            for ksp in self.config.petsc_shell_fallback_ksp_types:
+                                if ksp not in fallback_ksp:
+                                    fallback_ksp.append(ksp)
+                        fallback_attempts = []
+                        petsc_info = None
+                        u_free = None
+                        for ksp in fallback_ksp:
+                            fallback_attempts.append(ksp)
+                            u_free, petsc_info = _attempt_solve(ksp)
+                            lin_conv = petsc_info.get("converged")
+                            reason = petsc_info.get("reason")
+                            if lin_conv is None and reason is not None:
+                                lin_conv = reason > 0
+                            if lin_conv:
+                                break
+                            if self.config.petsc_shell_fallback and use_pmat and self.config.petsc_shell_fallback_rebuild_pmat:
+                                self._petsc_shell_pmat = pmat
+                                self._petsc_shell_pmat_rebuilds += 1
+                        lin_iters = petsc_info.get("iters")
+                        lin_res = petsc_info.get("residual_norm")
+                        lin_solve_dt = petsc_info.get("solve_time")
+                        pc_setup_dt = petsc_info.get("pc_setup_time")
+                        pmat_dt = petsc_info.get("pmat_build_time")
+                        lin_conv = petsc_info.get("converged")
+                        if lin_conv is None and petsc_info.get("reason") is not None:
+                            lin_conv = petsc_info.get("reason") > 0
+                        if lin_conv is None:
+                            lin_conv = True
+                        self._petsc_shell_last_iters = lin_iters
+                        info = SolverResult(
+                            converged=bool(lin_conv),
+                            iters=int(lin_iters) if lin_iters is not None else 0,
+                            linear_iters=int(lin_iters) if lin_iters is not None else None,
+                            linear_converged=bool(lin_conv),
+                            linear_residual=float(lin_res) if lin_res is not None else None,
+                            linear_solve_time=float(lin_solve_dt) if lin_solve_dt is not None else None,
+                            pc_setup_time=float(pc_setup_dt) if pc_setup_dt is not None else None,
+                            pmat_build_time=float(pmat_dt) if pmat_dt is not None else None,
+                            pmat_rebuilds=self._petsc_shell_pmat_rebuilds if use_pmat else None,
+                            pmat_mode=self.config.petsc_shell_pmat_mode if use_pmat else None,
+                            tol=self.config.tol,
+                            stop_reason=("converged" if lin_conv else "linfail"),
+                            nan_detected=bool(np.isnan(lin_res)) if lin_res is not None else False,
+                        )
+                        if len(fallback_attempts) > 1:
+                            info.linear_fallbacks = fallback_attempts
                     else:
                         raise ValueError(f"Unknown linear solve method: {self.config.method}")
 

fluxfem/solver/solver.py

@@ -5,8 +5,10 @@ import jax.numpy as jnp
 
 from .cg import cg_solve, cg_solve_jax
 from .newton import newton_solve
+from .petsc import petsc_solve, petsc_shell_solve
 from ..core.solver import spdirect_solve_cpu, spdirect_solve_gpu
 from .dirichlet import (
+    DirichletBC,
     condense_dirichlet_dense,
     condense_dirichlet_fluxsparse,
     expand_dirichlet_solution,
@@ -46,6 +48,12 @@ class LinearSolver:
         elif self.method == "spdirect_solve_gpu":
             x = spdirect_solve_gpu(A, b)
             return np.asarray(x), {"iters": 1, "converged": True}
+        elif self.method == "petsc":
+            x = petsc_solve(A, b)
+            return np.asarray(x), {"iters": None, "converged": True}
+        elif self.method == "petsc_shell":
+            x = petsc_shell_solve(A, b)
+            return np.asarray(x), {"iters": None, "converged": True}
         else:
             raise ValueError(f"Unknown linear method: {self.method}")
 
@@ -64,7 +72,10 @@ class LinearSolver:
         if dirichlet_mode not in ("condense", "enforce"):
             raise ValueError("dirichlet_mode must be 'condense' or 'enforce'.")
 
-        dir_dofs, dir_vals = dirichlet
+        if isinstance(dirichlet, DirichletBC):
+            dir_dofs, dir_vals = dirichlet.as_tuple()
+        else:
+            dir_dofs, dir_vals = dirichlet
         if dirichlet_mode == "enforce":
             if isinstance(A, FluxSparseMatrix):
                 A_bc, b_bc = enforce_dirichlet_sparse(A, b, dir_dofs, dir_vals)

fluxfem/solver/sparse.py

@@ -12,6 +12,79 @@ except Exception:  # pragma: no cover
     sp = None
 
 
+def coalesce_coo(rows, cols, data):
+    """
+    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):
+    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, n_dofs: int | None = None):
+    """
+    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):
+    """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 +153,7 @@ 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, cols=None, data=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 +161,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,6 +185,7 @@ 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):
@@ -121,11 +202,25 @@ class FluxSparseMatrix:
 
     def with_data(self, data):
         """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):
+        """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):
         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):
+        """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()")
@@ -167,6 +262,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: