COO

class brainunit.sparse.COO(args, *, shape, rows_sorted=False, cols_sorted=False)

Unit-aware Coordinate (COO) sparse matrix.

Stores a 2-D sparse matrix in COO (coordinate) format with optional physical-unit support via Quantity.

Parameters:

• args (Tuple[Array | saiunit.Quantity, Array, Array]) – data contains the non-zero values (jax.Array or Quantity), row contains the row indices, and col contains the column indices.

• shape (tuple[int, ...]) – The (nrows, ncols) shape of the matrix.

• rows_sorted (bool) – Whether the row indices are sorted. Default is False.

• cols_sorted (bool) – Whether the column indices are sorted. Default is False.

data

Non-zero values of shape (nse,).

Type:

jax.Array or Quantity

row

Row indices of shape (nse,).

Type:

jax.Array

col

Column indices of shape (nse,).

Type:

jax.Array

shape

Shape of the matrix (nrows, ncols).

Type:

tuple of int

nse

Number of stored elements.

Type:

int

dtype

Data type of the stored values.

Type:

dtype

See also

CSR

Unit-aware Compressed Sparse Row matrix.

CSC

Unit-aware Compressed Sparse Column matrix.

coo_fromdense

Create a COO matrix from a dense array.

coo_todense

Convert a COO matrix to a dense array.

Notes

This class has minimal compatibility with JAX transforms such as grad and jit, and offers limited functionality compared to jax.experimental.sparse.BCOO. Additionally, there are known failures when nse is larger than the true number of non-zeros in the represented matrix.

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.sparse as susparse
>>> dense = jnp.array([[1., 0., 2.], [0., 0., 3.]])
>>> coo = susparse.COO.fromdense(dense)
>>> coo.shape
(2, 3)
>>> coo.todense()
Array([[1., 0., 2.],
       [0., 0., 3.]], dtype=float32)

todense()

Convert this COO matrix to a dense array.

Returns:

Dense 2-D array equivalent to this sparse matrix.

Return type:

Array

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.sparse as susparse
>>> dense = jnp.array([[0., 3.], [4., 0.]])
>>> coo = susparse.COO.fromdense(dense)
>>> coo.todense()
Array([[0., 3.],
       [4., 0.]], dtype=float32)

with_data(data)

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

Parameters:

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

Returns:

A new COO matrix sharing the same row and col indices but holding the provided data.

Return type:

COO

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.sparse as susparse
>>> dense = jnp.array([[1., 0.], [0., 2.]])
>>> coo = susparse.COO.fromdense(dense)
>>> new_coo = coo.with_data(coo.data * 5)
>>> new_coo.todense()
Array([[ 5., 0.],
       [ 0., 10.]], dtype=float32)