qttools.datastructures#

Classes

DSDBCOO — A Distributed Stack of Distributed Block-accessible COO matrices.
DSDBCSR — A Distributed Stack of Distributed Block-accessible CSR matrices.
DSDBSparse — Base class for Distributed Stack of Distributed Block-accessible Sparse matrices.

Functions

btd_matmul — Matrix multiplication of two a @ b BTD DSDBSparse matrices.
btd_sandwich — Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.
bd_matmul — Matrix multiplication of two a @ b BD DSDBSparse matrices.
bd_sandwich — Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.
bd_matmul_distr — Matrix multiplication of two a @ b BD DSDBSparse matrices.
bd_sandwich_distr — Matrix multiplication of two a @ b BD DSDBSparse matrices.

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

Bases : DSDBSparse

A Distributed Stack of Distributed Block-accessible COO matrices.

Note

It is the caller's responsibility to ensure that the data is distributed correctly across the ranks.

Initializes a DSDBCOO matrix.

Parameters

data : NDArray — The local slice of the data. This should be an array of shape (*local_stack_shape, local_nnz).
rows : NDArray — The local row indices of the COO matrix.
cols : NDArray — The local column indices of the COO matrix.
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.
symmetry : bool, optional — Whether the matrix is symmetric. Default is False.
symmetry_op : callable, optional — The operation to use for the symmetry. Default is xp.conj.

Attributes

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_sizes — Sets new block sizes for the matrix.
spy — Returns the row and column indices of the non-zero elements.
symmetrize — Symmetrizes the matrix with a given operation.
zeros_like — Creates a new DSDBCOO matrix with the same shape and dtype.
from_sparray — Constructs a DSDBCOO matrix from a COO matrix.
to_dense — Converts the local data to a dense array.

source method DSDBCOO.block_sizes(block_sizes: NDArray) → None

Sets new block sizes for the matrix.

Parameters

block_sizes : NDArray — The new block sizes.

Raises

NotImplementedError
ValueError

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

Raises

NotImplementedError

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

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

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

Parameters

dsdbcoo : DSDBCOO — The matrix to copy the shape and dtype from.

Returns

dsdbcoo — The new DSDBCOO matrix.

Raises

ValueError

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

Constructs a DSDBCOO matrix from a COO matrix.

This essentially distributes the COO matrix across the participating ranks.

Parameters

sparray : sparse.coo_matrix — The COO matrix to distribute.
block_sizes : NDArray — The block sizes of the block-sparse matrix.
global_stack_shape : tuple — The global shape of the stack.
symmetry : bool, optional — Whether to enforce symmetry in the matrix. Default is False.
symmetry_op : callable, optional — The operation to use for the symmetry. Default is xp.conj.

Returns

DSDBCOO — The new DSDBCOO matrix.

Raises

ValueError

source method DSDBCOO.to_dense()

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

Raises

ValueError

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

Bases : DSDBSparse

A Distributed Stack of Distributed Block-accessible CSR matrices.

This DSDBSparse implementation uses a block-compressed sparse row format to store the sparsity pattern of the matrix. The data is sorted by block-row and -column. We use a row pointer map together with the column indices to access the blocks efficiently.

Note

It is the caller's responsibility to ensure that the data is distributed correctly across the ranks.

Initializes the DBCSR matrix.

Parameters

data : NDArray — The local slice of the data. This should be an array of shape (*local_stack_shape, nnz). It is the caller's responsibility to ensure that the data is distributed correctly across the ranks.
cols : NDArray — The column indices.
rowptr_map : dict — The row pointer map.
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.
symmetry : bool, optional — Whether the matrix is symmetric. Default is False.
symmetry_op : callable, optional — The operation to use for the symmetry. Default is xp.conj.

Attributes

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_sizes — Sets new block sizes for the matrix.
ltranspose — Performs a local transposition of the matrix.
symmetrize — Symmetrizes the matrix.
spy — Returns the row and column indices of the non-zero elements.
zeros_like — Creates a new DSDBSparse matrix with the same shape and dtype.
from_sparray — Creates a new DSDBSparse matrix from a scipy.sparse array.
to_dense — Converts the local data to a dense array.

source method DSDBCSR.block_sizes(block_sizes: NDArray) → None

Sets new block sizes for the matrix.

Parameters

block_sizes : NDArray — The new block sizes.

Raises

NotImplementedError
ValueError

source method DSDBCSR.ltranspose(copy=False) → None | DSDBCSR

Performs a local transposition of the matrix.

Parameters

copy : bool, optional — Whether to return a new object. Default is False.

Returns

None | DSDBCSR — The transposed matrix. If copy is False, this is None.

Raises

NotImplementedError

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

Symmetrizes the matrix.

NOTE: Assumes that the matrix's sparsity pattern is symmetric.

Parameters

op : callable, optional — The operation to perform on the symmetric elements. Default is addition.

Raises

NotImplementedError

source method DSDBCSR.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 classmethod DSDBCSR.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

DSDBSparse — The new DSDBSparse matrix.

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

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.
symmetry : bool, optional — Whether to enforce symmetry in the matrix. Default is False.
symmetry_op : callable, optional — The operation to use for the symmetry. Default is xp.conj.

Returns

DSDBCSR — The new DSDBCSR matrix.

source method DSDBCSR.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).

Raises

ValueError

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

DSDBSparse — The new DSDBSparse matrix.

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

DSDBSparse — The new DSDBSparse matrix.

source btd_matmul(a: DSDBSparse, b: DSDBSparse, out: DSDBSparse, spillover_correction: bool = False)

Matrix multiplication of two a @ b BTD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block tridiagonal matrix.
b : DSDBSparse — The second block tridiagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a and b. It will compute up to pentadiagonal.
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.

Raises

ValueError

source btd_sandwich(a: DSDBSparse, b: DSDBSparse, out: DSDBSparse, spillover_correction: bool = False)

Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block tridiagonal matrix.
b : DSDBSparse — The second block tridiagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a, and b. It will compute up to heptadiagonal.
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.

Raises

ValueError

source bd_matmul(a: DSDBSparse, b: DSDBSparse | list[DSDBSparse], out: DSDBSparse | None, b_op: Callable | None = None, in_num_diag: int = 3, out_num_diag: int = 5, spillover_correction: bool = False, accumulator_dtype=None)

Matrix multiplication of two a @ b BD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block diagonal matrix.
b : DSDBSparse — The second block diagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a and b. It will compute up to out_num_diag diagonals.
in_num_diag : int — The number of diagonals in input matrices
out_num_diag : int — The number of diagonals in output matrices
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.
accumulator_dtype : data type, optional — The data type of the temporary accumulator matrices. The default is complex128.
TODO : replace @ by appropriate gemm

Raises

ValueError

source bd_sandwich(a: DSDBSparse, b: DSDBSparse, out: DSDBSparse | None, in_num_diag: int = 3, out_num_diag: int = 7, spillover_correction: bool = False, accumulator_dtype=None, accumulate: bool = False)

Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block tridiagonal matrix.
b : DSDBSparse — The second block tridiagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a, and b. It will compute up to out_num_diag diagonals.
in_num_diag : int — The number of diagonals in input matrices
out_num_diag : int — The number of diagonals in output matrices
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.
accumulator_dtype : data type, optional — The data type of the temporary accumulator matrices. The default is complex128.
TODO : replace @ by appropriate gemm

Raises

ValueError

source bd_matmul_distr(a: DSDBSparse | BlockMatrix, b: DSDBSparse | BlockMatrix, out: DSDBSparse | None, a_num_diag: int = 3, b_num_diag: int = 3, out_num_diag: int = 5, start_block: int = 0, end_block: int = None, spillover_correction: bool = False, accumulator_dtype=None)

Matrix multiplication of two a @ b BD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block diagonal matrix.
b : DSDBSparse — The second block diagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a and b. It will compute up to out_num_diag diagonals.
in_num_diag : int — The number of diagonals in input matrices
out_num_diag : int — The number of diagonals in output matrices
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.
accumulator_dtype : data type, optional — The data type of the temporary accumulator matrices. The default is complex128.
TODO : replace @ by appropriate gemm

Raises

ValueError

source bd_sandwich_distr(a: DSDBSparse, b: DSDBSparse, out: DSDBSparse | None, in_num_diag: int = 3, out_num_diag: int = 7, start_block: int = 0, end_block: int = None, spillover_correction: bool = False, accumulator_dtype=None)

Matrix multiplication of two a @ b BD DSDBSparse matrices.

Parameters

a : DSDBSparse — The first block diagonal matrix.
b : DSDBSparse — The second block diagonal matrix.
out : DSDBSparse — The output matrix. This matrix must have the same block size as a and b. It will compute up to out_num_diag diagonals.
in_num_diag : int — The number of diagonals in input matrices
out_num_diag : int — The number of diagonals in output matrices
spillover_correction : bool, optional — Whether to apply spillover corrections to the output matrix. This is necessary when the matrices represent open-ended systems. The default is False.
accumulator_dtype : data type, optional — The data type of the temporary accumulator matrices. The default is complex128.
TODO : replace @ by appropriate gemm

Raises

NotImplementedError