jax_fem.solver module

jax_fem.solver module#

jax_fem.solver.solver(problem, solver_options={})[source]#

Solve a nonlinear problem (Newton by default, or arc-length / dynamic relaxation).

The solver imposes Dirichlet B.C. with “row elimination” method. Conceptually,

\[\begin{split}r(u) = D \, r_{\text{unc}}(u) + (I - D)u - u_b \\ A = \frac{\text{d}r}{\text{d}u} = D \frac{\text{d}r}{\text{d}u} + (I - D)\end{split}\]

where:

  • \(r_{\text{unc}}: \mathbb{R}^N\rightarrow\mathbb{R}^N\) is the residual function without considering Dirichlet boundary conditions.

  • \(u\in\mathbb{R}^N\) is the FE solution vector.

  • \(u_b\in\mathbb{R}^N\) is the vector for Dirichlet boundary conditions, e.g.,

    \[\begin{split}u_b = \begin{bmatrix} 0 \\ 0 \\ 2 \\ 3 \end{bmatrix}\end{split}\]

  • \(D\in\mathbb{R}^{N\times N}\) is the auxiliary matrix for masking, e.g.,

    \[\begin{split}D = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \end{bmatrix}\end{split}\]

  • \(I\in\mathbb{R}^{N\times N}\) is the ientity matrix, e.g.,

    \[\begin{split}I = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix}\end{split}\]

  • \(A\in\mathbb{R}^{N\times N}\) is the tangent stiffness matrix (the global Jacobian matrix).

Notes

  • TODO: Show some comments for linear multipoint constraint handling.

Parameters:

  • problem (Problem) – The nonlinear problem to solve

  • solver_options (dict) –

    Configuration for the nonlinear solve. Use exactly one top-level method key— newton, arc_length, or dynamic_relax. Nest the linear solver and method-specific options inside that block.

    Newton (default nonlinear backend):

    solver_options = {
    'newton': {
        'tol': 1e-5,
        'rel_tol': 1e-8,
        'line_search_flag': False,
        'initial_guess': initial_guess,
        'linear': {'petsc_solver': {}},
    },
}

    Linear solvers (keys under linear in any method block). Four backends are currently available:

    Examples nested under newton:

    solver_options = {'newton': {'linear': {'jax_solver': {}}}}

solver_options = {'newton': {'linear': {'spsolve_solver': {}}}}

solver_options = {
    'newton': {
        'linear': {
            'petsc_solver': {
                'ksp_type': 'bcgsl',  # e.g. 'minres', 'gmres', 'tfqmr'
                'pc_type': 'ilu',     # e.g. 'jacobi'
            },
        },
    },
}

solver_options = {'newton': {'linear': {'amgx_solver': {'cfg_path': 'path/to/amgx.json'}}}}

    Defaults. Omitted keys are filled in as follows.

    Newton (inside a newton block, or implied when no method key is given):

    • tol → 1e-6 (absolute residual \(\ell_2\) norm)

    • rel_tol → 1e-8 (relative to the initial residual)

    • line_search_flag → False

    • initial_guess → zero displacement vector

    • linear: The following are all equivalent for the linear solve:

      solver_options = {}
solver_options = {'newton': {}}
solver_options = {'newton': {'linear': {}}}
solver_options = {'newton': {'linear': {'jax_solver': {}}}}
solver_options = {'newton': {'linear': {'jax_solver': {'precond': True}}}}

    • {'jax_solver': {}} → precond → True

    • {'petsc_solver': {}} → ksp_type → 'bcgsl'; pc_type → 'ilu'

    • {'amgx_solver': {}} → cfg_path → None (built-in BICGSTAB + AMG)

    Arc-length (Crisfeld; control is required; set return_info to obtain continuation metadata):

    solver_options = {
    'arc_length': {
        'control': 'displacement',  # or 'force' (needs q_vec_aux)
        'return_info': True,
        'Delta_l': 0.1,
        'linear': {'petsc_solver': {}},
        'newton': {'tol': 1e-6},  # optional polish at lambda=1
    },
}

    Dynamic relaxation (useful for buckling paths):

    solver_options = {
    'dynamic_relax': {
        'tol': 1e-8,
        'linear': {'spsolve_solver': {}},
    },
}

    Legacy flat dict. For backward compatibility, a dict with no method key is still accepted and interpreted as Newton. Linear and Newton keys may appear at the top level, e.g.:

    solver_options = {'petsc_solver': {}, 'tol': 1e-5}

    is equivalent to specifying:

    solver_options = {'newton': {'linear': {'petsc_solver': {}}, 'tol': 1e-5}}

Returns:

sol_list

Return type:

list

jax_fem.solver.ad_wrapper(problem, solver_options={}, adjoint_solver_options={})[source]#

Automatic differentiation wrapper for the forward problem.

Parameters:

  • problem (Problem)

  • solver_options (dict) – Same layout as solver() (nonlinear method + nested linear).

  • adjoint_solver_options (dict) – Linear solver options for the adjoint solve only (flat dict, e.g. {'petsc_solver': {}}).

Returns:

fwd_pred

Return type:

callable

Contents