datastructures#

Includes data structures for distributed block-accessible sparse matrices and related routines.

Modules:

dsdbcoo –

Includes the distributed block-accessible COO matrix data structure.
dsdbcsr –

Includes the distributed block-accessible CSR matrix data structure.
dsdbsparse –

Includes the abstract base class for distributed block-accessible sparse matrices.
routines –

Includes routines for multiplication of distributed block banded matrices.

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

Functions:

bd_matmul –

Matrix multiplication of two a @ b BD DSDBSparse matrices.
bd_sandwich –

Matrix multiplication of three a @ b @ a BD DSDBSparse matrices.

DSDBCOO #

DSDBCOO(dtype: dtype[generic], rows: NDArray, cols: NDArray, block_sizes: NDArray, local_stack_shape: tuple | int, global_stack_shape: tuple | int, symmetry: str | None = None)

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.

Parameters:

dtype (dtype[generic]) –

The data type of the matrix.
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.
local_stack_shape (tuple or int) –

The local shape of the stack. If this is an integer, it is interpreted as a one-dimensional stack.
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.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.

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 symmetry.
empty_like –

Creates a new DSDBCOO matrix with the same shape and
from_sparray –

Constructs a DSDBCOO matrix from a sparse matrix.
to_dense –

Converts the local data to a dense array.

block_sizes #

block_sizes(block_sizes: NDArray) -> None

Sets new block sizes for the matrix.

Parameters:

block_sizes (NDArray) –

The new block sizes.

spy #

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.

Note

In the block distributed case, this returns the local sparsity pattern.

Returns:

rows ( NDArray ) –

Row indices of the non-zero elements.
cols ( NDArray ) –

Column indices of the non-zero elements.

symmetrize #

symmetrize(symmetry: str) -> None

Symmetrizes the matrix with a given symmetry.

Note

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

Parameters:

symmetry (str) –

The symmetry to enforce. This can be "symmetric", "hermitian", "skew-symmetric", or "skew-hermitian".

empty_like `classmethod` #

empty_like(dsdbsparse: DSDBCOO) -> DSDBCOO

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

Note

There is no data allocated in the new matrix. The sparsity pattern is the same as the original matrix.

Parameters:

dsdbsparse (DSDBCOO) –

The matrix to copy the shape and dtype from.

Returns:

DSDBCOO –

The new DSDBCOO matrix.

from_sparray `classmethod` #

from_sparray(sparray: spmatrix, block_sizes: NDArray, global_stack_shape: tuple, symmetry: str | None = None, dtype: dtype[generic] = complex128, allocate: bool = True) -> DSDBCOO

Constructs a DSDBCOO matrix from a sparse matrix.

This essentially distributed the matrix across the stack and block communicators.

Parameters:

sparray (spmatrix) –

The sparse matrix from which to use the sparsity pattern.
block_sizes (NDArray) –

The block sizes of the block-sparse matrix.
global_stack_shape (tuple) –

The global shape of the stack.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.
dtype (dtype, default: complex128 ) –

The data type of the matrix. Default is xp.complex128.
allocate (bool, default: True ) –

Whether to allocate the data of the resulting matrix. Default is True.

Returns:

DSDBCOO –

The new DSDBCOO matrix.

to_dense #

to_dense()

Converts the local data to a dense array.

This is dumb, unless used for testing and debugging.

Warning

This creates a very large dense matrix.

Returns:

arr ( NDArray ) –

The dense array of shape (*local_stack_shape, *shape).

DSDBCSR #

DSDBCSR(dtype: dtype[generic], cols: NDArray, rowptr_map: dict, block_sizes: NDArray, local_stack_shape: tuple | int, global_stack_shape: tuple, symmetry: str | None = None)

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.

Parameters:

dtype (dtype[generic]) –

The data type of the matrix.
cols (NDArray) –

The column indices.
rowptr_map (dict) –

The row pointer map.
block_sizes (NDArray) –

The size of each block in the sparse matrix.
local_stack_shape (tuple or int) –

The local shape of the stack. If this is an integer, it is interpreted as a one-dimensional stack.
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.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.

Methods:

block_sizes –

Sets new block sizes for the matrix.
symmetrize –

Symmetrizes the matrix with a given symmetry.
spy –

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

Creates a new DSDBCSR matrix with the same shape and dtype.
from_sparray –

Creates a new DSDBCSR matrix from a scipy.sparse array.
to_dense –

Converts the local data to a dense array.

block_sizes #

block_sizes(block_sizes: NDArray) -> None

Sets new block sizes for the matrix.

Parameters:

block_sizes (NDArray) –

The new block sizes.

symmetrize #

symmetrize(symmetry: str) -> None

Symmetrizes the matrix with a given symmetry.

Note

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

Parameters:

symmetry (str) –

The symmetry to enforce. This can be "symmetric", "hermitian", "skew-symmetric", or "skew-hermitian".

spy #

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.

Note

In the block distributed case, this returns the local sparsity pattern.

Warning

This not performant.

Returns:

rows ( NDArray ) –

Row indices of the non-zero elements.
cols ( NDArray ) –

Column indices of the non-zero elements.

empty_like `classmethod` #

empty_like(dsdbsparse: DSDBCSR) -> DSDBCSR

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

Note

There is no data allocated in the new matrix. The sparsity pattern is the same as the original matrix.

Parameters:

dsdbsparse (DSDBCSR) –

The matrix to copy the shape and dtype from.

Returns:

DSDBCSR –

The new DSDBCSR matrix.

from_sparray `classmethod` #

from_sparray(sparray: spmatrix, block_sizes: NDArray, global_stack_shape: tuple, symmetry: str | None = None, dtype: dtype[generic] = complex128, allocate: bool = True) -> DSDBCSR

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

This essentially distributed the matrix across the stack and block communicators.

Parameters:

sparray (spmatrix) –

The sparse matrix from which to use the sparsity pattern.
block_sizes (NDArray) –

The block sizes of the block-sparse matrix.
global_stack_shape (tuple) –

The global shape of the stack.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.
dtype (dtype, default: complex128 ) –

The data type of the matrix. Default is xp.complex128.
allocate (bool, default: True ) –

Whether to allocate the data of the resulting matrix. Default is True.

Returns:

DSDBCSR –

The new DSDBCSR matrix.

to_dense #

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

DSDBSparse #

DSDBSparse(dtype: dtype[generic], block_sizes: NDArray, nnz: int, local_stack_shape: tuple | int, global_stack_shape: tuple | int, index_type: int32 | int64, symmetry: str | None = None)

Bases: ABC

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

Parameters:

dtype (dtype[generic]) –

The data type of the matrix.
block_sizes (NDArray) –

The size of each block in the sparse matrix.
nnz (int) –

The number of non-zero elements in the sparse matrix.
local_stack_shape (tuple or int) –

The local shape of the stack. If this is an integer, it is interpreted as a one-dimensional stack.
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.
index_type (int32 | int64) –

The index type to use for the sparse matrix. This is relevant for the low level kernels to avoid unnecessary type conversions.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.

Methods:

__repr__ –

Returns a string representation of the object.
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 symmetry.
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.
empty_like –

Creates a new DSDBSparse matrix with the same shape and

Attributes:

block_sizes (ArrayLike) –

Returns the global block sizes.
block_offsets (ArrayLike) –

Returns the block sizes.
blocks (_DSDBlockIndexer) –

Returns a block indexer.
stack (_DStackView) –

Returns a stack indexer.
data (NDArray) –

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

block_sizes `property` `writable` #

block_sizes: ArrayLike

Returns the global block sizes.

block_offsets `property` #

block_offsets: ArrayLike

Returns the block sizes.

blocks `property` #

blocks: _DSDBlockIndexer

Returns a block indexer.

stack `property` #

stack: _DStackView

Returns a stack indexer.

data `property` `writable` #

data: NDArray

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

repr #

__repr__() -> str

Returns a string representation of the object.

diagonal #

diagonal(stack_index: tuple = (Ellipsis,)) -> NDArray

Returns or sets the diagonal elements of the matrix.

Note

In the block distributed case, this returns the local diagonal elements.

Parameters:

stack_index (tuple, default: (Ellipsis,) ) –

The index in the stack. Default is (Ellipsis,).

Returns:

diagonal ( NDArray ) –

The diagonal elements of the matrix.

fill_diagonal #

fill_diagonal(val: NDArray, stack_index: tuple = (Ellipsis,)) -> NDArray

Returns or sets the diagonal elements of the matrix.

Parameters:

val (NDArray) –

The value(s) to set along the diagonal.
stack_index (tuple, default: (Ellipsis,) ) –

The index in the stack. Default is (Ellipsis,).

Returns:

diagonal ( NDArray ) –

The diagonal elements of the matrix.

dtranspose #

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, default: False ) –

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.

spy `abstractmethod` #

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.

Note

In the block distributed case, this returns the local sparsity pattern including the offset.

Returns:

rows ( NDArray ) –

Row indices of the non-zero elements.
cols ( NDArray ) –

Column indices of the non-zero elements.

symmetrize `abstractmethod` #

symmetrize(symmetry: str) -> None

Symmetrizes the matrix with a given symmetry.

Note

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

Parameters:

symmetry (str) –

The symmetry to enforce. This can be "symmetric", "hermitian", "skew-symmetric", or "skew-hermitian".

to_dense `abstractmethod` #

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

free_data #

free_data() -> None

Frees the local data.

allocate_data #

allocate_data(stack_size: int | None = None) -> None

Allocates the local data.

Note

This should not be called with a non-None stack size if the data will be dtransposed. The data is not zeroed. It is the user responsibility to ensure that the data is initialized correctly.

Parameters:

stack_size (int, default: None ) –

The size of the stack dimension to allocate. If None, the full stack size is used. Default is None.

from_sparray `abstractmethod` `classmethod` #

from_sparray(sparray: spmatrix, block_sizes: NDArray, global_stack_shape: tuple, symmetry: str | None = None, dtype: dtype[generic] = complex128, allocate: bool = True) -> DSDBSparse

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

This essentially distributed the matrix across the stack and block communicators.

Parameters:

sparray (spmatrix) –

The sparse matrix from which to use the sparsity pattern.
block_sizes (NDArray) –

The block sizes of the block-sparse matrix.
global_stack_shape (tuple) –

The global shape of the stack.
symmetry (str | None, default: None ) –

The symmetry of the matrix. This can be "symmetric", "hermitian", "skew-symmetric", "skew-hermitian", or None. Default is None.
dtype (dtype, default: complex128 ) –

The data type of the matrix. Default is xp.complex128.
allocate (bool, default: True ) –

Whether to allocate the data of the resulting matrix. Default is True.

Returns:

DSDBSparse –

The new DSDBSparse matrix.

empty_like `abstractmethod` `classmethod` #

empty_like(dsdbsparse: DSDBSparse) -> DSDBSparse

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

Note

There is no data allocated in the new matrix. The sparsity pattern is the same as the original matrix.

Parameters:

dsdbsparse (DSDBSparse) –

The matrix to copy the shape and dtype from.

Returns:

DSDBSparse –

The new DSDBSparse matrix.

bd_matmul #

bd_matmul(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 = None) -> BlockMatrix

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 | None) –

The output matrix. This matrix must have the same block size as a and b. It will compute up to out_num_diag diagonals.
a_num_diag (int, default: 3 ) –

The number of diagonals in the first input matrix.
b_num_diag (int, default: 3 ) –

The number of diagonals in the second input matrix.
out_num_diag (int, default: 5 ) –

The number of diagonals in output matrices
start_block (int, default: 0 ) –

The index of the first block to compute.
end_block (int | None, default: None ) –

The index of the last block to compute. If None, it will compute up to the last block.

Returns:

BlockMatrix –

The resulting block matrix of the multiplication. Even if the output is not None, the method returns the corresponding BlockMatrix for convenience.

bd_sandwich #

bd_sandwich(a: DSDBSparse | _DStackView, b: DSDBSparse | _DStackView, out: DSDBSparse | _DStackView, in_num_diag: int = 3, out_num_diag: int = 7, start_block: int = 0, end_block: int = None) -> None

Matrix multiplication of three a @ b @ a 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, default: 3 ) –

The number of diagonals in input matrices
out_num_diag (int, default: 7 ) –

The number of diagonals in output matrices
start_block (int, default: 0 ) –

The index of the first block to compute.
end_block (int, default: None ) –

The index of the last block to compute. If None, it will compute up to the last block.

datastructures#

DSDBCOO #

block_sizes #

spy #

symmetrize #

empty_like classmethod #

from_sparray classmethod #

to_dense #

DSDBCSR #

block_sizes #

symmetrize #

spy #

empty_like classmethod #

from_sparray classmethod #

to_dense #

DSDBSparse #

block_sizes property writable #

block_offsets property #

blocks property #

stack property #

data property writable #

__repr__ #

diagonal #

fill_diagonal #

dtranspose #

spy abstractmethod #

symmetrize abstractmethod #

to_dense abstractmethod #

free_data #

allocate_data #

from_sparray abstractmethod classmethod #

empty_like abstractmethod classmethod #

bd_matmul #

bd_sandwich #

empty_like `classmethod` #

from_sparray `classmethod` #

empty_like `classmethod` #

from_sparray `classmethod` #

block_sizes `property` `writable` #

block_offsets `property` #

blocks `property` #

stack `property` #

data `property` `writable` #

repr #

spy `abstractmethod` #

symmetrize `abstractmethod` #

to_dense `abstractmethod` #

from_sparray `abstractmethod` `classmethod` #

empty_like `abstractmethod` `classmethod` #