splineax 0.2.0__py3-none-any.whl

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/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/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/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)