SparseMatrix

class saiunit.sparse.SparseMatrix(args, *, shape)

Base class for sparse matrices in saiunit.

This base class defines the interface that all sparse matrix implementations in the saiunit package should follow. Concrete subclasses must implement the abstract methods defined here.

data

The non-zero values in the sparse matrix.

Type:

jax.Array

Notes

This class provides NotImplementedError for most operations, requiring concrete subclasses to implement them according to their specific sparse format.

Examples

SparseMatrix is not instantiated directly. Use a concrete subclass such as CSR, CSC, or COO.

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.sparse as susparse
>>> dense = jnp.array([[1., 0.], [0., 2.]])
>>> csr = susparse.CSR.fromdense(dense)
>>> isinstance(csr, susparse.SparseMatrix)
True

sum(axis=None)

Sum of the elements of the sparse matrix.

Parameters:

axis (int | Sequence[int]) – Axis or axes along which the sum is computed. The default (None) computes the sum of the flattened array. Currently only None is supported.

Returns:

The sum of all elements in the sparse matrix.

Return type:

jax.Array or Quantity

Raises:

NotImplementedError – If axis is not None.

with_data(data)

Create a new sparse matrix with the same sparsity structure but different data.

Parameters:

data (Array | ndarray | Number | saiunit.Quantity) – The new non-zero values. Must have the same shape, dtype, and unit as the current self.data.

Returns:

A new sparse matrix of the same type with the provided data.

Return type:

SparseMatrix

Raises:

NotImplementedError – If called on the abstract base class directly.

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.sparse as susparse
>>> dense = jnp.array([[1., 0.], [0., 2.]])
>>> csr = susparse.CSR.fromdense(dense)
>>> new_csr = csr.with_data(csr.data * 3)
>>> new_csr.todense()
Array([[3., 0.],
       [0., 6.]], dtype=float32)

yw_to_w(y_dim_arr, w_dim_arr)

The protocol method to convert the product of the sparse matrix and a vector to the sparse matrix data.

This protocol method is primarily used in brainscale.

Parameters:

• y_dim_arr (Array | ndarray | saiunit.Quantity) – The first vector.

• w_dim_arr (Array | ndarray | saiunit.Quantity) – The second vector.

Return type:

Array | saiunit.Quantity

Returns:

The outer product of the two vectors.