COO

class brainevent.COO(data, row=None, col=None, ptr=None, *, shape, rows_sorted=False, cols_sorted=False, backend=None, buffers=None)

Coordinate Format (COO) sparse matrix.

This class represents a sparse matrix in coordinate format, where non-zero elements are stored as triplets (row, column, value).

The class supports arithmetic with dense arrays and scalars, and sparse-dense matrix multiplication via @. Sparse-sparse operations are intentionally limited and may raise NotImplementedError.

data

Array of the non-zero values in the matrix.

Type:

jax.Array, Quantity

row

Array of row indices for each non-zero element.

Type:

jax.Array

col

Array of column indices for each non-zero element.

Type:

jax.Array

shape

Shape of the matrix (rows, columns).

Type:

tuple[int, int]

nse

Number of stored elements (property).

Type:

int

dtype

Data type of the matrix elements (property).

Type:

dtype

info

Additional information about the matrix structure (property).

Type:

COOInfo

_bufs

Tuple of (data, row, col, ptr) arrays (property).

Type:

tuple

rows_sorted

Whether row indices are sorted.

Type:

bool

cols_sorted

Whether column indices are sorted within each row.

Type:

bool

ptr

Row (or column) pointer array computed when indices are sorted.

Type:

jax.Array or None

Note

This class is registered as a PyTree node for JAX, allowing it to be used with JAX transformations and compiled functions.

property T

Get the transpose of the COO matrix.

Returns:

The transposed COO matrix.

Return type:

COO

apply(fn)

Apply a function to the data and return a new sparse matrix with the same structure.

Unlike with_data(), which requires the new data to have the same shape, dtype, and unit, apply allows transformations that change dtype or unit.

Parameters:

fn (callable) – A function to apply to self.data.

Returns:

A new COO matrix with fn(self.data) and the same structure.

Return type:

COO

apply2(other, fn, *, reverse=False)

Apply a binary function while preserving sparse structure semantics.

Parameters:

• other (Any) – Right-hand operand for normal operations, or left-hand operand when reverse=True.

• fn (callable) – Binary function from operator or a compatible callable.

• reverse (bool) – If False, compute fn(self, other) semantics using _binary_op. If True, compute fn(other, self) semantics using _binary_rop. Defaults to False.

Returns:

Result of the operation.

Return type:

COO or Data

property dtype

Data type of the stored non-zero values.

Returns:

Element data type, e.g. jnp.float32.

Return type:

numpy.dtype

classmethod fromdense(mat, *, nse=None, index_dtype=<class 'numpy.int32'>, backend=None)

Create a COO (Coordinate Format) sparse matrix from a dense matrix.

This method converts a dense matrix to a sparse COO representation.

Parameters:

• mat (Array | ndarray | Quantity | Number) – The dense matrix to be converted to COO format.

• nse (int | None) – The number of non-zero elements in the matrix. If None, it will be calculated from the input matrix. Default is None.

• index_dtype (str | type[Any] | dtype | SupportsDType) – The data type to be used for the row and column indices. Default is np.int32.

Returns:

A new COO sparse matrix object representing the input dense matrix.

Return type:

COO

property info

Structural metadata for the COO matrix.

Returns:

A named tuple containing shape, rows_sorted, and cols_sorted.

Return type:

COOInfo

property nse

Number of stored elements (non-zeros).

Returns:

Total count of explicitly stored entries.

Return type:

int

sort_indices()

Return a copy of the COO matrix with sorted indices.

The matrix is sorted by row indices and column indices per row. If self.rows_sorted is True, this returns self without a copy.

Return type:

COO

tocsr()

Convert the COO matrix to CSR (Compressed Sparse Row) format.

Returns:

A CSR matrix containing the same data as the original COO matrix.

Return type:

CSR

todense()

Convert the COO matrix to a dense array.

Returns:

A dense representation of the COO matrix.

Return type:

Array | ndarray | Quantity | Number

transpose(axes=None)

Transpose the COO matrix.

Parameters:

axes (Tuple[int, ...] | None) – The axes to transpose over. Currently not implemented and will raise a NotImplementedError if provided.

Returns:

The transposed COO matrix.

Return type:

COO

Raises:

NotImplementedError – If axes argument is provided.

tree_flatten()

Flatten the COO matrix for JAX transformations.

This method is used by JAX to serialize the COO matrix object.

Returns:

A tuple containing: - A tuple with the matrix data. - A dictionary with auxiliary data (shape, sorting information, row and column indices).

Return type:

Tuple[Tuple[Array | Quantity], dict[str, Any]]

classmethod tree_unflatten(aux_data, children)

Reconstruct a COO matrix from flattened data.

This class method is used by JAX to deserialize the COO matrix object.

Parameters:

• aux_data (dict) – Auxiliary data containing shape, sorting information, and row and column indices.

• children (tuple) – A tuple containing the matrix data.

Returns:

The reconstructed COO matrix.

Return type:

COO

Raises:

ValueError – If the auxiliary data doesn’t contain the expected keys.

update_on_post(post_events, pre_trace, w_min=None, w_max=None, inplace=False, backend=None)

Update synaptic weights based on post-synaptic events (post-spike rule).

Applies a spike-timing-dependent plasticity (STDP) update to the stored weights. For each non-zero element (i, j) whose post-synaptic neuron i is active in post_events, the weight is updated according to the pre-synaptic trace value.

Parameters:

• post_events (EventRepresentation) – Post-synaptic spike events. Currently only BinaryArray is supported.

• pre_trace (Array | ndarray | Quantity | Number) – Pre-synaptic eligibility trace, shape (num_pre,).

• w_min (Array | ndarray | Quantity | Number | None) – Minimum weight bound.

• w_max (Array | ndarray | Quantity | Number | None) – Maximum weight bound.

• inplace (bool) – If True, mutate self.data in place and return None.

• backend (str | None) – Compute backend.

Returns:

Updated weight data, or None when inplace is True.

Return type:

jax.Array or Quantity or None

Raises:

NotImplementedError – If post_events is not a BinaryArray, or if the matrix uses unsupported pointer layout.

See also

update_on_pre

Pre-synaptic plasticity update.

update_coo_on_binary_post

Underlying COO post-spike kernel.

update_on_pre(pre_events, post_trace, w_min=None, w_max=None, inplace=False, backend=None)

Update synaptic weights based on pre-synaptic events (pre-spike rule).

Applies a spike-timing-dependent plasticity (STDP) update to the stored weights. For each non-zero element (i, j) in the COO matrix whose pre-synaptic neuron j is active in pre_events, the weight is updated according to the post-synaptic trace value.

Parameters:

• pre_events (EventRepresentation) – Pre-synaptic spike events. Currently only BinaryArray is supported.

• post_trace (Array | ndarray | Quantity | Number) – Post-synaptic eligibility trace, shape (num_post,).

• w_min (Array | ndarray | Quantity | Number | None) – Minimum weight bound. None means no lower clipping.

• w_max (Array | ndarray | Quantity | Number | None) – Maximum weight bound. None means no upper clipping.

• inplace (bool) – If True, mutate self.data in place and return None. If False (default), return the updated weight array.

• backend (str | None) – Compute backend (e.g. 'numba', 'pallas').

Returns:

Updated weight data, or None when inplace is True.

Return type:

jax.Array or Quantity or None

Raises:

NotImplementedError – If pre_events is not a BinaryArray, or if the matrix uses cols_sorted pointers (not yet supported).

See also

update_on_post

Post-synaptic plasticity update.

update_coo_on_binary_pre

Underlying COO pre-spike kernel.

with_data(data)

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

This method returns a new COO matrix with the same row and column indices as the current matrix, but with new data values.

Parameters:

data (Array | ndarray | Quantity | Number) – The new data to be used in the COO matrix. Must have the same shape, dtype, and unit as the current matrix’s data.

Returns:

A new COO matrix with the provided data and the same structure as the current matrix.

Return type:

COO

Raises:

AssertionError – If the shape, dtype, or unit of the new data doesn’t match the current matrix’s data.