Linear solvers

Distance computations requiring a linear system solve can be accelerated using specialized solvers. The user can select any solver compatible with the Lineax library to perform these computations, and use additional custom JAXScape solvers wrapped into Lineax solvers.

Using lineax solvers

You can use any Lineax solver with ResistanceDistance:

import lineax as lx

# Conjugate Gradient
solver = lx.CG(rtol=1e-6, atol=1e-6)
distance = ResistanceDistance(solver=solver)

See the Lineax documentation for available solvers and their options.

JAXScape solvers

jaxscape.solvers.pyamgsolver.PyAMGSolver

A linear solver that uses PyAMG to solve a sparse linear system.

Example

from jaxscape.solvers import PyAMGSolver

solver = PyAMGSolver(rtol=1e-6, maxiter=100_000)
distance = ResistanceDistance(solver=solver)
dist = distance(grid)

# Custom AMG method
import pyamg
solver = PyAMGSolver(pyamg_method=pyamg.ruge_stuben_solver)

Warning

PyAMG must be installed to use this solver.

jaxscape.solvers.cholmodsolver.CholmodSolver

A linear solver that uses CHOLMOD (via cholespy) to solve a sparse linear system. Uses direct Cholesky factorization for symmetric positive definite matrices.

Example

from jaxscape.solvers import CholmodSolver

solver = CholmodSolver()
distance = ResistanceDistance(solver=solver)
dist = distance(grid)

Warning

cholespy must be installed to use this solver.

Advanced: BCOOLinearOperator

JAXScape provides a custom linear operator compatible with Lineax that leverages the BCOO sparse matrix format.

jaxscape.solvers.operator.BCOOLinearOperator

lineax.MatrixLinearOperator wrapper for jax.experimental.sparse.BCOO matrices.