fol.solvers

Solve strategies provided by FoLax.

Finite element linear solver class

Authors: Reza Najian Asl, RezaNajian Date: May, 2024 License: FOL/LICENSE

class fol.solvers.fe_solver.FiniteElementSolver(fe_solver_name, fe_loss_function, fe_solver_settings={})

Bases: Solver

Base solver for finite element systems driven by a FiniteElementLoss.

This class provides a common interface for solvers that operate on FE loss functions and require assembling and solving linear systems of the form:

K(u, p) * delta = -R(u, p)

where K is a sparse tangent (Jacobian) matrix and R is a residual vector produced by a fol.loss_functions.fe_loss.FiniteElementLoss.ComputeJacobianMatrixAndResidualVector instance.

The class mainly manages linear-solver selection and configuration. It does not define a problem-specific Solve routine; derived solver classes implement their own Solve method (e.g., linear residual-based solvers, nonlinear Newton solvers, adjoint solvers) and call LinearSolve to solve the assembled linear system.

Parameters:

• fe_solver_name (str) – Name identifier for the solver instance.

• fe_loss_function (FiniteElementLoss) – Finite element loss object that provides problem-specific residuals and tangent matrices.

• fe_solver_settings (dict, optional) – Solver configuration dictionary. If it contains the key "linear_solver_settings", the settings are merged with defaults.

fe_loss_function

Loss object used to assemble residuals and tangent matrices.

Type:

FiniteElementLoss

fe_solver_settings

User-provided solver settings dictionary.

Type:

dict

linear_solver_settings

Linear solver configuration with default keys: "solver", "tol", "atol", "maxiter", and "pre-conditioner".

Type:

dict

LinearSolve

Function handle assigned in Initialize() to the selected linear solve routine.

Type:

callable

Notes

Supported linear solver backends are selected via linear_solver_settings["solver"]. The available options are: "JAX-bicgstab", "JAX-direct", and PETSc-based Krylov methods "PETSc-bcgsl", "PETSc-tfqmr", "PETSc-minres", "PETSc-gmres". If PETSc is requested but unavailable, the solver falls back to "JAX-bicgstab".

Initialize()

Initialize and select the linear solver backend.

This method merges user-provided linear_solver_settings (if present) with defaults and assigns the LinearSolve callable to one of the available backends.

Returns:

None

Raises:

ValueError – If an unknown solver name is provided in linear_solver_settings["solver"].

JaxBicgstabLinearSolver(tangent_matrix, residual_vector, dofs_vector)

Solve the linear system using JAX BiCGSTAB.

The solver computes a DOF increment by solving:

tangent_matrix * delta_dofs = -residual_vector

using the iterative BiCGSTAB method from jax.scipy.sparse.linalg. The optional initial guess is taken as dofs_vector.

Parameters:

• tangent_matrix (jax.experimental.sparse.BCOO) – Global tangent (Jacobian) matrix in BCOO sparse format.

• residual_vector (jax.numpy.ndarray) – Global residual vector.

• dofs_vector (jax.numpy.ndarray) – Initial guess for the iterative solver.

Returns:

Increment vector delta_dofs solving the linear system.

Return type:

jax.numpy.ndarray

JaxDirectLinearSolver(tangent_matrix, residual_vector, dofs_vector)

Solve the linear system using a direct sparse solve.

The input BCOO matrix is converted to a SciPy CSR sparse matrix and solved using jax.experimental.sparse.linalg.spsolve with the right hand side -residual_vector.

Parameters:

• tangent_matrix (jax.experimental.sparse.BCOO) – Global tangent (Jacobian) matrix in BCOO sparse format.

• residual_vector (jax.numpy.ndarray) – Global residual vector.

• dofs_vector (jax.numpy.ndarray) – Unused by the direct solver (kept for a consistent interface).

Returns:

Increment vector delta_dofs solving the linear system.

Return type:

jax.numpy.ndarray

PETScLinearSolver(tangent_matrix, residual_vector, dofs_vector)

Solve the linear system using PETSc KSP.

The input BCOO matrix is converted to a SciPy CSR matrix and then wrapped into a PETSc AIJ matrix. The right hand side is set to -residual_vector. The KSP type is selected from the solver name provided in Initialize() (e.g., gmres, minres, tfqmr). The PETSc preconditioner type is set via linear_solver_settings["pre-conditioner"].

Parameters:

• tangent_matrix (jax.experimental.sparse.BCOO) – Global tangent (Jacobian) matrix in BCOO sparse format.

• residual_vector (jax.numpy.ndarray) – Global residual vector.

• dofs_vector (jax.numpy.ndarray) – Unused by PETSc (kept for a consistent interface).

Returns:

Increment vector delta_dofs returned by PETSc.

Return type:

numpy.ndarray

Raises:

RuntimeError – If PETSc is not available in the current environment.

Adjoint-based Finite Element solver

Authors: Reza Najian Asl, RezaNajian Date: January, 2025 License: FOL/LICENSE

class fol.solvers.adjoint_fe_solver.AdjointFiniteElementSolver(adj_fe_solver_name, fe_response, adj_fe_solver_settings={})

Bases: FiniteElementSolver

Adjoint finite element solver for gradient and sensitivity analysis.

This class solves the adjoint system associated with a forward finite element problem. The adjoint formulation is typically used in PDE-constrained optimization, inverse problems, and sensitivity analysis to efficiently compute gradients of a scalar objective (loss) function with respect to control variables.

The adjoint system is constructed using a FiniteElementResponse object, which provides the loss functional, the adjoint Jacobian matrix, and the adjoint right-hand side vector.

The class inherits from FiniteElementSolver and reuses its linear system solution infrastructure.

Mathematical Formulation

Let the forward problem be defined as:

R(u, m) = 0

where:

u : state (degrees of freedom), m : control variables.

Let the scalar objective (loss) be:

J(u, m)

The adjoint variable λ satisfies the adjoint equation:

(∂R/∂u)ᵀ λ = - ∂J/∂u

where:

∂R/∂u : Jacobian of the residual with respect to the state, ∂J/∂u : gradient of the loss with respect to the state.

Boundary conditions are assumed to be applied directly to the adjoint Jacobian and right-hand side prior to solving the system.

param adj_fe_solver_name:

Name of the adjoint finite element solver instance.

type adj_fe_solver_name:

str

param fe_response:

Object responsible for defining the loss functional and assembling the adjoint Jacobian matrix and right-hand side vector.

type fe_response:

FiniteElementResponse

param adj_fe_solver_settings:

Dictionary containing solver-specific settings such as linear solver type, tolerances, or preconditioning options. Default is an empty dictionary.

type adj_fe_solver_settings:

dict, optional

fe_response

Reference to the finite element response object used to assemble adjoint quantities.

Type:

FiniteElementResponse

Solve(current_control_vars, current_dofs, current_adjoint_dofs)

Assemble and solve the adjoint linear system.

This method constructs the adjoint Jacobian matrix and right-hand side vector based on the current control variables and forward solution. The resulting linear system is then solved for the adjoint degrees of freedom.

The right-hand side vector is multiplied by -1 to ensure consistency with the solver sign convention used in LinearSolve().

Parameters:

• current_control_vars (jnp.array) – Current values of the control variables.

• current_dofs (jnp.array) – Current forward solution degrees of freedom.

• current_adjoint_dofs (jnp.array) – Initial guess or previous values of the adjoint degrees of freedom.

Returns:

The updated adjoint degrees of freedom obtained by solving the adjoint linear system.

Return type:

jnp.array

Linear residual-based Finite Element solver

Authors: Reza Najian Asl, RezaNajian Date: May, 2024 License: FOL/LICENSE

class fol.solvers.fe_linear_residual_based_solver.FiniteElementLinearResidualBasedSolver(fe_solver_name, fe_loss_function, fe_solver_settings={})

Bases: FiniteElementSolver

Linear residual-based finite element solver.

This solver advances the solution of a (typically) linear finite element system using a residual-based update. Given the current state (DOFs) and control variables, it:

1. Applies Dirichlet boundary conditions to the current DOF vector.

2. Assembles the Jacobian matrix and residual vector of the FE loss/residual.

3. Solves a linear system to obtain an increment (update) in DOFs.

4. Returns the updated DOF vector.

The assembly and boundary-condition operations are delegated to the underlying loss function (self.fe_loss_function), while the linear solve is handled by FiniteElementSolver.LinearSolve().

Linear Residual Solver Formulation

Let u be the vector of degrees of freedom (state) and m the control variables. After enforcing Dirichlet boundary conditions on u, the loss function provides:

• J: the Jacobian matrix (e.g., tangent stiffness matrix)

• r: the residual vector

The solver computes an increment Δu from a linear system of the form:

J(u, m) Δu = r(u, m)

and updates the state as:

u_new = u_cu + Δu

where u_cu is the given/current DOF vector after applying Dirichlet boundary conditions.

Notes

• This class is intended for linear (or linearized) problems where a single residual-based correction is sufficient per solve call.

• Sign conventions for the residual (whether the system is J Δu = r or J Δu = -r) are assumed to be consistent with the implementation of FiniteElementSolver.LinearSolve() and the loss function assembly.

Solve(current_control_vars, current_dofs)

Assemble and solve a linear residual-based update for the DOFs.

This method applies Dirichlet boundary conditions to the current DOFs, then uses the FE loss function to compute the Jacobian matrix and residual vector. It solves the resulting linear system to obtain the DOF increment and returns the updated DOF vector.

Parameters:

• current_control_vars (jnp.array) – Current values of the control variables (design/parameters) used to assemble the system.

• current_dofs (jnp.array) – Current degrees of freedom (state vector) serving as the starting point for the update.

Returns:

Updated degrees of freedom after applying Dirichlet boundary conditions and adding the computed increment.

Return type:

jnp.array

Nonlinear finite element solver with incremental Newton–Raphson iterations and Gauss-point state variable updates

Authors: Reza Najian Asl, RezaNajian Date: July, 2024 License: FOL/LICENSE

class fol.solvers.fe_nonlinear_residual_based_solver_with_history_update.FiniteElementNonLinearResidualBasedSolverWithStateUpdate(fe_solver_name, fe_loss_function, fe_solver_settings={}, history_plot_settings={})

Bases: FiniteElementNonLinearResidualBasedSolver

Nonlinear residual-based finite element solver with Gauss-point state updates.

This solver extends FiniteElementNonLinearResidualBasedSolver by coupling global Newton–Raphson iterations for the degrees of freedom with local updates of Gauss-point (integration-point) state variables such as internal or history-dependent quantities.

The solution procedure is performed using incremental load steps defined by nonlinear_solver_settings["load_incr"]. For each load step, the solver applies boundary conditions, performs Newton iterations, updates local state variables, and records convergence information.

This class is intended for nonlinear, history-dependent problems such as elastoplasticity, damage, or viscoelasticity, where consistent updates of integration-point state variables are required.

Solve(current_control_vars, current_dofs, current_state=None, return_all_steps=False)

Solve the nonlinear finite element system using incremental loading.

The solver advances the solution through a sequence of load steps. Within each load step, Newton–Raphson iterations are performed to enforce global equilibrium. At every Newton iteration, the Jacobian matrix, residual vector, and updated Gauss-point state variables are assembled by the loss function.

Convergence is declared when at least one of the following conditions is satisfied: the residual norm falls below the absolute tolerance, the update norm falls below the relative tolerance, or the maximum number of Newton iterations is reached.

Parameters:

• current_control_vars (jax.numpy.ndarray) – Control variables used to assemble the system (e.g., material parameters).

• current_dofs (jax.numpy.ndarray) – Initial degrees of freedom representing the current global state.

• current_state (jax.numpy.ndarray, optional) – Initial Gauss-point state variables. If not provided, a zero state field is initialized based on the element type and integration scheme.

• return_all_steps (bool, optional) – If False (default), only the final load step solution and state are returned. If True, the solution and state for all load steps are returned.

Returns:

The updated degrees of freedom, the corresponding Gauss-point state variables, and a dictionary containing convergence history for each load step.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, dict]

Raises:

ValueError – If the residual norm becomes NaN during the Newton iterations, indicating a breakdown of the nonlinear solve.

Nonlinear finite element solver with incremental Newton–Raphson iterations

Authors: Reza Najian Asl, RezaNajian Date: July, 2024 License: FOL/LICENSE

class fol.solvers.fe_nonlinear_residual_based_solver.FiniteElementNonLinearResidualBasedSolver(fe_solver_name, fe_loss_function, fe_solver_settings={}, history_plot_settings={})

Bases: FiniteElementLinearResidualBasedSolver

Nonlinear residual-based finite element solver using incremental Newton–Raphson.

This solver extends FiniteElementLinearResidualBasedSolver to support nonlinear problems by repeatedly assembling the tangent (Jacobian) matrix and residual vector and applying Newton–Raphson updates to the global degrees of freedom (DOFs).

Loading is applied incrementally according to self.nonlinear_solver_settings["load_incr"]. For each load step, the solver iterates until convergence is detected or the maximum number of Newton iterations is reached. Convergence history (residual norms and update norms) is stored per load step and can optionally be plotted via PlotHistoryDict().

Typical applications include problems with geometric nonlinearities, nonlinear material behavior (e.g., hyperelasticity, plasticity, damage), or nonlinear boundary-condition effects.

Initialize()

Initialize the solver and merge nonlinear and plotting settings.

This method calls the base class initialization and updates the nonlinear solver settings from fe_solver_settings["nonlinear_solver_settings"] if provided. It also merges the user-specified history plotting settings with the default plotting configuration.

Solve(current_control_vars, current_dofs_np)

Solve the nonlinear FE system using incremental Newton–Raphson iterations.

The load is applied over load_incr steps. For each load step, Dirichlet boundary conditions are applied at the current load factor, the Jacobian and residual are assembled via the loss function, and a Newton update is computed using LinearSolve(). The DOFs are updated until convergence is detected or the iteration limit is reached. Convergence metrics are recorded per load step and can be plotted via PlotHistoryDict().

Convergence is considered achieved when at least one of the following conditions is satisfied: - Residual norm < abs_tol - Update norm < rel_tol - Iteration count reaches maxiter

Parameters:

• current_control_vars (jax.numpy.ndarray) – Control variables passed to the loss function (e.g., material or design parameters).

• current_dofs_np (jax.numpy.ndarray) – Initial DOF vector. Converted to a JAX array internally.

Returns:

Converged DOF vector at the end of the final load step.

Return type:

jax.numpy.ndarray

Raises:

ValueError – If the residual norm becomes NaN during Newton iterations.