autopdex.solver.linear_solve_pyamg

autopdex.solver.linear_solve_pyamg(mat, rhs, solver, pc_type, verbose, free_dofs, **kwargs)[source]

Solve a linear system using the PyAMG solver (requires pyamg to be installed).

This function solves a linear system using PyAMG, an algebraic multi-grid solver 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.

  • solver (str) – The type of solver to use (‘cg’, ‘bcgs’, or ‘gmres’).

  • pc_type (str) – The type of preconditioner to use (‘ruge stuben’ or ‘smoothed aggregation’).

  • 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.

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

Returns:

The solution vector.

Return type:

jax.numpy.array