Solver#

Solvers#

class fluxfem.solver.LinearSolver(method: str = 'spsolve', tol: float = 1e-08, maxiter: int = 200)#

Lightweight wrapper for solving linear systems with optional Dirichlet BCs.

Supports dense arrays or FluxSparseMatrix and can either condense or enforce Dirichlet conditions before solving with the chosen backend.

class fluxfem.solver.NonlinearSolver(space: FESpaceClosure, res_form: Any, params: Any, *, tol: float = 1e-08, maxiter: int = 20, linear_method: str = 'spsolve', line_search: bool = False, max_ls: int = 10, ls_c: float = 0.0001, linear_tol: float | None = None, dirichlet: DirichletBC | tuple[ndarray, ndarray] | None = None, external_vector: ndarray | None = None, linear_maxiter: int | None = None, jacobian_pattern: Any | None = None)#

Backward-compatible Newton-based nonlinear solver.

This is a thin wrapper around newton_solve kept for legacy code paths. Prefer NonlinearAnalysis + NewtonSolveRunner for new workflows.

class fluxfem.solver.LinearAnalysis(space: Any, bilinear_form: Callable[[jax.tree_util.register_pytree_node_class, Any], jax.numpy.ndarray], params: Any, base_rhs_vector: ndarray, dirichlet: tuple[ndarray, ndarray] | None = None, pattern: Any | None = None, dtype: Any = jax.numpy.float64)#

Bundle linear problem data for a single solve or a load-scaled sequence.

The matrix is assembled once from bilinear_form; the RHS is scaled by the load factor.

class fluxfem.solver.NonlinearAnalysis(space: Any, residual_form: Callable[[jax.tree_util.register_pytree_node_class, jax.numpy.ndarray, Any], jax.numpy.ndarray], params: Any, base_external_vector: ndarray | None = None, dirichlet: tuple[ndarray, ndarray] | None = None, extra_terms: list[Callable[[ndarray], tuple[ndarray, ndarray] | tuple[ndarray, ndarray, dict[str, Any]] | None]] | None = None, assembly_policy: Any | None = None, jacobian_pattern: Any | None = None, dtype: Any = jax.numpy.float64)#

Bundle problem data needed for a Newton solve with load scaling.

Variables:

  • space (Any) – FE space containing topology/dofs.

  • residual_form (Any) – Internal residual form (e.g., neo_hookean_residual_form).

  • params (Any) – Parameters forwarded to the residual form.

  • base_external_vector (Any | None) – 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]).

  • assembly_policy (Any | None) – Optional volume assembly execution policy forwarded to residual/Jacobian assembly.

  • jacobian_pattern (Any | None) – Optional sparsity pattern to reuse between load steps.

  • dtype (Any) – dtype for the solution vector (defaults to float64).

class fluxfem.solver.LinearSolveRunner(analysis: LinearAnalysis, config: LinearSolveConfig)#

Solve linear systems for one or more load factors using a unified interface.

class fluxfem.solver.NewtonSolveRunner(analysis: NonlinearAnalysis, config: NewtonLoopConfig)#

Run one or more Newton solves across load factors.

This orchestrates load stepping, assembles external load per step, and returns the full (Dirichlet-expanded) solution and per-step history.

class fluxfem.solver.LinearSolveConfig(method: str = 'spsolve', tol: float = 1e-08, 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)#

Control parameters for the linear solve with optional load scaling.

class fluxfem.solver.NewtonLoopConfig(tol: float = 1e-08, atol: float = 0.0, maxiter: int = 20, line_search: bool = False, max_ls: int = 10, ls_c: float = 0.0001, linear_solver: str = 'spsolve', linear_maxiter: int | None = None, linear_tol: float | None = None, linear_preconditioner: Any | None = None, petsc_ksp_type: str = 'gmres', petsc_pc_type: str = 'none', petsc_rtol: float | None = None, petsc_atol: float | None = None, petsc_max_it: int | None = None, petsc_options: dict[str, Any] | None = None, petsc_use_pmat: bool = False, matfree_mode: str = 'linearize', load_sequence: Sequence[float] | None = None, n_steps: int = 1)#

Control parameters for the Newton loop with load stepping.

Coupled Systems#

class fluxfem.solver.CoupledSystem(K_u: jax.tree_util.register_pytree_node_class, F_u: jax.numpy.ndarray)#

JAX-native coupled system with sparse-first assembly.

class fluxfem.solver.CoupledSystemBuilder(system: CoupledSystem)#

Minimal JAX-native builder for structural blocks and remote springs.

class fluxfem.solver.DirichletSpec(field: 'str', value: 'float | Sequence[float] | np.ndarray | jnp.ndarray | None' = None, nodes: 'int | Sequence[int] | np.ndarray | None' = None, components: 'int | Sequence[int] | np.ndarray | None' = None, local_dofs: 'int | Sequence[int] | np.ndarray | None' = None)#
class fluxfem.solver.ConstraintSpec(kind: 'str', master: 'str', slave: 'str', C: 'np.ndarray | jnp.ndarray | None' = None, rho: 'float' = 0.0, F_contact: 'np.ndarray | jnp.ndarray | None' = None, contact_obj: 'Any | None' = None, family: 'str | None' = None, enforcement: 'str | None' = None, law: 'str | None' = None, formulation: 'str | None' = None, value_dim: 'int | None' = None, weak_form: 'Any | None' = None, state: 'Any | None' = None, params: 'Any | None' = None, residual: 'Any | None' = None, scale: 'float' = 1.0, residual_sign: 'float' = -1.0, normal_source: 'str' = 'master', sparse: 'bool | None' = None, batch_jac: 'bool | None' = None, ref_point: 'np.ndarray | jnp.ndarray | None' = None, slave_coords: 'np.ndarray | jnp.ndarray | None' = None, weights: 'np.ndarray | jnp.ndarray | None' = None, normalize_weights: 'bool' = True, embedding: 'Any | None' = None)#

Sparse#

class fluxfem.solver.SparsityPattern(rows: jnp.ndarray, cols: jnp.ndarray, n_dofs: int, idx: jnp.ndarray | None = None, diag_idx: jnp.ndarray | None = None, perm: jnp.ndarray | None = None, indptr: jnp.ndarray | None = None, indices: jnp.ndarray | None = None)#

Jacobian sparsity pattern (rows/cols) that is independent of the solution.

class fluxfem.solver.FluxSparseMatrix(rows_or_pattern: SparsityPattern | ArrayLike, cols: ArrayLike | None = None, data: ArrayLike | None = None, n_dofs: int | None = None, meta: dict | None = None)#

Sparse matrix wrapper (COO) with a fixed pattern and mutable values. - pattern stores rows/cols/n_dofs (optionally idx for dense scatter) - data stores the numeric values for the current nonlinear iterate

class fluxfem.solver.FluxSparseOperator(rows: ndarray, cols: ndarray, data: ndarray, shape: tuple[int, int], meta: dict | None = None)

Sparse operator wrapper (COO) for rectangular or square operators.

This is intentionally narrower than FluxSparseMatrix: - stores raw COO rows/cols/data - tracks a general shape=(n_rows, n_cols) - supports dense conversion and forward/adjoint matvec

It is intended as the first rectangular sparse abstraction for Petrov-Galerkin-style bilinear assembly without destabilizing existing square-matrix solver paths.

Dirichlet#

fluxfem.solver.enforce_dirichlet_dense(K, F, dofs, vals)#

Apply Dirichlet conditions directly to stiffness/load (dense).

fluxfem.solver.enforce_dirichlet_sparse(A: jax.tree_util.register_pytree_node_class, F, dofs, vals)#

Apply Dirichlet conditions to FluxSparseMatrix + load (CSR).

fluxfem.solver.free_dofs(n_dofs: int, dir_dofs) → ndarray#

Return free DOF indices given total DOFs and Dirichlet DOFs.

fluxfem.solver.condense_dirichlet_fluxsparse(A: jax.tree_util.register_pytree_node_class, F, dofs, vals)#

Condense Dirichlet DOFs for a FluxSparseMatrix. Returns: (K_ff, F_free, free_dofs, dir_dofs, dir_vals)

fluxfem.solver.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)

fluxfem.solver.expand_dirichlet_solution(u_free, free_dofs, dir_dofs, dir_vals, n_total)#

Expand condensed solution back to full vector.

Iterative#

fluxfem.solver.cg_solve(A: Any, b: jnp.ndarray, *, x0: jnp.ndarray | None = None, tol: float = 1e-08, maxiter: int | None = 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).

fluxfem.solver.cg_solve_jax(A: Any, b: jnp.ndarray, *, x0: jnp.ndarray | None = None, tol: float = 1e-08, maxiter: int | None = 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).

Block matrices#

fluxfem.solver.block_diag(*, order: Sequence[str | int] | None = None, **blocks: Any) → dict[str | int, Any]#

Build a dict of diagonal blocks with an optional explicit field order.

fluxfem.solver.make_block_matrix(*, diag: Mapping[str | int, Any] | Sequence[Any], rel: Mapping[tuple[str | int, str | int], Any] | None = None, add_contiguous: Any | None = None, sizes: Mapping[str | int, int] | None = None, symmetric: bool = False, transpose_rule: str = 'T') → FluxBlockMatrix#

Build a lazy FluxBlockMatrix from diagonal blocks, optional relations, and a full matrix.

class fluxfem.solver.FluxBlockMatrix(blocks: dict[str | int, dict[str | int, Any]], *, sizes: Mapping[str | int, int], order: Sequence[str | int] | None = None, symmetric: bool = False, transpose_rule: str = 'T')#

Lazy block-matrix container that assembles on demand.

Nonlinear#

fluxfem.solver.newton_solve(space, res_form: ResidualForm[Any], u0: ArrayLike, params: Any, *, tol: float = 1e-08, atol: float = 0.0, maxiter: int = 20, linear_solver: str = 'spsolve', linear_maxiter: int | None = None, linear_tol: float | None = None, linear_preconditioner: object | None = None, petsc_ksp_type: str = 'gmres', petsc_pc_type: str = 'none', petsc_rtol: float | None = None, petsc_atol: float | None = None, petsc_max_it: int | None = None, petsc_options: dict[str, Any] | None = None, petsc_use_pmat: bool = False, matfree_mode: str = 'linearize', 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 = 0.0001, external_vector: np.ndarray | None = None, jacobian_pattern: SparsityPattern | None = None, extra_terms: list[ExtraTerm] | None = None, assembly_policy: Any | None = None) → tuple[np.ndarray, SolverResult]#

Gridap-style Newton–Raphson solver on free DOFs only.

  • Unknown vector = free DOFs (Dirichlet eliminated).

  • Residual/Jacobian are assembled on full DOFs; we slice to free DOFs.

  • Convergence: ||R_free||_inf < max(atol, tol * ||R_free0||_inf).

  • external_vector: optional global RHS (internal - external).

  • CG path accepts an operator with matvec that acts on free DOFs via a wrapper.

  • cg_matfree uses JVP/linearize to form a matrix-free matvec (no global Jacobian).

  • linear_preconditioner: forwarded to cg_solve/cg_solve_jax or PETSc shell (None | “jacobi” | “block_jacobi” | “diag0” | callable).

  • petsc_* options configure the inner PETSc shell solve when linear_solver=”petsc_shell”.

  • matfree_cache: optional dict for reusing matrix-free preconditioners across calls.

  • linear_tol: CG tolerance (defaults to 0.1 * tol if not provided).

  • jacobian_pattern: optional SparsityPattern to reuse sparsity across load steps.

  • extra_terms: optional list of callbacks returning (K, f[, metrics]) for extra terms.

fluxfem.solver.solve_nonlinear(space, residual_form: Callable[[jax.tree_util.register_pytree_node_class, jax.numpy.ndarray, Any], jax.numpy.ndarray], params: Any, *, dirichlet: tuple[ndarray, ndarray] | None = None, base_external_vector: ndarray | None = None, extra_terms: list[Callable[[ndarray], tuple[ndarray, ndarray] | tuple[ndarray, ndarray, dict[str, Any]] | None]] | None = None, assembly_policy: Any | None = None, dtype=jax.numpy.float64, maxiter: int = 20, tol: float = 1e-08, atol: float = 1e-10, linear_solver: str = 'spsolve', linear_maxiter: int | None = None, linear_tol: float | None = None, linear_preconditioner=None, petsc_ksp_type: str = 'gmres', petsc_pc_type: str = 'none', petsc_rtol: float | None = None, petsc_atol: float | None = None, petsc_max_it: int | None = None, petsc_options: dict[str, Any] | None = None, petsc_use_pmat: bool = False, matfree_mode: str = 'linearize', line_search: bool = False, max_ls: int = 10, ls_c: float = 0.0001, n_steps: int = 1, jacobian_pattern=None, u0: ndarray | None = None) → tuple[ndarray, list[LoadStepResult]]#

Convenience wrapper: build NonlinearAnalysis and run NewtonSolveRunner.