The factor contract¶

These functions factor and solve a sparse matrix in code of your own, with any backend (linear solvers).

import jax.numpy as jnp
import numpy as np
from diff_soup.linear import Dense, SparseTopology, analyze, factor, solve

rows, cols = np.nonzero(np.ones((3, 3)))
analysis = analyze(SparseTopology(rows=rows, cols=cols, n_dofs=3), backend=Dense())
f = factor(analysis, jnp.asarray([4.0, -1.0, 0.0, -1.0, 4.0, -1.0, 0.0, -1.0, 3.0]))
x, f = solve(f, jnp.asarray([1.0, 2.0, -1.0]))
x_t, f = solve(f, jnp.asarray([1.0, 2.0, -1.0]), transpose=True)

Functions¶

Each function on a factor returns the factor to use next.

diff_soup.linear.analyze ¶

analyze(topology: SparseTopology, *, backend: LinearSolver, options: Any = None, layout: Optional[Layout] = None, precision: str = 'fp64', symmetric: bool = False, spd: bool = False) -> Analysis

Analyze a sparsity pattern for one backend, before any values are factored.

Equal arguments return the same Analysis, so the work runs once per pattern. A request the backend cannot serve raises here, before anything is traced.

Parameters:

Name	Type	Description	Default
`topology`	`SparseTopology`	The `SparseTopology` of the matrix.	required
`backend`	`LinearSolver`	The backend object, for example `Dense()` or `KLU(ordering="colamd")`.	required
`options`	`Any`	Backend settings as a dict or a `CuDSSOptions`, overriding the fields of `backend` key by key.	`None`
`layout`	`Optional[Layout]`	The device placement, or `None` for one device.	`None`
`precision`	`str`	The precision of the factors, `"fp64"` or `"fp32"`.	`'fp64'`
`symmetric`	`bool`	Declare the matrix symmetric, which is never detected from the values.	`False`
`spd`	`bool`	Declare the matrix symmetric positive definite, which implies `symmetric`.	`False`

Returns:

Type	Description
`Analysis`	The `Analysis` to pass to `factor`.

Raises:

Type	Description
`ConfigurationError`	For a setting, precision or layout the backend does not support, or `symmetric` with Dirichlet rows.

diff_soup.linear.factor ¶

factor(analysis: Analysis, values: ArrayLike, *, ctx: Optional[SolveContext] = None, orientation: str = 'A', tol: Optional[Tolerance] = None, after: Any = ()) -> Factor

Factor a matrix from its values.

Parameters:

Name	Type	Description	Default
`analysis`	`Analysis`	The `Analysis` of the pattern.	required
`values`	`ArrayLike`	One value per pattern entry, shape `(nnz,)`, with `jax.vmap` for a batch of matrices.	required
`ctx`	`Optional[SolveContext]`	The `SolveContext` that owns the factor, required by KLU, SuperLU and cuDSS.	`None`
`orientation`	`str`	`"A"` factors the matrix and `"AT"` its transpose.	`'A'`
`tol`	`Optional[Tolerance]`	A `Tolerance` for a backend of your own, ignored by the built-in backends.	`None`
`after`	`Any`	A factor, or a sequence of factors, whose operations this one must follow.	`()`

Returns:

Type	Description
`Factor`	The `Factor`.

Raises:

Type	Description
`FactorContractError`	For values of the wrong shape, or no `ctx` where the backend needs one.

diff_soup.linear.empty_factor ¶

empty_factor(analysis: Analysis, *, ctx: Optional[SolveContext] = None, orientation: str = 'A', tol: Optional[Tolerance] = None, dtype: Any = None, storage: Optional[str] = None, after: Any = ()) -> Factor

A factor that holds no factorization yet, the start of a loop that refactors it.

Its solves return nan with SolveStatus.EMPTY until a refactor with a true predicate fills it.

Parameters:

Name	Type	Description	Default
`analysis`	`Analysis`	The `Analysis` of the pattern.	required
`ctx`	`Optional[SolveContext]`	The `SolveContext` that will own the factor, as for `factor`.	`None`
`orientation`	`str`	`"A"` for the matrix and `"AT"` for its transpose.	`'A'`
`tol`	`Optional[Tolerance]`	As for `factor`.	`None`
`dtype`	`Any`	The dtype of the values, by default float64 or float32 to match the analysis.	`None`
`storage`	`Optional[str]`	A name of 1 to 64 characters, and the factors of one name in one context share one storage.	`None`
`after`	`Any`	A factor, or a sequence of factors, whose operations this one must follow.	`()`

Returns:

Type	Description
`Factor`	The empty `Factor`.

Example

def step(f, v):
    f = refactor(f, v, True)
    x, f = solve(f, b)
    return f, x

f, xs = jax.lax.scan(step, empty_factor(analysis), values)

diff_soup.linear.refactor ¶

refactor(f: Factor, values: ArrayLike, predicate: ArrayLike, *, adopt: bool = False) -> Factor

Factor new values where predicate is true, and keep the factor where it is false.

Under jax.vmap each matrix of the batch takes its own predicate.

Parameters:

Name	Type	Description	Default
`f`	`Factor`	The factor to update.	required
`values`	`ArrayLike`	The new values, in the dtype of `f.values`.	required
`predicate`	`ArrayLike`	A boolean scalar, `True` to refactor.	required
`adopt`	`bool`	`True` refactors in the storage of `f`, and `False` in storage of its own.	`False`

Returns:

Type	Description
`Factor`	The updated `Factor`.

Raises:

Type	Description
`FactorContractError`	For values of another shape or dtype, or a predicate that is not a scalar.

diff_soup.linear.solve ¶

solve(f: Factor, b: ArrayLike, *, transpose: bool = False, tol: Optional[Tolerance] = None, weights: Optional[ArrayLike] = None, x0: Optional[ArrayLike] = None) -> tuple[Array, Factor]

Solve M x = b with a factor of M, or M^T x = b with transpose=True.

A factor that is singular, non-finite or empty solves to nan. An iterative solve that stops short of tol returns its last iterate, and solve_checked reports the status.

Parameters:

Name	Type	Description	Default
`f`	`Factor`	The factor of `M`.	required
`b`	`ArrayLike`	The right-hand side, shape `(n,)`, or `(nrhs, n)` for several.	required
`transpose`	`bool`	A Python bool, `True` to solve with the transpose.	`False`
`tol`	`Optional[Tolerance]`	The `Tolerance` of an iterative solve, by default the backend's own.	`None`
`weights`	`Optional[ArrayLike]`	The error weights of an iterative solve's root-mean-square norm.	`None`
`x0`	`Optional[ArrayLike]`	The initial guess of an iterative solve, zeros by default.	`None`

Returns:

Type	Description
`tuple[Array, Factor]`	The solution `x` and the factor to use next.

diff_soup.linear.solve_checked ¶

solve_checked(f: Factor, b: ArrayLike, *, transpose: bool = False, tol: Optional[Tolerance] = None, weights: Optional[ArrayLike] = None, x0: Optional[ArrayLike] = None) -> tuple[Array, Factor, SolveReport]

Solve as solve does, and also return a SolveReport with the status and the residual.

Parameters:

Name	Type	Description	Default
`f`	`Factor`	The factor of `M`.	required
`b`	`ArrayLike`	The right-hand side, shape `(n,)`, or `(nrhs, n)` for several.	required
`transpose`	`bool`	A Python bool, `True` to solve with the transpose.	`False`
`tol`	`Optional[Tolerance]`	The `Tolerance` of an iterative solve, by default the backend's own.	`None`
`weights`	`Optional[ArrayLike]`	The error weights, which also make the reported residual a weighted root-mean-square norm.	`None`
`x0`	`Optional[ArrayLike]`	The initial guess of an iterative solve, zeros by default.	`None`

Returns:

Type	Description
`tuple[Array, Factor, SolveReport]`	The solution `x`, the factor to use next, and the `SolveReport`.

diff_soup.linear.invalidate ¶

invalidate(factor: Factor) -> Factor

Mark a factor's stored factorization out of date, so its next solve refactors from its values.

Use it when the storage may hold another factorization than the factor's values describe. A factor held in arrays, as on Dense, has nothing to mark.

Parameters:

Name	Type	Description	Default
`factor`	`Factor`	The factor.	required

Returns:

Type	Description
`Factor`	The factor to use next.

diff_soup.linear.refine ¶

refine(f: Factor, b: ArrayLike, *, values: ArrayLike, residual_target: float = DEFAULT_RESIDUAL_TARGET, max_steps: int = DEFAULT_MAX_STEPS, transpose: bool = False) -> tuple[Array, Factor, RefineReport]

Solve with a factor, then refine the solution against the values of the matrix.

Use it with a float32 factor of a float64 matrix, or a factor of a nearby matrix. It stops at residual_target, after max_steps, or when a step fails to halve the residual.

Parameters:

Name	Type	Description	Default
`f`	`Factor`	A factor from a direct backend.	required
`b`	`ArrayLike`	The right-hand side.	required
`values`	`ArrayLike`	The values of the matrix to solve, one per pattern entry.	required
`residual_target`	`float`	The relative residual `\|\|b - M x\|\| / \|\|b\|\|` to reach, default `1e-8`.	`DEFAULT_RESIDUAL_TARGET`
`max_steps`	`int`	The most correction steps, default `5`.	`DEFAULT_MAX_STEPS`
`transpose`	`bool`	Solve with the transpose instead.	`False`

Returns:

Type	Description
`tuple[Array, Factor, RefineReport]`	The solution, the factor to use next, and a `RefineReport`.

Raises:

Type	Description
`FactorContractError`	For a factor from an iterative backend.

diff_soup.linear.after ¶

after(f: Factor, *others: Factor) -> Factor

Order the next operations on a factor after those on other factors.

Use it for factors that share one named storage, so that work on one waits for the others.

Parameters:

Name	Type	Description	Default
`f`	`Factor`	The factor whose operations must wait.	required
`*others`	`Factor`	The factors to wait for.	`()`

Returns:

Type	Description
`Factor`	`f`, ordered after `others`.

Types¶

diff_soup.linear.SparseTopology ¶

The sparsity pattern of a square matrix, given by the coordinates of its entries.

The values passed to factor follow the order of rows and cols, one per entry, and duplicate coordinates add. A Dirichlet row becomes an identity row, so the pattern must hold its diagonal entry (i, i).

Parameters:

Name	Type	Description	Default
`rows`	`Any`	The row index of each entry.	required
`cols`	`Any`	The column index of each entry.	required
`n_dofs`	`int`	The number of rows and columns.	required
`bc_indices`	`Any`	The Dirichlet rows, or `None` for none.	`None`

rows `property` ¶

rows: ndarray

The row index of each entry, as a read-only int32 NumPy array.

cols `property` ¶

cols: ndarray

The column index of each entry, as a read-only int32 NumPy array.

bc_indices `property` ¶

bc_indices: Optional[ndarray]

The Dirichlet rows as a read-only int32 NumPy array, or None.

full `classmethod` ¶

full(n: int) -> 'SparseTopology'

The pattern of a dense n by n matrix, with its entries in row-major order.

Parameters:

Name	Type	Description	Default
`n`	`int`	The number of rows and columns.	required

Returns:

Type	Description
`'SparseTopology'`	The `SparseTopology`.

diff_soup.linear.Analysis ¶

The analysis of a sparsity pattern for one backend, as returned by analyze.

Pass it to factor or empty_factor. It is immutable and hashable, and unpickling it analyzes the pattern again.

Attributes:

Name	Type	Description
`topology`		The `SparseTopology`.
`backend`		The backend object.
`options`		The settings, as sorted `(key, value)` pairs.
`layout`		The device placement.
`precision`		`"fp64"` or `"fp32"`.
`symmetric`		Whether the matrix was declared symmetric.
`spd`		Whether the matrix was declared symmetric positive definite.
`n`	`int`	The number of rows.
`nnz`	`int`	The number of pattern entries.

diff_soup.linear.Factor ¶

A factorization of a matrix, as returned by factor, empty_factor and refactor.

It is a JAX pytree, so it passes through jax.jit, jax.vmap and loop carries. Pass on the factor each operation returns, which keeps the operations in order.

Attributes:

Name	Type	Description
`analysis`	`Analysis`	The `Analysis` the factor was built from.
`orientation`	`str`	`"A"` for the matrix and `"AT"` for its transpose.
`precision`	`str`	`"fp64"` or `"fp32"`.
`exact`	`bool`	Whether solves apply the exact inverse.
`values`	`Array`	The values the factor was built from.
`generation`	`Array`	`0` after `factor` and `-1` after `empty_factor`, plus one for every refactor that took new values.
`storage`	`Optional[str]`	The storage name given to `empty_factor`, or `None`.

diff_soup.linear.SolveContext ¶

The owner of the factors that KLU, SuperLU and cuDSS store outside JAX arrays.

Pass it as ctx= to factor or empty_factor. Releasing it, or leaving its with block, frees every factor it owns, and their solves raise from then on.

Parameters:

Name	Type	Description	Default
`name`	`Optional[str]`	A label shown in its `repr`.	`None`

Example

with SolveContext() as ctx:
    f = factor(analysis, values, ctx=ctx)
    x, f = solve(f, b)

release ¶

release() -> None

Free every factor this context owns.

A context that holds stored factors first waits for all running JAX work in the process. A second call does nothing.

Raises:

Type	Description
`FactorContractError`	For the default context of the solvers, which is never released.

resident_bytes ¶

resident_bytes() -> int

The bytes of stored factors this context owns.

Returns:

Type	Description
`int`	The total, `0` for factors held in arrays and after release.

diff_soup.linear.SolveReport ¶

The status and the residual of one solve_checked call.

Its fields are arrays, so a report passes through jax.jit and jax.vmap.

Attributes:

Name	Type	Description
`status`	`Array`	The `SolveStatus` code, as an int32.
`residual`	`Array`	The largest `\|b - M x\|`, or the weighted root-mean-square residual when the solve had `weights`.
`iterations`	`Array`	The iteration count, `0` for a direct solve.
`applications`	`Array`	The count of operator or factor applications, `1` for a direct solve.
`ok`	`Array`	`True` where `status` is `SUCCESS`.
`converged`	`Array`	The same as `ok`.

diff_soup.linear.SolveStatus ¶

The status code of a factorization or a solve, one int32 per matrix.

Attributes:

Name	Type	Description
`SUCCESS`		The solve succeeded.
`SINGULAR`		The factorization met a zero pivot, and its solves return `nan`.
`NONFINITE`		The values or the factor hold a non-finite number, and solves return `nan`.
`MAX_ITERATIONS`		An iterative solve reached its iteration cap.
`STAGNATION`		An iterative solve stopped making progress.
`BREAKDOWN`		An iterative solve broke down.
`PRECONDITIONER_FAILURE`		The preconditioner failed to factor or to apply.
`EMPTY`		The factor holds no factorization yet, and its solves return `nan`.

diff_soup.linear.Tolerance ¶

The stopping rule of an iterative solve.

A solve of M x = b stops when ||w * (b - M x)||_rms <= atol + rtol * ||w * b||_rms, with w the solve's weights, or ones without them.

Attributes:

Name	Type	Description
`rtol`	`float`	The relative tolerance, at least 0.
`atol`	`float`	The absolute tolerance, at least 0.
`max_iter`	`int`	The iteration cap, at least 0.
`restart`	`int`	The iterations between GMRES restarts, at least 1, default `30`.

diff_soup.linear.RefineReport ¶

The outcome of one refine call.

Attributes:

Name	Type	Description
`steps`	`Array`	The number of corrections applied.
`stalled`	`Array`	`True` when the final residual is above the target or not finite.
`residual`	`Array`	The final relative residual `\|\|b - M x\|\| / \|\|b\|\|`, the worst over a block of right-hand sides.

diff_soup.linear.FactorContractError ¶

Raised for a factor built or used in a way its backend does not allow.

It is a ConfigurationError. Examples are a released SolveContext, values of the wrong shape, and a KLU factor in a function compiled for a GPU.

Matrix-free operators¶

The iterative solves work on a matrix given by its action.

diff_soup.linear.LinearOperator ¶

A matrix given by its action, for the iterative solves.

Parameters:

Name	Type	Description	Default
`mv`	`Callable[[Array], Array]`	The function `x -> M x`.	required
`rmv`	`Optional[Callable[[Array], Array]]`	The function `x -> M^T x`, or `None` when no transpose is needed.	required
`shape`	`tuple[int, int]`	The shape, `(rows, columns)`.	required
`dtype`	`Any`	The dtype of the entries.	required
`tags`	`Iterable[str]`	Structure tags from `"square"`, `"symmetric"`, `"spd"` and `"nonsymmetric"`.	`('square', 'nonsymmetric')`

apply ¶

apply(x: ArrayLike, *, transpose: bool = False) -> Array

Apply the operator, M x, or M^T x with transpose=True.

Parameters:

Name	Type	Description	Default
`x`	`ArrayLike`	The vector to apply it to.	required
`transpose`	`bool`	Apply the transpose instead.	`False`

Returns:

Type	Description
`Array`	The product.

Raises:

Type	Description
`FactorContractError`	For `transpose=True` when `rmv` is `None`.

diff_soup.linear.operator_from_values ¶

operator_from_values(analysis: Analysis, values: ArrayLike, *, transpose: bool = False) -> LinearOperator

The operator of an assembled matrix, from its analysis and values.

Its structure tags follow the symmetric and spd declarations of the analysis.

Parameters:

Name	Type	Description	Default
`analysis`	`Analysis`	The `Analysis` of the pattern.	required
`values`	`ArrayLike`	One value per pattern entry, shape `(nnz,)`, or `(batch, nnz)`.	required
`transpose`	`bool`	Build the operator of the transpose instead.	`False`

Returns:

Type	Description
`LinearOperator`	The `LinearOperator`.

diff_soup.linear.operator_from_jvp ¶

operator_from_jvp(residual: Callable[[Array], Array], primal: ArrayLike, *, shape: Optional[tuple[int, int]] = None, dtype: Any = None, tags: Optional[Iterable[str]] = None) -> LinearOperator

The Jacobian of residual at primal, as a matrix-free operator.

It linearizes residual once, and applying the operator never evaluates residual again. The Jacobian is never assembled.

Parameters:

Name	Type	Description	Default
`residual`	`Callable[[Array], Array]`	A function from a vector `(n,)` to a vector `(m,)`.	required
`primal`	`ArrayLike`	The point `(n,)` where the Jacobian is taken.	required
`shape`	`Optional[tuple[int, int]]`	The expected shape `(m, n)`, checked against the residual, or `None`.	`None`
`dtype`	`Any`	The dtype of the entries, by default the residual's output dtype.	`None`
`tags`	`Optional[Iterable[str]]`	Structure tags, by default `("square", "nonsymmetric")` for a square Jacobian.	`None`

Returns:

Type	Description
`LinearOperator`	The `LinearOperator`.

Raises:

Type	Description
`FactorContractError`	If `primal` or the output is not a vector, or `shape` disagrees.

diff_soup.linear.gmres ¶

gmres(operator: LinearOperator, b: ArrayLike, *, precondition: Any = None, tol: Tolerance, weights: Optional[ArrayLike] = None, x0: Optional[ArrayLike] = None, transpose: bool = False) -> tuple[Array, SolveReport]

Solve M x = b with restarted GMRES on a matrix-free operator.

Parameters:

Name	Type	Description	Default
`operator`	`LinearOperator`	The `LinearOperator` `M`.	required
`b`	`ArrayLike`	The right-hand side `(n,)`, or `(nrhs, n)` for independent systems.	required
`precondition`	`Any`	A `Factor`, a function `v -> P^-1 v`, or `None` for no preconditioner.	`None`
`tol`	`Tolerance`	The `Tolerance` that stops the iteration.	required
`weights`	`Optional[ArrayLike]`	The error weights of the stopping rule's root-mean-square norm.	`None`
`x0`	`Optional[ArrayLike]`	The initial guess, zeros by default.	`None`
`transpose`	`bool`	Solve `M^T x = b`, which needs the operator's `rmv`.	`False`

Returns:

Type	Description
`tuple[Array, SolveReport]`	The solution and a `SolveReport`, the worst over the right-hand sides of a block.

The factor contract¶

Functions¶

diff_soup.linear.analyze ¶

diff_soup.linear.factor ¶

diff_soup.linear.empty_factor ¶

diff_soup.linear.refactor ¶

diff_soup.linear.solve ¶

diff_soup.linear.solve_checked ¶

diff_soup.linear.invalidate ¶

diff_soup.linear.refine ¶

diff_soup.linear.after ¶

Types¶

diff_soup.linear.SparseTopology ¶

rows property ¶

cols property ¶

bc_indices property ¶

full classmethod ¶

diff_soup.linear.Analysis ¶

diff_soup.linear.Factor ¶

diff_soup.linear.SolveContext ¶

release ¶

resident_bytes ¶

diff_soup.linear.SolveReport ¶

diff_soup.linear.SolveStatus ¶

diff_soup.linear.Tolerance ¶

diff_soup.linear.RefineReport ¶

diff_soup.linear.FactorContractError ¶

Matrix-free operators¶

diff_soup.linear.LinearOperator ¶

apply ¶

diff_soup.linear.operator_from_values ¶

diff_soup.linear.operator_from_jvp ¶

diff_soup.linear.gmres ¶

rows `property` ¶

cols `property` ¶

bc_indices `property` ¶

full `classmethod` ¶