jax_fem.solver module

jax_fem.solver module#

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

Solve the nonlinear problem using Newton’s method with configurable linear solvers.

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 dictionary for solver parameters and algorithms. Three solvers are currently available:

    The empty choice

    solver_options = {}

    will default to JAX solver as

    solver_options = {'jax_solver': {}}

    which will further default to

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

    The UMFPACK solver can be specified as

    solver_options = {'umfpack_solver': {}}

    The PETSc solver can be specified as

    solver_options = {
     'petsc_solver': {},
}

    which will default to

    solver_options = {
     'petsc_solver': {
         'ksp_type': 'bcgsl', # other choices can be, e.g., 'minres', 'gmres', 'tfqmr'
         'pc_type': 'ilu', # other choices can be, e.g., 'jacobi'
     }
}

    Other available options are

    solver_options = {
    'line_search_flag': False, # Line search method
    'initial_guess': initial_guess, # Same shape as sol_list
    'tol': 1e-5, # Absolute tolerance for residual vector (l2 norm), used in Newton's method
    'rel_tol': 1e-8, # Relative tolerance for residual vector (l2 norm), used in Newton's method
}

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 (dictionary)

  • adjoint_solver_options (dictionary)

Returns:

fwd_pred

Return type:

callable

Contents