Skip to content

Linear solvers¤

ResistanceDistance and RSPDistance internally solve one or more sparse linear systems. The choice of solver directly affects both runtime and memory usage.

By default, when no solver is provided, JAXScape falls back to dense matrix inversion — this is only suitable for small graphs. For any non-trivial landscape, passing an explicit solver is strongly recommended. CholmodSolver is the preferred choice when memory permits.

lineax solvers¤

JAXScape accepts any Lineax solver. For example, to use the Conjugate Gradient method:

import lineax as lx
from jaxscape import ResistanceDistance

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

See the Lineax documentation for the full list of available solvers and their options.

JAXScape solvers¤

JAXScape provides two optional high-performance solvers as extras:

Solver Method Memory Best for
CholmodSolver Direct Cholesky factorization Higher Speed-critical workflows
PyAMGSolver Algebraic multigrid (iterative) Moderate Memory-constrained problems

Installation: 

uv add jaxscape --extra cholespy   # Cholesky solver
uv add jaxscape --extra pyamg      # PyAMG solver
uv add jaxscape --extra solvers    # Both solvers

CI/CD coverage

These optional solvers are not included in the standard CI test suite.

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.

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.

Advanced: BCOOLinearOperator¤

JAXScape exposes a lineax-compatible linear operator that wraps JAX's native BCOO sparse matrix format. This allows any Lineax solver to operate directly on sparse matrices without converting to a dense representation, and is used internally by both CholmodSolver and PyAMGSolver.

jaxscape.solvers.operator.BCOOLinearOperator ¤

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