brainevent.coo2csr

Contents

brainevent.coo2csr#

brainevent.coo2csr(row_ids, col_ids, *, shape)[source]#

Convert COO format index arrays to CSR format.

Transforms a sparse matrix representation from Coordinate (COO) format (explicit row and column index arrays) to Compressed Sparse Row (CSR) format. The implementation automatically selects NumPy or JAX operations based on the type of the input arrays.

Parameters:

  • row_ids (Array | ndarray) – Row index array in COO format. Contains the row index for each non-zero element.

  • col_ids (Array | ndarray) – Column index array in COO format. Contains the column index for each non-zero element.

  • shape (Tuple[int, int]) – A (n_rows, n_cols) tuple specifying the dimensions of the sparse matrix. Keyword-only argument.

Returns:

  • csr_indptr (jax.Array or numpy.ndarray) – Row pointer array in CSR format. For a matrix with m rows, this has length m + 1. Element csr_indptr[i] gives the position in csr_indices where row i starts.

  • csr_indices (jax.Array or numpy.ndarray) – Column index array in CSR format. Contains the column index for each non-zero element, ordered by row.

  • order (jax.Array or numpy.ndarray) – Permutation array that reorders data values from COO order to CSR order. If data is the COO data array, then data[order] gives the values in CSR order.

See also

csr_to_coo_index

Convert CSR indices to COO format.

coo_to_csc_index

Convert COO indices to CSC format.

csr_to_csc_index

Convert CSR indices directly to CSC format.

Notes

When JAX arrays are provided, the computation is wrapped in jax.ensure_compile_time_eval() so that it executes at trace time rather than at runtime. Entries are grouped by row using a stable sort, so the relative order of elements within each row is preserved (the column indices within a row are not themselves sorted).

Examples

>>> import numpy as np
>>> from brainevent._misc import coo2csr
>>> row_ids = np.array([0, 2, 1, 0, 2])
>>> col_ids = np.array([0, 3, 1, 2, 0])
>>> indptr, indices, order = coo2csr(row_ids, col_ids, shape=(3, 4))
>>> print(indptr)
[0 2 3 5]
>>> print(indices)
[0 2 1 3 0]
Contents