fluxfem 0.1.1a0__py3-none-any.whl

fluxfem/solver/cg.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import jax.scipy as jsp
+
+try:
+    from jax.experimental import sparse as jsparse
+except Exception:  # pragma: no cover
+    jsparse = None
+
+from .sparse import FluxSparseMatrix
+
+
+def _matvec_builder(A):
+    if jsparse is not None and isinstance(A, jsparse.BCOO):
+        return lambda x: A @ x
+    if isinstance(A, FluxSparseMatrix):
+        return A.matvec
+    if hasattr(A, "matvec"):
+        return A.matvec
+    if callable(A):
+        return lambda x: A(x)
+    if isinstance(A, tuple) and len(A) == 4:
+        return FluxSparseMatrix.from_bilinear(A).matvec
+
+    def mv(x):
+        return jnp.asarray(A) @ x
+
+    return mv
+
+
+def _diag_builder(A, n: int):
+    """
+    Build diagonal for a Jacobi preconditioner when available.
+    """
+    if jsparse is not None and isinstance(A, jsparse.BCOO):
+        idx = A.indices
+        data = A.data
+        rows = idx[:, 0]
+        cols = idx[:, 1]
+        mask = rows == cols
+        diag_contrib = jnp.where(mask, data, jnp.zeros_like(data))
+        return jax.ops.segment_sum(diag_contrib, rows, n)
+    if isinstance(A, FluxSparseMatrix):
+        return A.diag()
+    if hasattr(A, "diagonal"):
+        return jnp.asarray(A.diagonal())
+    if isinstance(A, tuple) and len(A) == 4:
+        return _diag_builder(FluxSparseMatrix.from_bilinear(A), n)
+    if callable(A):
+        raise ValueError("Jacobi preconditioner requires access to matrix diagonal")
+    arr = jnp.asarray(A)
+    if arr.ndim == 2 and arr.shape[0] == arr.shape[1]:
+        return jnp.diag(arr)
+    raise ValueError("Cannot build Jacobi preconditioner: diagonal unavailable")
+
+
+def _cg_solve_single(
+    A,
+    b,
+    *,
+    x0=None,
+    tol: float = 1e-8,
+    maxiter: int | None = None,
+    preconditioner=None,
+):
+    """
+    Conjugate gradient (Ax=b) in JAX.
+    A: FluxSparseMatrix / (rows, cols, data, n) / dense array
+    b: RHS (jnp or np)
+    preconditioner: None | "jacobi" | callable(r) -> z
+    returns: (x, info dict)
+    """
+    b = jnp.asarray(b)
+    n = b.shape[0]
+    if x0 is None:
+        x0 = jnp.zeros_like(b)
+    if maxiter is None:
+        maxiter = max(10 * n, 1)
+
+    mv = _matvec_builder(A)
+    precon = None
+
+    if preconditioner is None:
+        pass
+    elif preconditioner == "jacobi":
+        diag = _diag_builder(A, n)
+        inv_diag = jnp.where(diag != 0.0, 1.0 / diag, 0.0)
+
+        def precon(r):
+            return inv_diag * r
+
+    elif preconditioner == "block_jacobi":
+        # Expect 3 DOFs per node
+        if n % 3 != 0:
+            raise ValueError("block_jacobi requires n_dofs % 3 == 0")
+        if jsparse is not None and isinstance(A, jsparse.BCOO):
+            rows = A.indices[:, 0]
+            cols = A.indices[:, 1]
+            data = A.data
+        elif isinstance(A, FluxSparseMatrix):
+            rows = jnp.asarray(A.pattern.rows)
+            cols = jnp.asarray(A.pattern.cols)
+            data = jnp.asarray(A.data)
+        else:
+            raise ValueError("block_jacobi requires FluxSparseMatrix or BCOO")
+
+        block_rows = rows // 3
+        block_cols = cols // 3
+        lr = rows % 3
+        lc = cols % 3
+        mask = block_rows == block_cols
+        block_rows = block_rows[mask]
+        lr = lr[mask]
+        lc = lc[mask]
+        data = data[mask]
+        n_block = n // 3
+        blocks = jnp.zeros((n_block, 3, 3), dtype=data.dtype)
+        blocks = blocks.at[block_rows, lr, lc].add(data)
+        blocks = blocks + 1e-12 * jnp.eye(3)[None, :, :]
+        inv_blocks = jnp.linalg.inv(blocks)
+
+        def precon(r):
+            rb = r.reshape((n_block, 3))
+            zb = jnp.einsum("bij,bj->bi", inv_blocks, rb)
+            return zb.reshape((-1,))
+
+    elif callable(preconditioner):
+        precon = preconditioner
+    else:
+        raise ValueError(f"Unknown preconditioner type: {preconditioner}")
+
+    def body_fun(state):
+        k, x, r, p, rz_old = state
+        Ap = mv(p)
+        alpha = rz_old / jnp.dot(p, Ap)
+        x_new = x + alpha * p
+        r_new = r - alpha * Ap
+        z_new = r_new if precon is None else precon(r_new)
+        rz_new = jnp.dot(r_new, z_new)
+        beta = rz_new / rz_old
+        p_new = z_new + beta * p
+        return k + 1, x_new, r_new, p_new, rz_new
+
+    r0 = b - mv(x0)
+    z0 = r0 if precon is None else precon(r0)
+    rz0 = jnp.dot(r0, z0)
+    init_state = (0, x0, r0, z0, rz0)
+
+    def cond_fun(state):
+        k, x, r, p, rz_old = state
+        return jnp.logical_and(k < maxiter, jnp.dot(r, r) > tol * tol)
+
+    k, x, r, p, rz_old = jax.lax.while_loop(cond_fun, body_fun, init_state)
+    res_norm = jnp.sqrt(jnp.dot(r, r))
+    info = {"iters": k, "residual_norm": res_norm}
+    return x, info
+
+
+def cg_solve(
+    A,
+    b,
+    *,
+    x0=None,
+    tol: float = 1e-8,
+    maxiter: int | None = None,
+    preconditioner=None,
+):
+    """
+    Conjugate gradient (Ax=b) in JAX.
+    Supports single RHS (n,) or multiple RHS (n, n_rhs).
+    """
+    b_arr = jnp.asarray(b)
+    if b_arr.ndim == 1:
+        return _cg_solve_single(
+            A,
+            b_arr,
+            x0=x0,
+            tol=tol,
+            maxiter=maxiter,
+            preconditioner=preconditioner,
+        )
+
+    if b_arr.ndim != 2:
+        raise ValueError("cg_solve expects b with shape (n,) or (n, n_rhs).")
+
+    xs = []
+    infos = []
+    for i in range(b_arr.shape[1]):
+        x0_i = None if x0 is None else jnp.asarray(x0)[:, i]
+        x_i, info_i = _cg_solve_single(
+            A,
+            b_arr[:, i],
+            x0=x0_i,
+            tol=tol,
+            maxiter=maxiter,
+            preconditioner=preconditioner,
+        )
+        xs.append(x_i)
+        infos.append(info_i)
+    x_out = jnp.stack(xs, axis=1)
+    info = {
+        "iters": [int(i.get("iters", 0)) for i in infos],
+        "residual_norm": jnp.asarray([i.get("residual_norm", 0.0) for i in infos]),
+    }
+    return x_out, info
+
+
+def _cg_solve_jax_single(
+    A,
+    b,
+    *,
+    x0=None,
+    tol: float = 1e-8,
+    maxiter: int | None = None,
+    preconditioner=None,
+):
+    """
+    Conjugate gradient via jax.scipy.sparse.linalg.cg.
+    A: FluxSparseMatrix / (rows, cols, data, n) / dense array / callable
+    preconditioner: None | "jacobi" | callable(r) -> z
+    returns: (x, info dict)
+    """
+    b = jnp.asarray(b)
+    n = b.shape[0]
+    if x0 is None:
+        x0 = jnp.zeros_like(b)
+    if maxiter is None:
+        maxiter = max(10 * n, 1)
+
+    mv = _matvec_builder(A)
+
+    precon = None
+    if preconditioner is None:
+        pass
+    elif preconditioner == "jacobi":
+        diag = _diag_builder(A, n)
+        inv_diag = jnp.where(diag != 0.0, 1.0 / diag, 0.0)
+
+        def precon(r):
+            return inv_diag * r
+
+    elif callable(preconditioner):
+        precon = preconditioner
+    else:
+        raise ValueError(f"Unknown preconditioner type: {preconditioner}")
+
+    def M_fun(x):
+        return precon(x) if precon is not None else x
+
+    x, info_val = jsp.sparse.linalg.cg(
+        mv,
+        b,
+        x0=x0,
+        tol=tol,
+        atol=0.0,
+        maxiter=maxiter,
+        M=M_fun if precon is not None else None,
+    )
+    res = b - mv(x)
+    res_norm = jnp.sqrt(jnp.dot(res, res))
+
+    if isinstance(info_val, (int, jnp.integer)):
+        iters = int(info_val)
+        converged = iters == 0
+    elif info_val is None:
+        # jax.scipy may return None on success; treat as converged with unknown iter count.
+        iters = 0
+        converged = True
+    else:
+        # Unknown type; fall back to "not sure" but keep running.
+        iters = 0
+        converged = True
+    info = {"iters": iters, "residual_norm": res_norm, "converged": converged, "info": info_val}
+    return x, info
+
+
+def cg_solve_jax(
+    A,
+    b,
+    *,
+    x0=None,
+    tol: float = 1e-8,
+    maxiter: int | None = None,
+    preconditioner=None,
+):
+    """
+    Conjugate gradient via jax.scipy.sparse.linalg.cg.
+    Supports single RHS (n,) or multiple RHS (n, n_rhs).
+    """
+    b_arr = jnp.asarray(b)
+    if b_arr.ndim == 1:
+        return _cg_solve_jax_single(
+            A,
+            b_arr,
+            x0=x0,
+            tol=tol,
+            maxiter=maxiter,
+            preconditioner=preconditioner,
+        )
+
+    if b_arr.ndim != 2:
+        raise ValueError("cg_solve_jax expects b with shape (n,) or (n, n_rhs).")
+
+    xs = []
+    infos = []
+    for i in range(b_arr.shape[1]):
+        x0_i = None if x0 is None else jnp.asarray(x0)[:, i]
+        x_i, info_i = _cg_solve_jax_single(
+            A,
+            b_arr[:, i],
+            x0=x0_i,
+            tol=tol,
+            maxiter=maxiter,
+            preconditioner=preconditioner,
+        )
+        xs.append(x_i)
+        infos.append(info_i)
+    x_out = jnp.stack(xs, axis=1)
+    info = {
+        "iters": [int(i.get("iters", 0)) for i in infos],
+        "residual_norm": jnp.asarray([i.get("residual_norm", 0.0) for i in infos]),
+        "converged": [bool(i.get("converged", True)) for i in infos],
+    }
+    return x_out, info

fluxfem/solver/dirichlet.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import numpy as np
+import jax.numpy as jnp
+
+from .sparse import FluxSparseMatrix
+
+
+def _normalize_dirichlet(dofs, vals):
+    dir_arr = np.asarray(dofs, dtype=int)
+    if vals is None:
+        return dir_arr, np.zeros(dir_arr.shape[0], dtype=float)
+    return dir_arr, np.asarray(vals, dtype=float)
+
+
+def enforce_dirichlet_dense(K, F, dofs, vals):
+    """Apply Dirichlet conditions directly to stiffness/load (dense)."""
+    Kc = np.asarray(K, dtype=float).copy()
+    Fc = np.asarray(F, dtype=float).copy()
+    dofs, vals = _normalize_dirichlet(dofs, vals)
+    if Fc.ndim == 2:
+        Fc = Fc - (Kc[:, dofs] @ vals)[:, None]
+    else:
+        Fc = Fc - Kc[:, dofs] @ vals
+    for d, v in zip(dofs, vals):
+        Kc[d, :] = 0.0
+        Kc[:, d] = 0.0
+        Kc[d, d] = 1.0
+        if Fc.ndim == 2:
+            Fc[d, :] = v
+        else:
+            Fc[d] = v
+    return Kc, Fc
+
+
+def enforce_dirichlet_sparse(A: FluxSparseMatrix, F, dofs, vals):
+    """Apply Dirichlet conditions to FluxSparseMatrix + load (CSR)."""
+    K_csr = A.to_csr().tolil()
+    Fc = np.asarray(F, dtype=float).copy()
+    dofs, vals = _normalize_dirichlet(dofs, vals)
+    if Fc.ndim == 2:
+        Fc = Fc - (K_csr[:, dofs] @ vals)[:, None]
+    else:
+        Fc = Fc - K_csr[:, dofs] @ vals
+    for d, v in zip(dofs, vals):
+        K_csr.rows[d] = [d]
+        K_csr.data[d] = [1.0]
+        K_csr[:, d] = 0.0
+        K_csr[d, d] = 1.0
+        if Fc.ndim == 2:
+            Fc[d, :] = v
+        else:
+            Fc[d] = v
+    return K_csr.tocsr(), Fc
+
+
+def condense_dirichlet_fluxsparse(A: FluxSparseMatrix, F, dofs, vals):
+    """
+    Condense Dirichlet DOFs for a FluxSparseMatrix.
+    Returns: (K_ff, F_free, free_dofs, dir_dofs, dir_vals)
+    """
+    K_csr = A.to_csr()
+    dir_arr, dir_vals_arr = _normalize_dirichlet(dofs, vals)
+    mask = np.ones(K_csr.shape[0], dtype=bool)
+    mask[dir_arr] = False
+    free = np.nonzero(mask)[0]
+    K_ff = K_csr[free][:, free]
+    K_fd = K_csr[free][:, dir_arr] if dir_arr.size > 0 else None
+    F_full = np.asarray(F, dtype=float)
+    F_free = F_full[free]
+    if K_fd is not None and dir_arr.size > 0:
+        if F_free.ndim == 2:
+            F_free = F_free - (K_fd @ dir_vals_arr)[:, None]
+        else:
+            F_free = F_free - K_fd @ dir_vals_arr
+    return K_ff, F_free, free, dir_arr, dir_vals_arr
+
+
+def free_dofs(n_dofs: int, dir_dofs) -> np.ndarray:
+    """
+    Return free DOF indices given total DOFs and Dirichlet DOFs.
+    """
+    dir_set = np.asarray(dir_dofs, dtype=int)
+    mask = np.ones(int(n_dofs), dtype=bool)
+    mask[dir_set] = False
+    return np.nonzero(mask)[0]
+
+
+def condense_dirichlet_dense(K, F, dofs, vals):
+    """
+    Eliminate Dirichlet dofs for dense/CSR matrices and return condensed system.
+    Returns: (K_cc, F_c, free_dofs, dir_dofs, dir_vals)
+    """
+    K_np = np.asarray(K, dtype=float)
+    F_np = np.asarray(F, dtype=float)
+    n = K_np.shape[0]
+
+    dir_set, dir_vals = _normalize_dirichlet(dofs, vals)
+    mask = np.ones(n, dtype=bool)
+    mask[dir_set] = False
+    free_dofs = np.nonzero(mask)[0]
+
+    K_ff = K_np[np.ix_(free_dofs, free_dofs)]
+    K_fd = K_np[np.ix_(free_dofs, dir_set)]
+    F_f = F_np[free_dofs]
+    if F_f.ndim == 2:
+        F_f = F_f - (K_fd @ dir_vals)[:, None]
+    else:
+        F_f = F_f - K_fd @ dir_vals
+
+    return K_ff, F_f, free_dofs, dir_set, dir_vals
+
+
+def expand_dirichlet_solution(u_free, free_dofs, dir_dofs, dir_vals, n_total):
+    """Expand condensed solution back to full vector."""
+    dir_dofs, dir_vals = _normalize_dirichlet(dir_dofs, dir_vals)
+    u_free_arr = np.asarray(u_free, dtype=float)
+    if u_free_arr.ndim == 2:
+        u = np.zeros((n_total, u_free_arr.shape[1]), dtype=float)
+        u[free_dofs, :] = u_free_arr
+        u[dir_dofs, :] = np.asarray(dir_vals, dtype=float)
+    else:
+        u = np.zeros(n_total, dtype=float)
+        u[free_dofs] = u_free_arr
+        u[dir_dofs] = np.asarray(dir_vals, dtype=float)
+    return u

fluxfem/solver/history.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+
+@dataclass
+class NewtonIterRecord:
+    iter: int
+    res_inf: float
+    res_two: float
+    rel_res_inf: float
+    alpha: float
+    step_norm: float
+    lin_iters: Optional[int] = None
+    lin_converged: Optional[bool] = None
+    lin_residual: Optional[float] = None
+    nan_detected: bool = False
+    assemble_time: Optional[float] = None
+    linear_time: Optional[float] = None
+
+
+@dataclass
+class LoadStepResult:
+    load_factor: float
+    info: Any
+    solve_time: float
+    u: Any
+    iter_history: List[NewtonIterRecord] = field(default_factory=list)
+    exception: Optional[str] = None
+    meta: dict[str, Any] = field(default_factory=dict)