如何将矩阵乘法改写成matrix free的形式

目标是希望得到如下的矩阵A

$$A=\nabla C{M}^{-1}\nabla C^T+\tilde{\alpha }$$

其中gradC是mx3n的矩阵，它是CSR格式的，每行有12个非零值。M是对角矩阵，3nx3n。alpha_tilde是对角矩阵。

gradC存储：Mx12的数组。同时有一个辅助的数组为col_ind，也是Mx12。col_ind记录的是每行的列下标。

M存储：实际上存每个M的倒数。inv_mass数组为Nx1的。通过np.repeat(3)将其扩容三倍。例如之前是1,2,3, 扩容后是1,1,1,2,2,2,3,3,3

    inv_mass3 = np.repeat(inv_mass, 3)

col_ind: col_ind是从tet_indices得到的。它也需要扩容， 例如之前是1,2,3, 扩容后是1,2,3,4,5,6,7,8,9

    for i in range(M):
        for j in range(4):
            for k in range(3):
                col_indices[i, j*3+k] = tet_indices[i, j] * 3 + k

首先来乘gradC@inv_mass

def spmat12_diamat_prod(A, D, col_ind):
    """
    A@D
    A(Mx12): sparse matrix, csr format, each line has 12 nonzeros
    D(M): diagonal matrix, dense format
    col_ind(Mx12): column indices of A
    """
    M = A.shape[0]
    C = np.zeros_like(A)
    for i in range(M):
        for d in range(12):
            index = col_ind[i,d]
            C[i,d] = A[i,d] * D[index]
    return C

然后来乘gradC_M@gradCT

def spmat12_spmat12T_prod(A, B, col_ind):
    """
    A@B
    A(Mx12): sparse matrix, csr format, each line has 12 nonzeros
    B(Mx12): sparse matrix, csr format, each line has 12 nonzeros
    col_ind(Mx12): column indices of A and B
    """
    M = A.shape[0]
    C = np.zeros((M, M), dtype=np.float32)
    for i in range(M):
        for j in range(M):
            for ki in range(12):
                for kj in range(12):
                    ni = col_ind[i, ki]
                    nj = col_ind[j, kj]
                    if ni == nj:
                        C[i, j] += A[i, ki] * B[j, kj]
    return C

最后来加上alpha_tilde

def spmat12_add_diag(A, col_ind, diag):
    """
    A + alpha * I
    A(Mx12): sparse matrix, csr format, each line has 12 nonzeros
    diag(M): diagonal matrix, dense format
    """
    M = A.shape[0]
    C = np.zeors((M, M), dtype=np.float32)
    for i in range(M):
        C[i,i] = A[i,i] + diag[i]
    return C