dedalus.tools.linalg_gpu

Linear algebra routines using cupy.

Module Contents

class CustomCupyUpperTriangularSolver(matrix)

Hacky class to save spsm_descr for reuse in spsm for triangular solves.

solve(b, lower=True, overwrite_A=False, overwrite_b=False, unit_diagonal=False)

Solves a sparse triangular system A x = b.

Parameters:

A (cupyx.scipy.sparse.spmatrix) – Sparse matrix with dimension (M, M).
b (cupy.ndarray) – Dense vector or matrix with dimension (M) or (M, K).
lower (bool) – Whether A is a lower or upper triangular matrix. If True, it is lower triangular, otherwise, upper triangular.
overwrite_A (bool) – (not supported)
overwrite_b (bool) – Allows overwriting data in b.
unit_diagonal (bool) – If True, diagonal elements of A are assumed to be 1 and will not be referenced.

Returns:

Solution to the system A x = b. The shape is the same as b.

Return type:

cupy.ndarray

apply_csr_mid_kernel(data, indices, indptr, x3, y3, N1, N2i, N2o, N3)

cupy_apply_csr(matrix, array, axis, out): Apply CSR matrix to arbitrary axis of array.

cupy_solve_upper_csr(matrix, array, axis, out): Solve upper triangular CSR matrix along specified axis of an array.

cupy_solve_upper_csr_first(matrix, array, out): Solve upper triangular CSR matrix along first axis of 2D array.

cupy_solve_upper_csr_last(matrix, array, out): Solve upper triangular CSR matrix along last axis of 2D array.

cupy_solve_upper_csr_mid(matrix, array, out): Solve upper triangular CSR matrix along middle axis of 3D array.

cupy_solve_upper_csr_vec(matrix, vec, out): Solve upper triangular CSR matrix along a vector.

custom_SuperLU_solve(self, rhs, trans='N', spsm_descr=None): Custom SuperLU solve wrapper to save spsm_descr, since spsm_analysis takes lots of time.

custom_spsm(a, b, alpha=1.0, lower=True, unit_diag=False, transa=False, spsm_descr=None): Custom spsm wrapper to save spsm_descr, since spsm_analysis takes lots of time.