Skip to content

routines#

Includes routines for multiplication of distributed block banded matrices.

Classes:

  • BlockMatrix –

    Block-sparse matrix class, including halo blocks for communication.

Functions:

  • arrow_partition_halo_comm –

    Communicate halo blocks between neighboring ranks assuming arrow

  • bd_matmul –

    Matrix multiplication of two a @ b BD DSDBSparse matrices.

  • bd_sandwich –

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

BlockMatrix #

BlockMatrix(dsdbsparse: DSDBSparse, local_keys: set[tuple[int, int]], origin: tuple[int, int], mapping: dict | None = None)

Block-sparse matrix class, including halo blocks for communication.

Any local block keys are stored in the blocks of the dsdbsparse, while non-local block keys are stored in a separate dictionary.

Parameters:

  • dsdbsparse (DSDBSparse) –

    The underlying DSDBSparse matrix.

  • local_keys (set[tuple[int, int]]) –

    The set of block keys that are local to the current rank.

  • origin (tuple[int, int]) –

    The global index of the first local block. This is used to compute the local block keys from the global block keys.

  • mapping (dict[tuple[int, int], ndarray], default: None ) –

    A mapping from non-local block keys to their corresponding data arrays. This is used to store halo blocks that are communicated between ranks. The default is None, which means that there are no non-local blocks.

Methods:

  • __getitem__ –

    Gets the block corresponding to the given key.

  • __setitem__ –

    Sets the block corresponding to the given key.

  • toarray –

    Converts the BlockMatrix to a dense array.

__getitem__ #

__getitem__(key)

Gets the block corresponding to the given key.

__setitem__ #

__setitem__(key, val)

Sets the block corresponding to the given key.

toarray #

toarray()

Converts the BlockMatrix to a dense array.

Debugging method to check the correctness of the block matrix.

arrow_partition_halo_comm #

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

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.

Parameters:

  • a (BlockMatrix) –

    The first block matrix.

  • b (BlockMatrix) –

    The second block matrix.

  • a_num_diag (int) –

    The number of diagonals in the first block matrix.

  • b_num_diag (int) –

    The number of diagonals in the second block matrix.

  • start_block (int) –

    The index of the first block to communicate.

  • end_block (int) –

    The index of the last block to communicate.

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.