tensormesh.sparse¶

Sparse matrix data type and solver entry points. Built on top of torch-sla; see Sparse Solvers for the design and backend matrix.

SparseMatrix¶

class SparseMatrix(edata: Tensor, row: Tensor, col: Tensor, shape: Tuple[int, int])[source]¶

Bases: SparseTensor

COO sparse matrix with FEM-flavoured helpers.

Subclass of torch_sla.SparseTensor. Inherits @ (spmm), solve, autograd-aware values, .to_dense() and friends from the parent; adds layout hashing, block-COO assembly, scipy interoperability and type-preserving arithmetic.

Parameters:

edata (Tensor) – 1D float tensor of shape [nnz]: the non-zero values.
row (Tensor) – 1D integer tensor of shape [nnz]: row indices.
col (Tensor) – 1D integer tensor of shape [nnz]: column indices.
shape (Tuple[int, int]) – (nrows, ncols) of the dense equivalent.

layout_hash¶

SHA-256 of (row, col) concatenated bytes. Two matrices with the same sparsity pattern share this hash; useful for caching any quantity that depends only on the topology (Condenser permutations, AMG hierarchies, etc.).

Type:: str

Examples

import torch
from tensormesh.sparse import SparseMatrix

edata = torch.tensor([1.0, 2.0, 3.0])
row   = torch.tensor([0, 1, 2])
col   = torch.tensor([1, 2, 0])
A     = SparseMatrix(edata, row, col, shape=(3, 3))

# Inherited from SparseTensor:
y = A @ torch.tensor([1.0, 2.0, 3.0])
x = A.double().solve(torch.ones(3, dtype=torch.float64))

# FEM-specific block-COO assembly: 10 element matrices of size 3x3.
block_data = torch.randn(10, 3, 3)
elem_row   = torch.arange(10)
elem_col   = torch.arange(10)
K = SparseMatrix.from_block_coo(block_data, elem_row, elem_col, (10, 10))

__init__(edata: Tensor, row: Tensor, col: Tensor, shape: Tuple[int, int])[source]¶

property edata: Tensor¶: Alias for values (legacy TensorMesh API).

property row: Tensor¶: Alias for row_indices.

property col: Tensor¶: Alias for col_indices.

property edges: Tensor¶: Stacked (row, col) indices of shape [2, nnz].

property layout_mask: Tensor¶: Dense [nrows, ncols] mask with 1.0 on the sparsity pattern.

property grad: SparseMatrix | None¶

Gradient w.r.t. values, wrapped as a SparseMatrix.

Returns None when no gradient has accumulated on values.

to(*args, **kwargs) → SparseMatrix[source]¶

Move tensor to device and/or convert dtype.

Parameters:

device (str or device, optional) – Target device (e.g., ‘cuda’, ‘cpu’, ‘cuda:0’).
dtype (dtype, optional) – Target data type (e.g., torch.float32, torch.float64).

Returns:

New SparseTensor on the target device/dtype.

Return type:

SparseTensor

Examples

>>> A = SparseTensor(val, row, col, shape)
>>> A_cuda = A.to('cuda')
>>> A_float32 = A.to(dtype=torch.float32)
>>> A_cuda_float32 = A.to('cuda', torch.float32)

cuda(device=None) → SparseMatrix[source]¶

Move tensor to CUDA device.

Parameters:: device (int, optional) – CUDA device index. Default: current device.
Returns:: Tensor on CUDA.
Return type:: SparseTensor

cpu() → SparseMatrix[source]¶

Move tensor to CPU.

Returns:: Tensor on CPU.
Return type:: SparseTensor

float() → SparseMatrix[source]¶: Convert to float32.

double() → SparseMatrix[source]¶: Convert to float64.

half() → SparseMatrix[source]¶: Convert to float16.

detach() → SparseMatrix[source]¶: Detach from computation graph. Preserves subclass type.

has_same_layout(other: str | SparseMatrix) → bool[source]¶

Check whether other shares this matrix’s sparsity pattern.

Parameters:: other (str or SparseMatrix) – Either another SparseMatrix or a layout_hash string to compare against.
Returns:: True iff both have identical (row, col) arrays.
Return type:: bool
Raises:: TypeError – If other is neither a string nor a SparseMatrix.

degree(axis: int = 0) → Tensor[source]¶

Non-zero count per row (axis=0) or per column (axis=1).

Parameters:: axis (int, default 0) – Aggregate over axis; 0 counts nnz per row, 1 per column.
Returns:: 1D int tensor of length shape[axis].
Return type:: Tensor

transpose() → SparseMatrix[source]¶: Return A^T as a SparseMatrix sharing this matrix’s values.

property T: SparseMatrix¶: Shorthand for transpose().

to_scipy_coo() → coo_matrix[source]¶: Detach and convert to a scipy.sparse.coo_matrix on CPU.

to_sparse_coo() → Tensor[source]¶: Convert to a torch.sparse_coo_tensor() (autograd-tracked).

static from_scipy_coo(matrix: coo_matrix, device: str = 'cpu', dtype: dtype = torch.float) → SparseMatrix[source]¶

Wrap a scipy.sparse.coo_matrix as a SparseMatrix.

Parameters:

matrix (coo_matrix) – Source matrix.
device (str, default "cpu") – Target torch device.
dtype (torch.dtype, default torch.float) – Target value dtype.

Return type:

SparseMatrix

static from_sparse_coo(matrix: Tensor) → SparseMatrix[source]¶

Wrap a coalesced torch.sparse_coo_tensor() as a SparseMatrix.

Parameters:: matrix (Tensor) – Sparse COO tensor; will be coalesced in place.
Return type:: SparseMatrix

static from_dense(tensor: Tensor) → SparseMatrix[source]¶

Pull non-zero entries out of a dense 2D tensor into COO.

Parameters:: tensor (Tensor) – 2D tensor; entries exactly equal to zero are dropped.
Return type:: SparseMatrix

static from_block_coo(edata: Tensor, row: Tensor, col: Tensor, shape: Tuple[int, int]) → SparseMatrix[source]¶

Expand block-COO storage into a flat SparseMatrix.

Each block is a small dense matrix attached to a (row, col) pair in a coarser graph. This is the layout produced by FEM element assembly when the element data is vector-valued (e.g. linear elasticity), with one [dim, dim] block per (node_i, node_j) pair.

Blocks are assumed square; the function uses edata.shape[1] as the block size and ignores shape[2].

Parameters:

edata (Tensor) – Block data of shape [n_elements, block_size, block_size].
row (Tensor) – Block indices of shape [n_elements]; entries refer to the coarse graph.
col (Tensor) – Block indices of shape [n_elements]; entries refer to the coarse graph.
shape (Tuple[int, int]) – (block_rows, block_cols) of the coarse graph.

Returns:

Flat COO matrix of shape (shape[0] * block_size, shape[1] * block_size).

Return type:

SparseMatrix

static random(m: int, n: int, density: float = 0.1, device: str = 'cpu', dtype: dtype = torch.float) → SparseMatrix[source]¶

Random m x n sparse matrix with the requested density.

Parameters:

m (int) – Dense shape.
n (int) – Dense shape.
density (float, default 0.1) – Fraction of entries that are non-zero. Drawn via scipy.sparse.random().
device (str, default "cpu")
dtype (torch.dtype, default torch.float)

Return type:

SparseMatrix

static random_layout(m: int, n: int, density: float = 0.1, device: str = 'cpu') → Tuple[Tensor, Tensor, Tuple[int, int]][source]¶

Random (row, col, shape) triple with no value tensor attached.

Useful when several matrices need to share the same sparsity pattern; combine with random_from_layout().

Returns:: (row, col, (m, n)).
Return type:: Tuple[Tensor, Tensor, Tuple[int, int]]

static random_from_layout(layout: Tuple[Tensor, Tensor, Tuple[int, int]], device: str = 'cpu', dtype: dtype = torch.float) → SparseMatrix[source]¶

Attach random uniform [0, 1) values to an existing layout.

Parameters:

layout (Tuple[Tensor, Tensor, Tuple[int, int]]) – (row, col, shape) triple, e.g. from random_layout().
device (str, default "cpu")
dtype (torch.dtype, default torch.float)

Return type:

SparseMatrix

static eye(n: int, value: float = 1., device: str = 'cpu', dtype: dtype = torch.float) → SparseMatrix[source]¶

n x n sparse identity, optionally scaled by value.

Parameters:

n (int) – Dimension.
value (float, default 1.0) – Diagonal value; value * I_n is stored.
device (str, default "cpu")
dtype (torch.dtype, default torch.float)

Return type:

SparseMatrix

static full(m: int, n: int, value: float = 1., device: str = 'cpu', dtype: dtype = torch.float) → SparseMatrix[source]¶

Constant m x n matrix stored densely-as-sparse.

Mostly a building block for combine_matrix() block layouts; value == 0 is short-circuited to an empty COO tensor.

Parameters:

m (int) – Shape.
n (int) – Shape.
value (float, default 1.0) – Constant entry.
device (str, default "cpu")
dtype (torch.dtype, default torch.float)

Return type:

SparseMatrix

static combine_vector(matrices: List[SparseMatrix], axis: int = 0) → SparseMatrix[source]¶

Stack a list of sparse matrices along axis.

Equivalent of torch.cat() for SparseMatrix. Dense torch.Tensor entries are auto-converted via from_dense().

Parameters:

matrices (List[SparseMatrix or Tensor]) – Items must agree on shape[1 - axis].
axis (int, default 0) – 0 stacks vertically (along rows), 1 horizontally.

Return type:

SparseMatrix

static combine_matrix(matrices: List[List[SparseMatrix]]) → SparseMatrix[source]¶

Assemble a 2D block layout of sparse matrices into one matrix.

Each entry of matrices[i][j] may be:

None — treated as a zero block of inferred size;
int or float — expanded via full() to a constant block of the inferred size;
torch.Tensor — auto-converted via from_dense();
SparseMatrix — used as-is.

Block sizes are inferred from the first non-scalar entry in each row / column; a row or column with only None/scalar entries will have inferred size zero.

Parameters:: matrices (List[List[SparseMatrix or Tensor or int or float or None]]) – Rectangular nested list, [n_rows][n_cols].
Returns:: Assembled matrix of shape (sum(row_sizes), sum(col_sizes)).
Return type:: SparseMatrix

static combine(matrices) → SparseMatrix[source]¶

Dispatch to combine_matrix() or combine_vector().

If the first element is a list or tuple, the input is treated as a 2D block layout (calls combine_matrix()); otherwise it is a 1D stack along axis 0 (calls combine_vector()).

class SparseTensor[source]¶: COO sparse tensor with autograd support — the base class that SparseMatrix extends. Re-exported from torch-sla; see the torch-sla docs for the full API.

Solvers¶

spmm(edata, row, col, shape, B, backend=None)[source]¶

Compute A @ B for sparse A and dense matrix (or vector) B.

When B is 1D this transparently delegates to spmv.

Parameters:

edata (Tensor) – 1D tensor of shape [nnz]: non-zero values of A.
row (Tensor) – 1D int tensor of shape [nnz]: row indices of A.
col (Tensor) – 1D int tensor of shape [nnz]: column indices of A.
shape (Tuple[int, int]) – Shape (M, N) of the dense A.
B (Tensor) – Dense operand of shape [N, K] (matrix) or [N] (vector).
backend (str or None, default None) – See spmv for the backend / device matrix.

Returns:

[M, K] or [M] matching the rank of B.

Return type:

Tensor

Legacy entry points (deprecated)¶

Deprecated since version 0.x: The two functions below pre-date the torch-sla integration and are scheduled for removal. The canonical solver path is the methods on SparseMatrix itself (inherited from torch_sla.SparseTensor):

Linear solves → SparseMatrix.solve (i.e. torch_sla.SparseTensor.solve) — see Sparse Solvers.
Nonlinear solves → SparseMatrix.nonlinear_solve (Newton / Picard / Anderson with line search and autograd Jacobian).

Both will be retired once the remaining in-tree call sites have migrated. New code should not use them.

spsolve(edata, row, col, shape, b, backend='auto', method='cg', preconditioner='jacobi', tol=1e-5, max_iter=10000, x0=None, is_spd=True, verbose=False)[source]¶

Solve the sparse linear system A x = b (legacy entry point).

Deprecated since version 0.x: This free function pre-dates the torch-sla integration and is scheduled for removal. The canonical path is to wrap the data in a SparseMatrix and call its solve method (inherited from torch_sla.SparseTensor), which auto-detects symmetry / positive-definiteness and routes to torch_sla.spsolve directly. See Sparse Solvers.

Low-level entry point: takes raw COO arrays instead of a SparseMatrix object. With torch-sla installed, dispatches to a differentiable sparse-linear-algebra backend; without it, falls back to a curated mini-stack of SciPy / SuperLU / CuPy / PETSc wrappers.

Parameters:

edata (Tensor) – 1D tensor of shape [nnz]: non-zero values of A.
row (Tensor) – 1D int tensor of shape [nnz]: row indices of A.
col (Tensor) – 1D int tensor of shape [nnz]: column indices of A.
shape (Tuple[int, int]) – Dense shape (m, n) of A.
b (Tensor) – Right-hand side. Shape [n] for a single RHS, or [n, n_batch] for batched RHS (auto-routed through SuperLU).
backend (str, default "auto") – Torch-sla path: "auto" (CPU → "scipy", CUDA → "pytorch"), "scipy", "pytorch", "eigen", "cudss", "cupy". Fallback path (no torch-sla): "auto", "petsc", "cupy" — others are accepted but the method/preconditioner hints below are ignored.
method (str, default "cg") – Iterative algorithm — "cg", "bicgstab", "minres", "gmres", "lgmres" — or one of the direct factorizations "lu", "umfpack", "cholesky", "ldlt". See the installed torch_sla.spsolve signature for the canonical list. Honoured only on the torch-sla path. On the fallback path, each wrapper uses a fixed algorithm.
preconditioner (str, default "jacobi") – "jacobi", "ilu", or "none". Same caveat as method — torch-sla path only.
tol (float, default 1e-5) – Convergence tolerance (iterative methods).
max_iter (int, default 10000) – Iteration budget (iterative methods).
x0 (Tensor, optional) – Initial guess. Currently consumed only by some fallback wrappers and ignored by torch-sla.
is_spd (bool, default True) – Hint to the torch-sla path that A is symmetric positive definite. Picks CG as the default method; set False for indefinite / non-symmetric A and combine with method="bicgstab" or "gmres".
verbose (bool, default False) – Print which backend/method was picked.

Returns:

Solution x, same shape and dtype as b.

Return type:

Tensor

Notes

Both paths are autograd-aware: gradients of x flow back into edata and b via an adjoint sparse solve. On the torch-sla path this is built in to the library; on the fallback path each wrapper supplies its own torch.autograd.Function backward.

Examples

>>> from tensormesh.sparse import spsolve
>>> x = spsolve(edata, row, col, (n, n), b)                   # auto
>>> x = spsolve(edata, row, col, (n, n), b, method="lu")      # direct
>>> x = spsolve(edata, row, col, (n, n), b, backend="cudss")  # GPU direct
>>> x = spsolve(edata, row, col, (n, n), b,
...             is_spd=False, method="bicgstab")              # non-SPD

nonlinear_solve(f: Callable[[...], Tensor], j: Callable[[...], SparseMatrix], u0: Tensor, params: Tuple[Tensor, ...], max_iter: int = 100, tol: float = 1e-6, verbose: bool = False) → Tensor[source]¶

Solve F(u, params) = 0 for u (legacy in-tree implementation).

Deprecated since version 0.x: This in-tree Newton-Raphson driver is scheduled for removal. torch-sla ships a richer torch_sla.SparseTensor.nonlinear_solve (also accessible as SparseMatrix.nonlinear_solve) with Newton / Picard / Anderson modes, Armijo line search, and an autograd-based Jacobian — so the user no longer has to supply j. Migrate to A.nonlinear_solve(residual, u0, *params); see Sparse Solvers.

Drives Newton-Raphson on the forward pass and the implicit-function theorem on the backward pass: gradients w.r.t. params cost roughly one extra linear solve regardless of how many Newton iterations were taken.

Parameters:

f (Callable) – F(u, *params) -> torch.Tensor. Must support autograd in params for gradients to flow.
j (Callable) – J(u, *params) -> SparseMatrix: the Jacobian dF/du.
u0 (Tensor) – Initial guess.
params (Tuple[Tensor, ...]) – Optimizable parameters forwarded to f / j.
max_iter (int, default 100) – Newton iteration budget.
tol (float, default 1e-6) – Convergence tolerance on ||F(u)||.
verbose (bool, default False) – Print residual norm each Newton step.

Returns:

Converged u; gradients route back into params via the adjoint solve inside NonLinearSolveFunction.

Return type:

Tensor

Backend availability flags¶

Module-level booleans, set at import time, indicating which optional backends were detected. torch-sla itself is a hard dependency — its absence raises at import time rather than flipping a flag — so it is not listed here.

is_petsc_available¶

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

is_cupy_available¶

bool(x) -> bool