Metadata-Version: 2.3
Name: splineax
Version: 0.4.0
Summary: Sparse linear operators and solvers for Lineax
Author: Nardi Lam
Author-email: Nardi Lam <mail@nardilam.nl>
Requires-Dist: lineax>=0.1.1
Requires-Dist: jax>=0.4
Requires-Dist: equinox>=0.11
Requires-Dist: jaxtyping>=0.2
Requires-Dist: splineax-klujax>=0.5.0.post7
Requires-Dist: asdex>=0.5.1
Requires-Dist: entangle-jax>=0.1.1
Requires-Dist: spineax>=0.0.5 ; python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cudss'
Requires-Dist: pardiso-mkl-jax>=0.3.0 ; extra == 'pardiso'
Requires-Python: >=3.11
Provides-Extra: cudss
Provides-Extra: pardiso
Description-Content-Type: text/markdown

# splineax

Sparse linear operators and direct solvers for
[Lineax](https://github.com/patrick-kidger/lineax).

`splineax` lets you keep a linear system in its native sparse storage
(`jax.experimental.sparse.BCOO` / `BCSR`) and solve it with a sparse *direct* solver that
plugs straight into `lineax.linear_solve`. It also interfaces with [asdex](https://github.com/adrhill/asdex) for calculating sparse Jacobians and using them as operators.

- **Operators**: `BCOOLinearOperator` and `BCSRLinearOperator` to hold sparse matrices,
  and tags for `lineax.JacobianLinearOperator` and `lineax.FunctionLinearOperator` that
  allow solvers to materialize them as sparse matrices via
  [asdex](https://github.com/adrhill/asdex).
- **Stateful solver protocols**: `StatefulSolver` and `SparseLinearSolver` for writing
  solver-agnostic code that reuses factorizations over many solves and operators.
- **Solver library bindings**: `Spsolve` (any backend), `KLU` (CPU-only, SuiteSparse KLU),
  `Pardiso` (CPU-only, Intel oneMKL Pardiso, installed as extra), and `CuDSS` (CUDA
  GPU-only, NVIDIA cuDSS, installed as extra). `KLU`, `Pardiso`, and `CuDSS` reuse their
  factorization across solves.
- **Higher-level solvers**: `AutoSparseLinearSolver`, which picks an appropriate solver
  based on platform and settings, and `HybridDirectIterative`, which combines a direct
  solver with an iterative solver that uses the factorization as its preconditioner. For
  example, `IterativeRefinement` applies a Richardson iteration to the direct solver
  result in order to recover a high-precision solution using a limited-precision solver.
- **Lineax code interop**: `stateful_solve_transform` rewrites a function that calls
  `lineax.linear_solve` so its solves thread a solver state and reuse a factorization.

## Installation

```bash
pip install splineax
```

## Example

Solve a 10000 x 10000 system. As a dense matrix it would need 10^8 entries, but kept
sparse it has only ~3 x 10^4 nonzeros, and the solver never materialises the dense form.

```python
import jax.numpy as jnp
import lineax as lx
import numpy as np
from jax.experimental.sparse import BCOO

import splineax as splx

n = 10000
np.random.seed(0)

# A large, randomly sparse matrix with a heavy diagonal (so it is invertible).
diagonal_indices = np.stack([np.arange(n), np.arange(n)], axis=1)
off_diagonal_indices = np.unique(np.random.randint(0, n, size=(2 * n, 2)), axis=0)
indices = jnp.concatenate([diagonal_indices, off_diagonal_indices])
values = jnp.concatenate(
    [
        np.full(n, float(n)),
        np.random.uniform(low=-1, high=1, size=off_diagonal_indices.shape[0]),
    ]
)
matrix = BCOO((values, indices), shape=(n, n)).sum_duplicates()

operator = splx.BCOOLinearOperator(matrix)
vectors = [jnp.ones(n), jnp.arange(n) % 2]
solver = splx.AutoSparseLinearSolver()

# Solve once, then thread the returned state back in to reuse the factorization.
solution, state = splx.linear_solve(operator, vectors[0], solver)
assert jnp.allclose(matrix @ solution.value, vectors[0], atol=1e-4)

solution, state = splx.linear_solve(operator, vectors[1], solver, state=state)
assert jnp.allclose(matrix @ solution.value, vectors[1], atol=1e-4)

# Free the factorization when you are done with it.
state.release()
```

## Testing on a GPU

`CuDSS` wraps NVIDIA's cuDSS library, which is CUDA-only, so its tests skip on ordinary
(CPU) CI. The easiest way to run them for real is
[`notebooks/colab_gpu_tests.ipynb`](notebooks/colab_gpu_tests.ipynb):

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/nardi/splineax/blob/main/notebooks/colab_gpu_tests.ipynb)

Open it, pick a GPU runtime, and choose *Run all*. It clones a branch, installs the
locked environment with the `cudss` extra, and runs the test suite. On any other machine
with a CUDA 13 GPU, the same run is `uv run --python 3.12 --extra cudss pytest tests`. On
a GPU, the tests marked `cpu_only` skip, since they solve through `KLU` or `Pardiso`.

## Documentation

Build the docs locally with `uv run mkdocs serve`, or view the [user guide and API reference here](https://nardi.github.io/splineax).
