Skip to content

qttools.datastructures.dsdbsparse#

source module qttools.datastructures.dsdbsparse

Classes

  • BlockConfig — Configuration of block-sizes and block-slices for a DSDBSparse matrix.

  • DSDBSparse — Base class for Distributed Stack of Distributed Block-accessible Sparse matrices.

source class BlockConfig(block_sizes: NDArray, block_offsets: NDArray, inds_canonical2bcoo: NDArray | None = None, rowptr_map: dict | None = None, block_slice_cache: dict | None = None)

Bases : object

Configuration of block-sizes and block-slices for a DSDBSparse matrix.

Initializes the block config.

Parameters

  • block_sizes : NDArray — The size of each block in the sparse matrix.

  • block_offsets : NDArray — The block offsets of the block-sparse matrix.

  • inds_canonical2lock : NDArray, optional — A mapping from canonical to block-sorted indices. Default is None.

  • rowptr_map : dict, optional — A mapping from block-coordinates to row-pointers. Default is None.

  • block_slice_cache : dict, optional — A cache for the block slices. Default is None.

source class DSDBSparse(data: NDArray, block_sizes: NDArray, global_stack_shape: tuple | int, return_dense: bool = True, symmetry: bool | None = False, symmetry_op: Callable = xp.conj)

Bases : ABC

Base class for Distributed Stack of Distributed Block-accessible Sparse matrices.

Initializes a DSBDSparse matrix.

Parameters

  • data : NDArray — The local slice of the data. This should be an array of shape (*local_stack_shape, local_nnz). It is the caller's responsibility to ensure that the data is distributed correctly across the ranks.

  • block_sizes : NDArray — The size of each block in the sparse matrix.

  • global_stack_shape : tuple or int — The global shape of the stack. If this is an integer, it is interpreted as a one-dimensional stack.

  • return_dense : bool, optional — Whether to return dense arrays when accessing the blocks. Default is True.

Attributes

  • block_sizes : ArrayLike — Returns the global block sizes.

  • block_offsets : ArrayLike — Returns the block sizes.

  • blocks : _DSDBlockIndexer — Returns a block indexer.

  • sparse_blocks : _DSDBlockIndexer — Returns a block indexer.

  • stack : _DStackIndexer — Returns a stack indexer.

  • data : NDArray — Returns the local slice of the data, masking the padding.

Methods

  • block_diagonal — Returns the block diagonal of the matrix.

  • diagonal — Returns or sets the diagonal elements of the matrix.

  • fill_diagonal — Returns or sets the diagonal elements of the matrix.

  • dtranspose — Performs a distributed transposition of the datastructure.

  • spy — Returns the row and column indices of the non-zero elements.

  • symmetrize — Symmetrizes the matrix with a given operation.

  • to_dense — Converts the local data to a dense array.

  • free_data — Frees the local data.

  • allocate_data — Allocates the local data.

  • from_sparray — Creates a new DSDBSparse matrix from a scipy.sparse array.

  • zeros_like — Creates a new DSDBSparse matrix with the same shape and dtype.

source property DSDBSparse.block_sizes: ArrayLike

Returns the global block sizes.

source property DSDBSparse.block_offsets: ArrayLike

Returns the block sizes.

source property DSDBSparse.blocks: _DSDBlockIndexer

Returns a block indexer.

source property DSDBSparse.sparse_blocks: _DSDBlockIndexer

Returns a block indexer.

source property DSDBSparse.stack: _DStackIndexer

Returns a stack indexer.

source property DSDBSparse.data: NDArray

Returns the local slice of the data, masking the padding.

source method DSDBSparse.block_diagonal(offset: int = 0) → list[NDArray]

Returns the block diagonal of the matrix.

Note that this will cause communication in the block-communicator.

Parameters

  • offset : int, optional — Offset from the main diagonal. Positive values indicate superdiagonals, negative values indicate subdiagonals. Default is 0.

Returns

  • blocks : list — List of block diagonal elements. The length of the list is the number of blocks on the main diagonal minus the offset. Depending on return_dense, the elements are either sparse or dense arrays.

source method DSDBSparse.diagonal(stack_index: tuple = (Ellipsis,)) → NDArray

Returns or sets the diagonal elements of the matrix.

This temporarily sets the return_dense state to True. Note that this will cause communication in the block-communicator.

Returns

  • diagonal : NDArray — The diagonal elements of the matrix.

Raises

  • NotImplementedError

source method DSDBSparse.fill_diagonal(val: NDArray, stack_index: tuple = (Ellipsis,)) → NDArray

Returns or sets the diagonal elements of the matrix.

This temporarily sets the return_dense state to True. Note that this will cause communication in the block-communicator.

Returns

  • diagonal : NDArray — The diagonal elements of the matrix.

Raises

  • NotImplementedError

source method DSDBSparse.dtranspose(discard: bool = False) → None

Performs a distributed transposition of the datastructure.

This is done by reshaping the local data, then performing an in-place Alltoall communication, and finally reshaping the data back to the correct new shape.

The local reshaping of the data cannot be done entirely in-place. This can lead to pronounced memory peaks if all ranks start reshaping concurrently, which can be mitigated by using more ranks and by not forcing a synchronization barrier right before calling dtranspose.

Parameters

  • discard : bool, optional — Whether to perform a "fake" transposition. Default is False. This is useful if you want to get the correct data shape after a transposition, but do not want to perform the actual all-to-all communication.

source method DSDBSparse.spy() → tuple[NDArray, NDArray]

Returns the row and column indices of the non-zero elements.

This is essentially the same as converting the sparsity pattern to coordinate format. The returned sparsity pattern is not sorted.

Returns

  • rows : NDArray — Row indices of the non-zero elements.

  • cols : NDArray — Column indices of the non-zero elements.

source method DSDBSparse.symmetrize(op: Callable[[NDArray, NDArray], NDArray] = xp.add) → None

Symmetrizes the matrix with a given operation.

This is done by setting the data to the result of the operation applied to the data and its conjugate transpose.

Note

This assumes that the matrix's sparsity pattern is symmetric.

Parameters

  • op : callable, optional — The operation to apply to the data and its conjugate transpose. Default is xp.add, so that the matrix is Hermitian after calling.

source method DSDBSparse.to_dense() → NDArray

Converts the local data to a dense array.

This is dumb, unless used for testing and debugging.

Returns

  • arr : NDArray — The dense array of shape (*local_stack_shape, *shape).

source method DSDBSparse.free_data() → None

Frees the local data.

source method DSDBSparse.allocate_data() → None

Allocates the local data.

source classmethod DSDBSparse.from_sparray(arr: sparse.spmatrix, block_sizes: NDArray, global_stack_shape: tuple, symmetry: bool | None = False, symmetry_op: Callable = xp.conj) → DSDBSparse

Creates a new DSDBSparse matrix from a scipy.sparse array.

Parameters

  • arr : sparse.spmatrix — The sparse array to convert.

  • block_sizes : NDArray — The size of all the blocks in the matrix.

  • global_stack_shape : tuple — The global shape of the stack of matrices. The provided sparse matrix is replicated across the stack.

Returns

source classmethod DSDBSparse.zeros_like(dsdbsparse: DSDBSparse) → DSDBSparse

Creates a new DSDBSparse matrix with the same shape and dtype.

All non-zero elements are set to zero, but the sparsity pattern is preserved.

Parameters

  • dsdbsparse : DSDBSparse — The matrix to copy the shape and dtype from.

Returns