Skip to content

dsdbcoo#

Includes our Numba coo datastructure kernels.

Functions:

compute_block_slice #

compute_block_slice(rows: NDArray, cols: NDArray, block_offsets: NDArray, row: int, col: int)

Computes the slice in the data for the given block.

Parameters:

  • rows (NDArray) –

    The rows in the COO matrix.

  • cols (NDArray) –

    The columns in the COO matrix.

  • block_offsets (NDArray) –

    The block offsets.

  • row (int) –

    THe block row.

  • col (int) –

    The block column.

Returns:

  • start ( int ) –

    The start index of the block.

  • stop ( int ) –

    The stop index of the block.

densify_block #

densify_block(block: NDArray, rows: NDArray, cols: NDArray, data: NDArray, block_slice: slice, row_offset: int, col_offset: int)

Fills the dense block with the given data.

Note

If the blocks to be densified get very small, the overhead of starting the CPU threads can lead to worse performance in the jitted version than in the bare API implementation.

Parameters:

  • block (NDArray) –

    Preallocated dense block. Should be filled with zeros.

  • rows (NDArray) –

    The rows at which to fill the block.

  • cols (NDArray) –

    The columns at which to fill the block.

  • data (NDArray) –

    The data to fill the block with.

  • block_slice (slice) –

    The slice of the block to fill.

  • row_offset (int) –

    The row offset of the block.

  • col_offset (int) –

    The column offset of the block

sparsify_block #

sparsify_block(block: NDArray, rows: NDArray, cols: NDArray, data: NDArray)

Fills the data with the given dense block.

Parameters:

  • block (NDArray) –

    The dense block to sparsify.

  • rows (NDArray) –

    The rows at which to fill the block.

  • cols (NDArray) –

    The columns at which to fill the block.

  • data (NDArray) –

    The data to be filled with the block.

compute_block_sort_index #

compute_block_sort_index(coo_rows: NDArray, coo_cols: NDArray, block_sizes: NDArray) -> NDArray

Computes the block-sorting index for a sparse matrix.

Note

This method incurs a bit of memory overhead compared to a naive implementation. No assumptions on the sparsity pattern of the matrix are made here. See the source code for more details.

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 indexing that sorts the data by block-row and -column.