splineax 0.2.0__tar.gz

splineax-0.2.0/PKG-INFO

@@ -0,0 +1,84 @@
+Metadata-Version: 2.3
+Name: splineax
+Version: 0.2.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.post1
+Requires-Dist: asdex>=0.5.1
+Requires-Python: >=3.11
+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`, `BCSRLinearOperator`, `SparseJacobianLinearOperator`.
+- **Solvers**: `Spsolve` (any backend), `KLU` (CPU-only, SuiteSparse, factorization reuse),
+  and `AutoSparseLinearSolver` (picks one based on the platform).
+- A `SparseLinearSolver` protocol for separating factorization from solving.
+
+## Installation
+
+```bash
+pip install git+https://github.com/nardi/splineax.git@v0.1.1
+```
+
+## 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
+
+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 = splineax.BCOOLinearOperator(matrix)
+vectors = [jnp.ones(n), jnp.arange(n) % 2]
+solver = splineax.AutoSparseLinearSolver()
+
+# Calculate factorization once...
+with solver.factorize(operator) as factorized_state:
+    # ...and reuse for multiple solves.
+    solution = lx.linear_solve(
+        operator, vectors[0], solver=solver, state=factorized_state
+    )
+    assert jnp.allclose(matrix @ solution.value, vectors[0], atol=1e-4)
+
+    solution = lx.linear_solve(
+        operator, vectors[1], solver=solver, state=factorized_state
+    )
+    assert jnp.allclose(matrix @ solution.value, vectors[1], atol=1e-4)
+```
+
+## Documentation
+
+Build the docs locally with `uv run mkdocs serve`, or view the [user guide and API reference here](https://nardi.github.io/splineax).

splineax-0.2.0/README.md

@@ -0,0 +1,69 @@
+# 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`, `BCSRLinearOperator`, `SparseJacobianLinearOperator`.
+- **Solvers**: `Spsolve` (any backend), `KLU` (CPU-only, SuiteSparse, factorization reuse),
+  and `AutoSparseLinearSolver` (picks one based on the platform).
+- A `SparseLinearSolver` protocol for separating factorization from solving.
+
+## Installation
+
+```bash
+pip install git+https://github.com/nardi/splineax.git@v0.1.1
+```
+
+## 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
+
+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 = splineax.BCOOLinearOperator(matrix)
+vectors = [jnp.ones(n), jnp.arange(n) % 2]
+solver = splineax.AutoSparseLinearSolver()
+
+# Calculate factorization once...
+with solver.factorize(operator) as factorized_state:
+    # ...and reuse for multiple solves.
+    solution = lx.linear_solve(
+        operator, vectors[0], solver=solver, state=factorized_state
+    )
+    assert jnp.allclose(matrix @ solution.value, vectors[0], atol=1e-4)
+
+    solution = lx.linear_solve(
+        operator, vectors[1], solver=solver, state=factorized_state
+    )
+    assert jnp.allclose(matrix @ solution.value, vectors[1], atol=1e-4)
+```
+
+## Documentation
+
+Build the docs locally with `uv run mkdocs serve`, or view the [user guide and API reference here](https://nardi.github.io/splineax).

splineax-0.2.0/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "splineax"
+version = "0.2.0"
+description = "Sparse linear operators and solvers for Lineax"
+readme = "README.md"
+authors = [{ name = "Nardi Lam", email = "mail@nardilam.nl" }]
+requires-python = ">=3.11"
+dependencies = [
+    "lineax>=0.1.1",
+    "jax>=0.4",
+    "equinox>=0.11",
+    "jaxtyping>=0.2",
+    "splineax-klujax>=0.5.0.post1",
+    "asdex>=0.5.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.7,<0.12.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "mkdocs-material>=9.7.6",
+    "mkdocstrings[python]>=0.27",
+    "pytest>=9.1.0",
+    "pytest-markdown-docs>=0.9.2",
+    "ruff>=0.15.17",
+    "ty>=0.0.59",
+]
+
+[tool.pytest.ini_options]
+# Execute the code examples in the markdown docs (and README) as tests via
+# pytest-markdown-docs. The superfences syntax lets dependent code blocks share state
+# with a `{.python continuation}` fence, which also renders under Material for MkDocs.
+addopts = "--markdown-docs --markdown-docs-syntax superfences"
+testpaths = ["tests", "docs", "README.md"]
+
+[tool.ruff.lint]
+ignore = ["F722"]
+select = ["E4", "E7", "E9", "F", "I"]

splineax-0.2.0/src/splineax/__init__.py

@@ -0,0 +1,30 @@
+from .operators import (
+    BCOOLinearOperator as BCOOLinearOperator,
+)
+from .operators import (
+    BCSRLinearOperator as BCSRLinearOperator,
+)
+from .operators import (
+    JacobianColoring as JacobianColoring,
+)
+from .operators import (
+    SparseJacobianLinearOperator as SparseJacobianLinearOperator,
+)
+from .operators import (
+    SparseJacobianLinearOperatorColoring as SparseJacobianLinearOperatorColoring,
+)
+from .solvers import (
+    KLU as KLU,
+)
+from .solvers import (
+    AbstractSparseLinearSolver as AbstractSparseLinearSolver,
+)
+from .solvers import (
+    AutoSparseLinearSolver as AutoSparseLinearSolver,
+)
+from .solvers import (
+    SparseLinearSolver as SparseLinearSolver,
+)
+from .solvers import (
+    Spsolve as Spsolve,
+)

splineax-0.2.0/src/splineax/operators/__init__.py

@@ -0,0 +1,7 @@
+from ._bcoo import BCOOLinearOperator as BCOOLinearOperator
+from ._bcsr import BCSRLinearOperator as BCSRLinearOperator
+from ._jacobian import JacobianColoring as JacobianColoring
+from ._jacobian import SparseJacobianLinearOperator as SparseJacobianLinearOperator
+from ._jacobian import (
+    SparseJacobianLinearOperatorColoring as SparseJacobianLinearOperatorColoring,
+)

splineax-0.2.0/src/splineax/operators/_bcoo.py

@@ -0,0 +1,78 @@
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jax.experimental.sparse import BCOO
+from jaxtyping import Array, Inexact
+from lineax import AbstractLinearOperator, is_symmetric
+from lineax._tags import transpose_tags
+
+from ._operations import (
+    register_sparse_operator,
+    sparse_as_matrix,
+    sparse_in_structure,
+    sparse_mv,
+    sparse_out_structure,
+)
+
+
+class BCOOLinearOperator(AbstractLinearOperator):
+    """Wraps a `jax.experimental.sparse.BCOO` array into a linear operator.
+
+    If the matrix has shape `(a, b)` then matrix-vector multiplication (`self.mv`) is
+    defined in the usual way: as accepting a vector of shape `(b,)` and returning a
+    vector of shape `(a,)`.
+    """
+
+    matrix: Inexact[BCOO, "a b"]
+    tags: frozenset[object] = eqx.field(static=True)
+
+    def __init__(
+        self, matrix: Inexact[BCOO, "a b"], tags: object | frozenset[object] = ()
+    ):
+        """**Arguments:**
+
+        - `matrix`: a two-dimensional `BCOO` array. For an array with shape `(a, b)`
+            then this operator can perform matrix-vector products on a vector of shape
+            `(b,)` to return a vector of shape `(a,)`.
+        - `tags`: any tags indicating whether this matrix has any particular properties,
+            like symmetry or positive-definite-ness. Note that these properties are
+            unchecked and you may get incorrect values elsewhere if these tags are
+            wrong.
+        """
+        if matrix.ndim != 2:
+            raise ValueError(
+                "`BCOOLinearOperator(matrix=...)` should be 2-dimensional."
+            )
+        if not jnp.issubdtype(matrix.dtype, jnp.inexact):
+            matrix = BCOO(
+                (matrix.data.astype(float), matrix.indices), shape=matrix.shape
+            )
+        self.matrix = matrix
+        self.tags = tags if isinstance(tags, frozenset) else frozenset([tags])
+
+    def mv(self, vector: Inexact[Array, " b"]) -> Inexact[Array, " a"]:
+        return sparse_mv(self.matrix, vector)
+
+    def as_matrix(self) -> Inexact[Array, "a b"]:
+        return sparse_as_matrix(self.matrix)
+
+    def transpose(self) -> "BCOOLinearOperator":
+        if is_symmetric(self):
+            return self
+        matrix_T: BCOO = self.matrix.T
+        return BCOOLinearOperator(matrix_T, transpose_tags(self.tags))
+
+    def in_structure(self) -> jax.ShapeDtypeStruct:
+        return sparse_in_structure(self)
+
+    def out_structure(self) -> jax.ShapeDtypeStruct:
+        return sparse_out_structure(self)
+
+    def _conj(self) -> "BCOOLinearOperator":
+        matrix = BCOO(
+            (self.matrix.data.conj(), self.matrix.indices), shape=self.matrix.shape
+        )
+        return BCOOLinearOperator(matrix, self.tags)
+
+
+register_sparse_operator(BCOOLinearOperator)

splineax-0.2.0/src/splineax/operators/_bcsr.py

@@ -0,0 +1,82 @@
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jax.experimental.sparse import BCOO, BCSR
+from jaxtyping import Array, Inexact
+from lineax import AbstractLinearOperator, is_symmetric
+from lineax._tags import transpose_tags
+
+from ._operations import (
+    register_sparse_operator,
+    sparse_as_matrix,
+    sparse_in_structure,
+    sparse_mv,
+    sparse_out_structure,
+)
+
+
+class BCSRLinearOperator(AbstractLinearOperator):
+    """Wraps a `jax.experimental.sparse.BCSR` array into a linear operator.
+
+    If the matrix has shape `(a, b)` then matrix-vector multiplication (`self.mv`) is
+    defined in the usual way: as accepting a vector of shape `(b,)` and returning a
+    vector of shape `(a,)`.
+    """
+
+    matrix: Inexact[BCSR, "a b"]
+    tags: frozenset[object] = eqx.field(static=True)
+
+    def __init__(
+        self, matrix: Inexact[BCSR, "a b"], tags: object | frozenset[object] = ()
+    ):
+        """**Arguments:**
+
+        - `matrix`: a two-dimensional `BCSR` array. For an array with shape `(a, b)`
+            then this operator can perform matrix-vector products on a vector of shape
+            `(b,)` to return a vector of shape `(a,)`.
+        - `tags`: any tags indicating whether this matrix has any particular properties,
+            like symmetry or positive-definite-ness. Note that these properties are
+            unchecked and you may get incorrect values elsewhere if these tags are
+            wrong.
+        """
+        if matrix.ndim != 2:
+            raise ValueError(
+                "`BCSRLinearOperator(matrix=...)` should be 2-dimensional."
+            )
+        if not jnp.issubdtype(matrix.dtype, jnp.inexact):
+            matrix = BCSR(
+                (matrix.data.astype(jnp.float32), matrix.indices, matrix.indptr),
+                shape=matrix.shape,
+            )
+        self.matrix = matrix
+        self.tags = tags if isinstance(tags, frozenset) else frozenset([tags])
+
+    def mv(self, vector: Inexact[Array, " b"]) -> Inexact[Array, " a"]:
+        return sparse_mv(self.matrix, vector)
+
+    def as_matrix(self) -> Inexact[Array, "a b"]:
+        return sparse_as_matrix(self.matrix)
+
+    def transpose(self) -> "BCSRLinearOperator":
+        if is_symmetric(self):
+            return self
+        # `BCSR.transpose` is not implemented in JAX; round-trip through `BCOO`.
+        matrix_T_bcoo: BCOO = self.matrix.to_bcoo().T
+        matrix_T = BCSR.from_bcoo(matrix_T_bcoo)
+        return BCSRLinearOperator(matrix_T, transpose_tags(self.tags))
+
+    def in_structure(self) -> jax.ShapeDtypeStruct:
+        return sparse_in_structure(self)
+
+    def out_structure(self) -> jax.ShapeDtypeStruct:
+        return sparse_out_structure(self)
+
+    def _conj(self) -> "BCSRLinearOperator":
+        matrix = BCSR(
+            (self.matrix.data.conj(), self.matrix.indices, self.matrix.indptr),
+            shape=self.matrix.shape,
+        )
+        return BCSRLinearOperator(matrix, self.tags)
+
+
+register_sparse_operator(BCSRLinearOperator)