Skip to content

Linear solvers

Every Newton iteration solves one linear system, and the backend is what solves it. Pass one as linear= on a solver: a name for the defaults, or the object with settings (linear solvers).

Recommended defaults

Leave linear="auto", which picks KLU on a CPU host, cuDSS on a GPU, and Dense at or below ten unknowns. Pass a backend object only to change a setting.

Direct backends

Each one is a name on linear= and a class here.

diff_soup.linear.Dense

A dense LU in JAX, on the CPU and on GPUs.

It factors float64 and float32 matrices, real or complex. linear="dense" selects it, and "auto" picks it for systems of at most ten unknowns. Its factor needs no SolveContext.

diff_soup.linear.KLU

SuiteSparse KLU, a sparse direct solver on the CPU.

It factors float64 matrices only. linear="klu" selects it with default settings, and "auto" picks it on a CPU host.

Parameters:

Name Type Description Default
ordering

The fill-reducing ordering of each block, "amd" (default) or "colamd".

 required
btf

True (default) permutes to block triangular form first and factors each diagonal block alone.

 required
scale

The row scaling before pivoting, "max" (default), "sum" or "none".

 required
pivot_tolerance

The diagonal pivoting threshold in (0, 1], where 0.001 (default) favours the diagonal and 1.0 is partial pivoting.

 required

Raises:

Type Description
ConfigurationError

For a setting outside these values.

SuperLU_MT on more than one thread

SuperLU_MT's threaded factorization can return a wrong solution and report SolveStatus.SUCCESS, an upstream defect. A source build uses every core the process may run on unless OMP_NUM_THREADS caps it. Set OMP_NUM_THREADS=1, or use KLU.

diff_soup.linear.SuperLU

SuperLU_MT, a threaded sparse direct solver on the CPU.

It factors float64 matrices only. linear="superlu" selects it with default settings, and "auto" never picks it. It needs a source build with SuperLU_MT.

Parameters:

Name Type Description Default
ordering

The fill-reducing column ordering, "colamd" (default), "mmd_ata", "mmd_at_plus_a" or "natural".

 required
panel_size

The panel width, where 0 (default) keeps SuperLU_MT's own value.

 required
relax

The relaxed supernode size, where 0 (default) keeps SuperLU_MT's own value.

 required
pivot_threshold

The diagonal pivoting threshold in [0, 1], where 1.0 (default) is partial pivoting.

 required

Raises:

Type Description
ConfigurationError

For a setting outside these values.

CuDSS

NVIDIA cuDSS, a sparse direct solver on CUDA GPUs.

It factors float64 and float32 matrices. linear="cudss" selects it with default settings, and "auto" picks it on a GPU. ds.linear.CuDSS is None when the cuDSS extension is not installed.

Parameters:

Name Type Description Default
reorder

The fill-reducing ordering, "amd" or "metis", where None takes the layout's reorder.

 required
hybrid_memory

True keeps part of the factor in host memory when it outgrows the GPU, and None takes the layout's.

 required
hybrid_device_memory_limit_mb

The GPU memory a hybrid factor may use, in MB, which needs hybrid_memory=True.

 required
deterministic

True runs cuDSS's deterministic mode, which allows no refinement, downcast or symmetric factorization.

 required
ir_nsteps

The iterative refinement steps per solve, where None takes 0 for a float64 factor and 1 otherwise.

 required
symmetric_factorization

True factors an analysis declared symmetric with cuDSS's symmetric factorization instead of a general LU.

 required
zero_fill

True zeros the solution buffer before every solve, and None does so for float32 factors only.

 required
downcast_factors

True or "block_diagonal" stores the factor in float32 except its diagonal blocks, and "strict_diagonal" except its diagonal.

 required
emulation_mantissa_bits

From 1 to 64, the mantissa bits cuDSS uses to emulate a float64 factorization on GPUs that support it.

 required
batch_layout

"stacked" (default) factors a jax.vmap batch as one block-diagonal matrix, and "uniform" refactors only the matrices that ask.

 required

Raises:

Type Description
ConfigurationError

For a setting outside these values, or a combination cuDSS refuses.

diff_soup.linear.CuDSSOptions

The cuDSS settings as an options= object for analyze with backend=CuDSS().

The same settings are fields of CuDSS. Every field of this object, defaults included, overrides the backend's.

Attributes:

Name Type Description
reorder str

The fill-reducing ordering, "amd" or "metis".
hybrid_memory bool

True keeps part of the factor in host memory when it outgrows the GPU.
ir_nsteps Optional[int]

The iterative refinement steps per solve, where None takes 0 for a float64 factor and 1 otherwise.

Composed backends

These take arguments, so they are passed as objects.

diff_soup.linear.Krylov

Restarted GMRES, preconditioned by the factor of a direct backend.

It solves float64 systems to a tolerance rather than exactly. A refactor always takes the new values, and its predicate decides only whether the preconditioner is factored again. Pass it as linear=ds.linear.Krylov(ds.linear.CuDSS()).

Parameters:

Name Type Description Default
inner LinearSolver

The direct backend whose factor preconditions the iteration, for example CuDSS().

 required
tol Optional[Tolerance]

The default stopping rule, Tolerance(rtol=1e-10, atol=0.0, max_iter=100, restart=20) when None.

 None
lagged_max_iter int

The iteration cap of each linear solve inside the solvers, default 10.

 DEFAULT_LAGGED_MAX_ITER

diff_soup.linear.BlockTridiagonal

A direct solver in JAX for block-tridiagonal matrices, on the CPU and on GPUs.

Every pattern entry must join a block to itself or to a neighbouring block, or analyze raises. It factors float64 matrices only. Its factor needs no SolveContext.

Parameters:

Name Type Description Default
block_size int

The unknowns per block, for example 1 for one field on a one-dimensional grid.

 required

diff_soup.linear.MemberBlock

A direct solver in JAX for chains of sites that couple only through retained unknowns.

A member is a chain of sites whose entries join each site only to itself and its neighbours. Each solve checks its residual row by row, and a solve that misses after its sweeps reports SolveStatus.SINGULAR.

Parameters:

Name Type Description Default
layout MemberBlockLayout

The MemberBlockLayout of the members and the retained unknowns.

 required
residual_tolerance float

The relative residual each row of a solve must reach.

 1e-10
refinement_sweeps int

The most correction sweeps a solve runs to reach it.

 3

diff_soup.linear.MemberBlockLayout

The members and the retained unknowns of a MemberBlock solver.

Every unknown is either a site of exactly one member or retained.

Parameters:

Name Type Description Default
members Any

One (n_sites, block_size) index array per member, with one block_size for all.

 required
n_dofs int

The number of unknowns.

 required
retained Any

The retained unknowns, by default every unknown in no member.

 None
term Optional[LowRankTerm]

A LowRankTerm whose couplings each lie inside one member, or None.

 None
names Any

One name per member, used in error messages.

 None

diff_soup.linear.Woodbury

A direct solver for a sparse matrix plus a low-rank term, A + U W^T.

It factors A with inner and applies the Woodbury identity for the term. Analyze term.topology(base), and pass the values of A followed by the values of U. It needs a SolveContext exactly when inner does.

Parameters:

Name Type Description Default
inner LinearSolver

A direct backend for A, for example Dense() or CuDSS().

 required
term LowRankTerm

The LowRankTerm U W^T.

 required

Example

backend = Woodbury(Dense(), term)
analysis = analyze(term.topology(base), backend=backend)
f = factor(analysis, jnp.concatenate([a_values, u_values]))

diff_soup.linear.LowRankTerm

A low-rank term U W^T added to a sparse matrix, for the Woodbury solver.

The entries of U take their values from the matrix values, after those of the sparse part. The entries of W have fixed weights.

Parameters:

Name Type Description Default
n_dofs int

The number of rows of the matrix.

 required
rank int

The number of couplings, the columns of U and W.

 required
u_rows Any

The row of each entry of U.

 required
u_cols Any

The coupling of each entry of U.

 required
w_rows Any

The row of each entry of W.

 required
w_cols Any

The coupling of each entry of W.

 required
w_values Any

The weight of each entry of W.

 required

topology 

topology(base: SparseTopology) -> SparseTopology

The pattern to analyze, base followed by the entries of U.

Parameters:

Name Type Description Default
base SparseTopology

The SparseTopology of the sparse part.

 required

Returns:

Type Description
SparseTopology

The SparseTopology to pass to analyze with a Woodbury backend.
Abstract base class

The details matter only when writing a backend of your own.

diff_soup.linear.LinearSolver

The base class of the linear backends.

Subclass it only to write a backend of your own. Its class attributes declare what the backend supports, and analyze refuses a request outside them.

backend_name class-attribute 

backend_name: Optional[str] = None

The backend's name, such as "klu" or "dense", or None on the base class.

platforms class-attribute 

platforms: tuple[str, ...] = ()

The JAX platforms the backend runs on, such as ("cpu",), where empty means every platform.

factor_residency class-attribute 

factor_residency: str = 'arrays'

Where a factor is stored, "arrays" in the factor itself or "host" or "device" memory a SolveContext owns.

transpose_is_free class-attribute 

transpose_is_free: bool = False

Whether one factor serves transpose=True solves without another factorization.

vmap_factors class-attribute 

vmap_factors: str = 'unsupported'

How jax.vmap over the values runs, "native" as one batch, "sequential" one matrix at a time, or "unsupported".

supported_precisions class-attribute 

supported_precisions: tuple[str, ...] = ('fp64',)

The factor precisions analyze accepts, from "fp64" and "fp32".

is_exact class-attribute 

is_exact: bool = True

Whether a solve applies the exact inverse rather than iterating to a tolerance.

accepts_operator class-attribute 

accepts_operator: bool = False

Whether the backend works on a matrix-free operator rather than assembled values.