CSR

CSR#

class brainevent.CSR(data, indices=None, indptr=None, *, shape, backend=None, buffers=None)#

Event-driven and Unit-aware Compressed Sparse Row (CSR) matrix.

This class represents a sparse matrix in CSR format, which is efficient for row-wise operations and matrix-vector multiplications. It is compatible with JAX’s tree utilities and supports unit-aware computations.

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

data#

Array of the non-zero values in the matrix.

Type:: Data

indices#

Array of column indices for the non-zero values.

Type:: jax.Array

indptr#

Array of row pointers indicating where each row starts in the data and indices arrays.

Type:: jax.Array

shape#

The shape of the matrix as (rows, columns).

Type:: tuple[int, int]

nse#

Number of stored elements (non-zero entries).

Type:: int

dtype#

Data type of the matrix values.

Type:: dtype

Notes

In CSR format a matrix of shape (m, n) is stored as three arrays:

indptr of length m + 1 – the i-th row occupies entries indptr[i] to indptr[i+1] in the data and indices arrays.
indices – column indices of the stored elements.
data – the corresponding non-zero values.

The @ operator dispatches to optimised kernels depending on the right-hand operand type:

BinaryArray – event-driven binary CSR MV/MM.
Dense jax.Array / brainunit.Quantity – standard float CSR MV/MM with automatic dtype promotion.

Examples

import jax.numpy as jnp
import brainevent

data    = jnp.array([1.0, 2.0, 3.0])
indices = jnp.array([0, 2, 1])
indptr  = jnp.array([0, 1, 2, 3])
csr     = brainevent.CSR((data, indices, indptr), shape=(3, 3))

# Sparse-dense matrix-vector product
x = jnp.ones(3)
y = csr @ x

See also

CSC: Compressed Sparse Column format.
brainevent.COO: Coordinate sparse format.

apply(fn)[source]#

Apply a unary function to the stored data values.

Creates a new CSR matrix with fn(self.data) while preserving the sparsity structure (indices, indptr, shape, and cached diagonal positions).

Parameters:: fn (callable) – A function that accepts a single array argument and returns an array of the same shape. The dtype and unit may differ from the input.
Returns:: A new CSR matrix with transformed data.
Return type:: CSR

Examples

squared = csr.apply(lambda x: x ** 2)

classmethod fromdense(mat, *, nse=None, index_dtype=<class 'jax.numpy.int32'>, backend=None)[source]#

Create a CSR matrix from a dense matrix.

This method converts a dense matrix to a Compressed Sparse Row (CSR) format.

Parameters:

mat (array_like) – The dense matrix to be converted to CSR format.
nse (int | None) – The number of non-zero elements in the matrix. If None, it will be calculated from the input matrix.
index_dtype (dtype, optional) – The data type to be used for index arrays (default is jnp.int32).

Returns:

A new CSR matrix object created from the input dense matrix.

Return type:

CSR

Examples

import jax.numpy as jnp
import brainevent

dense = jnp.array([[1.0, 0.0], [0.0, 2.0]])
csr = brainevent.CSR.fromdense(dense)

solve(b, tol=1e-06, reorder=1)[source]#

Solve the linear system A x = b where A is this CSR matrix.

Uses a sparse direct solver via the underlying csr_solve routine.

Parameters:

b (Array | Quantity) – Right-hand side vector. Its first dimension must equal self.shape[0].
tol (float, optional) – Tolerance for singularity detection. Defaults to 1e-6.
reorder (int, optional) – Fill-reducing reordering scheme: 0 for no reordering, 1 for symrcm, 2 for symamd, 3 for csrmetisnd. Defaults to 1.

Returns:

Solution vector x satisfying A x = b.

Return type:

Array | Quantity

Raises:

AssertionError – If b.shape[0] != self.shape[0].

Examples

x = csr.solve(b)

tocoo()[source]#

Convert the CSR matrix to COO (Coordinate) format.

This method transforms the Compressed Sparse Row (CSR) matrix into a COO matrix, which stores sparse data as a collection of (row, column, value) triplets.

Returns:: A COO matrix containing the same data as the original CSR matrix.
Return type:: COO

See also

brainevent.COO: The Coordinate sparse matrix class.

Examples

csr_matrix = CSR((data, indices, indptr), shape=(3, 4))
coo_matrix = csr_matrix.tocoo()

todense()[source]#

Convert the CSR matrix to a dense matrix.

This method transforms the compressed sparse row (CSR) representation into a full dense matrix.

Returns:: A dense matrix of shape self.shape containing all the values (including zeros) of the sparse matrix.
Return type:: Array | Quantity

Examples

dense = csr.todense()

transpose(axes=None)[source]#

Transpose the CSR matrix.

This method returns the transpose of the CSR matrix as a CSC matrix. Because the transpose of a CSR matrix is a CSC matrix with the same underlying arrays, this operation is essentially free (no data is copied or rearranged).

Parameters:: axes (None) – This parameter is not used and must be None. Included for compatibility with numpy’s transpose function signature.
Returns:: The transpose of the CSR matrix as a CSC (Compressed Sparse Column) matrix.
Return type:: CSC
Raises:: AssertionError – If axes is not None, as this implementation doesn’t support custom axis ordering.

Examples

csc = csr.transpose()
# or equivalently:
csc = csr.T

with_data(data)[source]#

Create a new CSR matrix with updated data while keeping the same structure.

This method creates a new CSR matrix instance with the provided data, maintaining the original indices, indptr, and shape.

Parameters:: data (Array | ndarray | Quantity | Number) – The new data array to replace the existing data in the CSR matrix. It must have the same shape, dtype, and unit as the original data.
Returns:: A new CSR matrix instance with updated data and the same structure as the original.
Return type:: CSR
Raises:: AssertionError – If the shape, dtype, or unit of the new data doesn’t match the original data.

Examples

new_data = jnp.array([10.0, 20.0, 30.0])
new_csr = csr.with_data(new_data)

yw_to_w(y_dim_arr, w_dim_arr)[source]#

Compute a sparse transformation from y-w space to w space.

Performs a specialised sparse matrix-vector product optimised for event-driven neural simulations, accumulating contributions from the target (post-synaptic) dimension y_dim_arr weighted by the per-synapse values w_dim_arr according to the connectivity defined by this CSR matrix.

Parameters:

y_dim_arr (Array | ndarray | Quantity) – Values in the target (post-synaptic) dimension.
w_dim_arr (Array | ndarray | Quantity) – Per-synapse weight values.

Returns:

Accumulated result, preserving physical units when present.

Return type:

Array | Quantity

See also

yw_to_w_transposed: The transposed (adjoint) variant.

Notes

Internally calls csrmv_yw2y with transpose=False.

yw_to_w_transposed(y_dim_arr, w_dim_arr)[source]#

Compute the transposed sparse transformation from y-w space to w space.

This is the adjoint of yw_to_w(), useful for back-propagation or adjoint computations in event-driven neural simulations.

Parameters:

y_dim_arr (Array | ndarray | Quantity) – Values in the target (post-synaptic) dimension.
w_dim_arr (Array | ndarray | Quantity) – Per-synapse weight values.

Returns:

Accumulated result of the transposed operation, preserving physical units when present.

Return type:

Array | Quantity

See also

yw_to_w: The forward (non-transposed) variant.

Notes

Internally calls csrmv_yw2y with transpose=True.

CSR

Contents

CSR#