qttools.kernels.datastructure.numba.dsdbcoo

source module qttools.kernels.datastructure.numba.dsdbcoo

Functions

• find_inds — Finds the corresponding indices of the given rows and columns.

• compute_block_slice — Computes the slice in the data for the given block.

• densify_block — Fills the dense block with the given data.

• sparsify_block — Fills the data with the given dense block.

• compute_block_sort_index — Computes the block-sorting index for a sparse matrix.

source find_inds(self_rows: NDArray, self_cols: NDArray, rows: NDArray, cols: NDArray) → tuple[NDArray, NDArray, int]

Finds the corresponding indices of the given rows and columns.

This also counts the number of matches found, which is used to check if the indices contain duplicates.

Parameters

• self_rows : NDArray — The rows of this matrix.

• 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.

• max_counts : int — The maximum number of matches found.

source 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.

source 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

source 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.

source 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.