qttools.datastructures.routines

source module qttools.datastructures.routines

Classes

• BlockMatrix

Functions

• correct_out_range_index

• bd_matmul — Matrix multiplication of two a @ b BD DSDBSparse matrices.

• bd_sandwich — Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.

• btd_matmul — Matrix multiplication of two a @ b BTD DSDBSparse matrices.

• btd_sandwich — Compute the sandwich product a @ b @ a BTD DSDBSparse matrices.

• arrow_partition_halo_comm — Communicate halo blocks between neighboring ranks assuming arrow partitioning.

• arrow_partition_halo_comm_nccl — Communicate halo blocks between neighboring ranks assuming arrow partitioning.

• 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 correct_out_range_index(i: int, k: int, num_blocks: int)

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 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 class BlockMatrix(dsdbsparse: DSDBSparse, local_keys: set[tuple[int, int]], origin: tuple[int, int], mapping=None)

Bases : dict

Methods

• toarray

source method BlockMatrix.toarray()

source arrow_partition_halo_comm(a: BlockMatrix, b: BlockMatrix, a_num_diag: int, b_num_diag: int, start_block: int, end_block: int, comm: Intracomm)

Communicate halo blocks between neighboring ranks assuming arrow partitioning.

NOTE: The method works ONLY IF the ranks need to communicate ONLY with their immediate neighbors, i.e., rank - 1 and rank + 1.

source arrow_partition_halo_comm_nccl(a: BlockMatrix, b: BlockMatrix, a_num_diag: int, b_num_diag: int, start_block: int, end_block: int, comm: Intracomm, nccl_comm)

Communicate halo blocks between neighboring ranks assuming arrow partitioning.

NOTE: The method works ONLY IF the ranks need to communicate ONLY with their immediate neighbors, i.e., rank - 1 and rank + 1.

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