autopdex.solver.linear_solve_petsc

autopdex.solver.linear_solve_petsc(mat, rhs, n_fields, solver, pc_type, verbose, free_dofs, tol=1e-08, **kwargs)[source]

Solve a linear system using the PETSc solver (requires PETSc and petsc4py to be installed).

This function solves a linear system using PETSc, with options for different solvers and preconditioners. The input matrix is first converted to a SciPy CSR matrix, and rows and columns corresponding to Dirichlet DOFs are optionally removed.

Parameters:

  • mat (jax.experimental.sparse.BCOO) – The input matrix in JAX BCOO format.

  • rhs (jnp.ndarray) – The (reduced/free) right-hand side vector.

  • n_fields (int) – The number of fields.

  • solver (str) – The type of solver to use.

  • pc_type (str) – The type of preconditioner to use.

  • verbose (int) – Verbosity level. If >= 1, timing and solver information is printed.

  • free_dofs (array or None) – Boolean array indicating which degrees of freedom are free.

  • tol (float) – The relative tolerance for the solver.

  • **kwargs – Additional keyword arguments for the solver.

Returns:

The solution vector.

Return type:

jax.numpy.array

Notes

  • The solver settings can also be set from the command line. See PETSc and petsc4py documentation.