Solver Configuration

This page documents how to configure jaxamg.solve(...), including solver/preconditioner choices and common AmgX options.

Configuration input modes

jaxamg.solve(...) accepts configuration in three ways:

config=<dict>: a full configuration dictionary.
keyword arguments (solver=..., max_iters=..., etc.).
both together, where keyword arguments override values in config.

Example:

x, info = jaxamg.solve(
    A,
    b,
    config={"solver": "CG", "max_iters": 100},
    max_iters=50,  # overrides config["max_iters"]
)

Default config

If you do not provide a config, JAX-AMG uses:

{
    "solver": "PBICGSTAB",
    "preconditioner": {
        "solver": "AMG",
        "algorithm": "CLASSICAL",
        "selector": "PMIS",
        "interpolator": "D2",
        "smoother": "BLOCK_JACOBI",
        "presweeps": 1,
        "postsweeps": 1,
        "coarse_solver": "DENSE_LU_SOLVER",
        "strength_threshold": 0.5,
        "max_levels": 100,
        "cycle": "V",
    },
    "convergence": "RELATIVE_INI",
    "tolerance": 1e-6,
    "max_iters": 1000,
    "norm": "L2",
}

User-supplied keys are merged into these defaults, so you only need to specify what you want to change. For example, config={"preconditioner": {"solver": "AMG"}} inherits all the Classical AMG settings above.

JAX-AMG also enables residual tracking internally (monitor_residual=1, store_res_history=1) so info["residual"] is always available.

MPI config

The communicator key can be set to MPI for standard CPU-based MPI or MPI_DIRECT for GPU-aware MPI. The default is MPI. To use GPU-aware MPI, ensure that your MPI installation supports it. For more details, see the MPI Guide.

config = {
    "solver": "PBICGSTAB",
    "preconditioner": {"solver": "MULTICOLOR_DILU"},
    "communicator": "MPI_DIRECT",
    "max_iters": 200,
    "tolerance": 1e-8,
}

Flat vs nested config

JAX-AMG supports both flat (simple) and nested config dict:

Flat config (simple dict with top-level keys like solver, tolerance).

config = {
    "solver": "CG",
    "preconditioner": {"solver": "JACOBI_L1"},
    "tolerance": 1e-8,
    "max_iters": 200,
}

Nested config (config_version: 2 with solver: {...} scope).

config = {
    "config_version": 2,
    "solver": {
        "solver": "PCG",
        "preconditioner": {
            "solver": "AMG",
            "smoother": "JACOBI_L1",
        },
        "tolerance": 1e-6,
        "max_iters": 200,
    },
}

Common configuration keys

These are commonly used and work in JAX-AMG:

Key	Type	Meaning
`solver`	`str`	Main solver type, e.g. `CG`, `PBICGSTAB`, `FGMRES`, `AMG`.
`preconditioner`	`str` or `dict`	Optional preconditioner solver (simple or nested).
`tolerance`	`float`	Convergence tolerance.
`max_iters`	`int`	Max solver iterations.
`convergence`	`str`	Convergence criterion (see below).
`norm`	`str`	Residual norm type, commonly `L2`.
`monitor_residual`	`0` or `1`	Enable residual monitoring.
`store_res_history`	`0` or `1`	Store residual history.
`print_solve_stats`	`0` or `1`	Print per-solve stats from AmgX.
`obtain_timings`	`0` or `1`	Collect setup/solve timings in AmgX.
`communicator`	`str`	MPI communication mode (`MPI` or `MPI_DIRECT`).

Convergence criteria

Value	Meaning
`ABSOLUTE`	Absolute residual threshold check.
`RELATIVE_INI_CORE`	Core variant of `RELATIVE_INI` (same criterion family, different internal residual handling).
`RELATIVE_MAX_CORE`	Core variant of `RELATIVE_MAX` (same criterion family, different internal residual handling).
`RELATIVE_INI`	True relative residual ‖r_k‖ / ‖r_0‖. Computes the explicit residual at every iteration (default).
`RELATIVE_MAX`	Max-norm relative residual.
`COMBINED_REL_INI_ABS`	Combined relative-initial and absolute criterion.

AMG-specific keys

These go inside the AMG solver/preconditioner scope:

Key	Meaning
`algorithm`	AMG algorithm: `CLASSICAL` (default) or `AGGREGATION`.
`selector`	Coarsening selector. Classical: `PMIS` (default), `HMIS`. Aggregation: `SIZE_2`, `SIZE_4`, `SIZE_8`, `MULTI_PAIRWISE`.
`interpolator`	Interpolation strategy, e.g. `D2` (default for Classical).
`smoother`	Smoother solver, e.g. `BLOCK_JACOBI` (default), `JACOBI_L1`, `MULTICOLOR_GS`, `MULTICOLOR_DILU`.
`presweeps`, `postsweeps`	Number of smoothing sweeps (default: 1).
`coarse_solver`	Coarse-level solver: `DENSE_LU_SOLVER` (default) or `NOSOLVER`.
`strength_threshold`	Strength of connection threshold for coarsening (default: 0.5).
`max_levels`	Maximum multigrid levels (default: 100).
`cycle`	Cycle type: `V` (default), `W`, `F`, `CG`, `CGF`.

Supported Solvers and Preconditioners

AmgX supports a broader set of solvers, preconditioners, and smoothers:

Value	Method
`AMG`	Algebraic multigrid (AMG) method.
`CG`	Conjugate gradient (CG) method for symmetric positive-definite systems.
`PCG`	Preconditioned CG method.
`PCGF`	Flexible PCG that supports changing/variable preconditioners across iterations.
`BICGSTAB`	Bi-conjugate gradient stabilized (BiCGStab) method for general nonsymmetric systems.
`PBICGSTAB`	Preconditioned BiCGSTAB method.
`GMRES`	Generalized minimal residual (GMRES) method for nonsymmetric systems.
`FGMRES`	Flexible GMRES that supports changing/variable preconditioners across iterations.
`IDR`	Induced dimension reduction (IDR) method for nonsymmetric systems.
`IDRMSYNC`	A variant of IDR designed for parallel performance optimization.
`JACOBI_L1`	L1-scaled Jacobi smoother/preconditioner.
`BLOCK_JACOBI`	Block Jacobi method.
`GS`	Gauss-Seidel preconditioner/smoother.
`MULTICOLOR_GS`	Parallelized Gauss-Seidel using graph coloring.
`MULTICOLOR_ILU`	Parallelized incomplete LU factorization (ILU) with graph coloring.
`MULTICOLOR_DILU`	Parallelized diagonal ILU variant with graph coloring.
`CHEBYSHEV`	Chebyshev iterative method.
`CHEBYSHEV_POLY`	Polynomial Chebyshev method.
`CF_JACOBI`	C/F-point Jacobi variant for AMG.
`DENSE_LU_SOLVER`	Dense LU direct solve.
`NOSOLVER`	No solver/preconditioner.