CSR and CSC Sparse Matrices#
This chapter treats connectivity as Data: which edges are stored, how their weights are represented, and which orientation matches an operation.
import brainevent
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
Why Use Sparse Connectivity Data?#
A dense matrix stores every possible edge. CSR and CSC store only explicit nonzero weights plus integer indices. This can reduce storage when connectivity is sparse, but the actual speed and memory benefit depends on shape, sparsity, dtype, operation, backend, and compilation state.
COO Input, CSR Storage, and CSC Storage#
Coordinate (COO) data lists (row, column, value) triplets. CSR groups entries by row through indptr; CSC groups them by column. CSR naturally supports row-oriented access and events @ weights; CSC is useful when column-oriented access is primary.
dense = jnp.array([[0.0, 0.5, 0.0, -0.2],
[0.3, 0.0, 0.0, 0.0],
[0.0, 0.1, 0.4, 0.0]])
csr = brainevent.CSR.fromdense(dense)
csc = brainevent.CSC.fromdense(dense)
print(csr.indptr, csr.indices, csr.data)
print(csc.indptr, csc.indices, csc.data)
[0 2 3 5] [1 3 0 1 2] [ 0.5 -0.2 0.3 0.1 0.4]
[0 1 3 4 5] [1 0 2 2 0] [ 0.3 0.5 0.1 0.4 -0.2]
Constructing CSR and CSC Data#
fromdense is convenient for small examples. In data pipelines, construct from sparse source data when possible so a large dense intermediate is never materialized. Converting back with todense() is appropriate for validation and visualization at small scale.
assert jnp.array_equal(csr.todense(), dense)
assert jnp.array_equal(csc.todense(), dense)
print(csr.shape, csc.shape, csr.nse)
(3, 4) (3, 4) 5
Combining Sparse Data with Binary Events#
The sparse object describes connectivity; BinaryArray describes active presynaptic events. Keeping those responsibilities separate makes the same connectivity reusable with binary or dense activity.
events = brainevent.BinaryArray(jnp.array([True, False, True]))
sparse_output = events @ csr
dense_output = events.value.astype(dense.dtype) @ dense
print(sparse_output)
assert jnp.allclose(sparse_output, dense_output)
---------------------------------------------------------------------------
ModuleNotFoundError Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:569, in XLACustomKernel._register_fallback_lowering.<locals>.fallback_kernel_fn(*args, **kwargs)
568 try:
--> 569 kernel = entry.kernel_generator(**kwargs)
570 except Exception as exc:
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_csr/binary.py:393, in _csrmv_numba_kernel(weight_info, vector_info, transpose, **kwargs)
387 def _csrmv_numba_kernel(
388 weight_info: jax.ShapeDtypeStruct,
389 vector_info: jax.ShapeDtypeStruct,
390 transpose: bool,
391 **kwargs
392 ):
--> 393 import numba
394 if weight_info.size == 1:
ModuleNotFoundError: No module named 'numba'
The above exception was the direct cause of the following exception:
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_csr/binary.py:250, in binary_csrmv()
249 v, unitv = u.split_mantissa_unit(v)
--> 250 res = binary_csrmv_p_call(
251 data,
252 indices,
253 indptr,
254 v,
255 shape=shape,
256 transpose=transpose,
257 backend=backend,
258 workspace=workspace,
259 )[0]
260 return u.maybe_decimal(res * (unitd * unitv))
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_csr/binary.py:962, in binary_csrmv_p_call()
961 # Call the binary_csrmv_p custom operation to perform the matrix-vector multiplication.
--> 962 return binary_csrmv_p(
963 weights,
964 indices,
965 indptr,
966 vector,
967 task_begin,
968 task_end,
969 status,
970 # Initialize a zero vector with the output shape and data type.
971 outs=(out_info, task_begin_info, task_end_info, status_info),
972 shape=shape,
973 transpose=transpose,
974 backend=backend,
975 # Provide shape and data type information for indices.
976 indices_info=jax.ShapeDtypeStruct(indices.shape, indices.dtype),
977 # Provide shape and data type information for indptr.
978 indptr_info=jax.ShapeDtypeStruct(indptr.shape, indptr.dtype),
979 # Provide shape and data type information for weights.
980 weight_info=jax.ShapeDtypeStruct(weights.shape, weights.dtype),
981 # Provide shape and data type information for v.
982 vector_info=jax.ShapeDtypeStruct(vector.shape, vector.dtype),
983 task_begin_info=task_begin_info,
984 task_end_info=task_end_info,
985 status_info=status_info,
986 task_capacity=task_capacity,
987 )
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:338, in __call__()
337 flat_outs, tree_def = abstract_arguments(outs)
--> 338 r = self.primitive.bind(*ins, **kwargs, outs=tuple(flat_outs))
339 if len(r) != len(flat_outs):
JaxStackTraceBeforeTransformation: brainevent.KernelCompilationError: Backend 'numba' failed to construct a kernel for primitive 'binary_csrmv' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba', 'jax_raw']. Switch with backend='<name>' on this call, or binary_csrmv.set_default('cpu', '<name>') / brainevent.set_backend('cpu', '<name>') to change the default.
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
KernelCompilationError Traceback (most recent call last)
Cell In[4], line 2
1 events = brainevent.BinaryArray(jnp.array([True, False, True]))
----> 2 sparse_output = events @ csr
3 dense_output = events.value.astype(dense.dtype) @ dense
4 print(sparse_output)
5 assert jnp.allclose(sparse_output, dense_output)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_event/binary.py:214, in BinaryArray.__matmul__(self, oc)
212 return binary_densemm(oc, self.value.T, transpose=True).T
213 else:
--> 214 return oc.__rmatmul__(self)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_csr/main.py:1739, in CSR.__rmatmul__(self, other)
1737 if other.ndim == 1:
1738 matrix, workspace = _ensure_binary_workspace_and_get(self, "csr", self.indptr)
-> 1739 return binary_csrmv(matrix.data, matrix.indices, matrix.indptr, other,
1740 shape=matrix.shape, transpose=True, backend=matrix.backend, workspace=workspace)
1741 elif other.ndim == 2:
1742 other = other.T
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_misc.py:1770, in NameScope.__call__(self, *args, **kwargs)
1768 backend = kwargs.pop('backend', None)
1769 jit_fn = self._get_jit_fn(backend)
-> 1770 return jit_fn(*args, **kwargs)
[... skipping hidden 18 frame]
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:571, in XLACustomKernel._register_fallback_lowering.<locals>.fallback_kernel_fn(*args, **kwargs)
569 kernel = entry.kernel_generator(**kwargs)
570 except Exception as exc:
--> 571 raise KernelCompilationError(
572 f"Backend '{backend_to_use}' failed to construct a kernel "
573 f"for primitive '{self.name}' on platform '{platform}': "
574 f"{type(exc).__name__}: {exc}." + _construction_error_suffix()
575 ) from exc
577 try:
578 return kernel(*args)
KernelCompilationError: Backend 'numba' failed to construct a kernel for primitive 'binary_csrmv' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba', 'jax_raw']. Switch with backend='<name>' on this call, or binary_csrmv.set_default('cpu', '<name>') / brainevent.set_backend('cpu', '<name>') to change the default.
Memory, Correctness, and Performance#
Correctness should be checked separately from performance. The comparison below compiles each path, synchronizes device work, then times one steady-state call. Treat the result as a local measurement, not a general ranking.
from time import perf_counter
sparse_step = jax.jit(lambda x: x @ csr)
dense_step = jax.jit(lambda x: x.value.astype(dense.dtype) @ dense)
sparse_step(events).block_until_ready()
dense_step(events).block_until_ready()
start = perf_counter(); sparse_result = sparse_step(events); sparse_result.block_until_ready(); sparse_seconds = perf_counter() - start
start = perf_counter(); dense_result = dense_step(events); dense_result.block_until_ready(); dense_seconds = perf_counter() - start
assert jnp.allclose(sparse_result, dense_result)
print({'sparse_seconds': sparse_seconds, 'dense_seconds': dense_seconds})
Build a Sparse Event-Driven Network#
A layer can reuse one CSR matrix for a batch of event vectors. A two-dimensional BinaryArray performs the batch operation without a Python time-step loop.
event_batch = brainevent.BinaryArray(jnp.array([[1, 0, 1], [0, 1, 0]], dtype=bool))
network_output = event_batch @ csr
print(network_output.shape)
print(network_output)
Inspect the Connectivity Structure#
For small matrices, a dense image is a useful structural diagnostic. Avoid materializing large sparse matrices solely for plotting.
fig, ax = plt.subplots(figsize=(5, 2.5))
ax.spy(csr.todense(), markersize=14)
ax.set(xlabel='post-synaptic index', ylabel='pre-synaptic index', title='Stored connections')
plt.show()
Choosing CSR or CSC#
Choose by the dominant access direction, not by a universal performance claim. Start with CSR for row-oriented forward event propagation. Prefer CSC when repeated column-oriented operations dominate. Measure the real workload on the intended hardware.
Summary and Next Steps#
CSR and CSC encode the same sparse matrix with different orientation. Continue to Fixed Connection Count Structures when every neuron must have a fixed fan-in or fan-out, or to Just-in-Time Connection Matrices when reproducible random connectivity should be generated rather than stored.