Skip to content

dsdbcsr#

Includes our Numba csr datastructure kernels.

Functions:

find_inds #

find_inds(rowptr_map: dict[tuple, NDArray], block_offsets: NDArray, self_cols: NDArray, rows: NDArray, cols: NDArray) -> tuple[NDArray, NDArray]

Finds the corresponding indices of the given rows and columns.

Parameters:

  • rowptr_map (dict) –

    The row pointer map.

  • block_offsets (NDArray) –

    The block offsets.

  • self_cols (NDArray) –

    The columns of this matrix.

  • rows (NDArray) –

    The rows to find the indices for.

  • cols (NDArray) –

    The columns to find the indices for.

Returns:

  • inds ( NDArray ) –

    The indices of the given rows and columns.

  • value_inds ( NDArray ) –

    The matching indices of this matrix.

densify_block #

densify_block(block: NDArray, block_offset: NDArray, self_cols: NDArray, rowptr: NDArray, data: NDArray)

Fills the dense block with the given data.

Parameters:

  • block (NDArray) –

    Preallocated dense block. Should be filled with zeros.

  • block_offset (NDArray) –

    The block offset.

  • self_cols (NDArray) –

    The column indices of this matrix.

  • rowptr (NDArray) –

    The row pointer of this matrix block.

  • data (NDArray) –

    The data to fill the block with.

sparsify_block #

sparsify_block(block: NDArray, block_offset: NDArray, self_cols: NDArray, rowptr: NDArray, data: NDArray)

Fills the data with the given dense block.

Parameters:

  • block (NDArray) –

    The dense block to sparsify.

  • block_offset (NDArray) –

    The block offset.

  • self_cols (NDArray) –

    The column indices of this matrix.

  • rowptr (NDArray) –

    The row pointer of this matrix block.

  • data (NDArray) –

    The data to be filled with the block.

compute_rowptr_map #

compute_rowptr_map(coo_rows: NDArray, coo_cols: NDArray, block_sizes: NDArray) -> dict

Computes the rowptr map for a sparse matrix.

Parameters:

  • coo_rows (NDArray) –

    The row indices of the matrix in coordinate format.

  • coo_cols (NDArray) –

    The column indices of the matrix in coordinate format.

  • block_sizes (NDArray) –

    The block sizes of the block-sparse matrix we want to construct.

Returns:

  • sort_index ( NDArray ) –

    The block-sorting index for the sparse matrix.

  • rowptr_map ( dict ) –

    The row pointer map, describing the block-sparse matrix in blockwise column-sparse-row format.