一些关于线性代数的程序

一些关于线性代数的程序

1. 基于分块的大矩阵求逆

scipy不能对太大的矩阵进行求逆操作，下面这个函数通过分块求逆处理了这个问题，但是分块求逆对矩阵要求比较高，例如需要余子式可逆等：

import numpy as np
from matplotlib import pyplot as plt 
import matplotlib
from scipy.linalg import inv, eigh

def split_matrix(x):
    '''
    Split matrix into four part, return a view of the four parts\n
    x: 2D ndarray, the matrix to split\n
    return: tuple\n
        x1, x2, x3, x4: the left upper, right upper, left lower, right lower parts of the input matrix. They are all view, i.e. do not own data
    '''
    assert x.shape[0] == x.shape[1]
    ii = x.shape[0]//2
    # left upper, right upper, left lower, right lower
    return x[:ii,:ii], x[:ii,ii:], x[ii:,:ii], x[ii:,ii:]

def block_inv(x, max_size, hermitian=False, inv_func=inv, out=None, **kwargs):
    '''
    Inverse matrix based on block inverse\n
    x: 2D ndarray, 2D matrix to inverse\n
    max_size: the maximum size to inverse for each block\n
    hermitian: use out = (out + out.conj().T)/2. to keep the hermitian of matrix\n
    inv_func: callable, the function to calculate the inverse of each block matrix\n
    out: None or ndarray, the array to store the inverse of x, shape must be the same as x, if None, build it inside the function\n
    kwargs: for inv_func\n
    return: ndarray, inverse of x\n
    '''
    x = np.asarray(x)
    if out is None:
        out = np.zeros_like(x)
    else:
        assert np.array_equal(out.shape, x.shape)
    if np.size(x) < max_size:
        out[:] = inv_func(x, **kwargs)
        return out
    A, B, C, D = split_matrix(x)
    A_out, B_out, C_out, D_out = split_matrix(out)
    block_inv(A, max_size, hermitian=hermitian, inv_func=inv_func, out=A_out, **kwargs)
    block_inv( D - np.linalg.multi_dot( [C, A_out, B] ), max_size, hermitian=hermitian, inv_func=inv_func, out=D_out, **kwargs )
    B_out[:] = -np.linalg.multi_dot( [A_out, B, D_out] )
    C_out[:] = -np.linalg.multi_dot( [D_out, C, A_out] )
    #if hermitian:
    #    C_out[:] = B_out.T.conj()
    A_out[:] = A_out - np.linalg.multi_dot( [B_out, C, A_out] )
    #print(np.abs(B_out[:].conj().T - C_out[:]).max())
    if hermitian:
        out = (out + out.conj().T)/2.
    return out

if __name__ == '__main__':
    shp = (50001, 50001)
    a = np.random.rand(*shp)# + np.random.rand(*shp)*1.J
    a = np.dot(a.conj().T, a)
    # to keep the block is inversible
    diag = np.diag(a)
    diag0 = np.median(np.abs(diag))
    diag = np.diag(diag*0. + diag0)
    a = a + diag

    #w, v = eigh(a)
    #w = w-w.min()
    #w = w + 10.0
    #plt.matshow(a)
    #a = v@np.diag(w)@v.conj().T
    #plt.matshow(a)
    #plt.show()
    
    #a = np.asarray(a, dtype=np.complex256)
    #ainv0 = inv(a.copy())
    ainv1 = block_inv(a.copy(), max_size=1600)
    ainv2 = np.zeros_like(ainv1)
    #ainv2 = np.zeros([10, 10])
    block_inv(a.copy(), max_size=1600, hermitian=True, out=ainv2)
    print('=====================')
    #test0 = np.dot(a, ainv0)
    test1 = np.dot(a, ainv1)
    test2 = np.dot(a, ainv2)
    #print('===============')
    #print(np.abs(ainv0-ainv1).max())
    #print(np.abs(ainv0-ainv2).max())
    print('===============')
    #print(np.abs(test0 - np.eye(test0.shape[0])).max())
    print(np.abs(test1 - np.eye(test1.shape[0])).max())
    print(np.abs(test2 - np.eye(test2.shape[0])).max())
    #print('===============')
    #print(np.abs(a - a.conj().T).max())
    #print(np.abs(ainv0 - ainv0.conj().T).max())
    #print(np.abs(ainv1 - ainv1.conj().T).max())
    #print(np.abs(ainv2 - ainv2.conj().T).max())
    

2. 点乘

numpy.dot虽然快，但是对内存消耗比较大，特别是对于列数很多的矩阵点乘行数很多的矩阵，对内存不友好，下面函数中的dot在(N, M)点乘(M, L)，且$M\gg N$，$M\gg L$ 时表现更好：

import numpy as np
import matplotlib.pyplot as plt
import healpy as hp
from numba import njit, prange, set_num_threads
from timeit import timeit

@njit(parallel=True, cache=True)
def dot(x, y, out):
    '''
    best solution if x is (N,M), y is (M,L), and M>>L, M>>N
    '''
    n, _ = x.shape
    x = np.ascontiguousarray(x)
    y = np.ascontiguousarray(y.T)
    for ii in prange(n):
        for jj in prange(n):
            out[ii,jj] = np.dot(x[ii,:], y[jj,:])
            #out[ii,jj] = (x[ii,:]*x[jj,:].conj()).sum()

@njit(parallel=True, cache=True)
def doth(x, out):
    '''
    dot(x, x^dagger), more friendly to memory, but slower than dot
    '''
    n, _ = x.shape
    x = np.ascontiguousarray(x)
    for ii in prange(n):
        for jj in prange(ii+1):
            out[ii,jj] = np.dot(x[ii,:], x[jj,:].conj())
            #out[ii,jj] = (x[ii,:]*x[jj,:].conj()).sum()
            if ii!=jj:
                out[jj,ii] = np.conj(out[ii,jj])

@njit(parallel=True, cache=True)
def doth2(x, out):
    '''
    dot(x, x^dagger)
    NOTE: do not use it, VERY SLOW
    '''
    n, _ = x.shape
    for ii in prange(n):
        out[ii,:ii+1] = np.dot(x[ii,:], x[:ii+1,:].T.conj())
        out[:ii+1,ii] = out[ii,:ii+1].conj()
        #for jj in prange(ii+1):
        #    out[ii,:] = np.dot(x[ii,:], x[jj,:].conj())
        #    #out[ii,jj] = (x[ii,:]*x[jj,:].conj()).sum()
        #    if ii!=jj:
        #        out[jj,ii] = np.conj(out[ii,jj])

if __name__ == '__main__':


    # test dot
    set_num_threads(4)

    a = np.random.rand(100, 1000) + np.random.rand(100, 1000) * 1.J

    a1 = np.dot(a, a.conj().T)
    a2 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
    #a3 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
    a4 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
    doth(a, a2)
    #doth2(a, a3)
    dot(a, a.conj().T, a4)
    print(np.abs(a1 - a2).max())
    #print(np.abs(a1 - a3).max())
    #print(np.abs(a2 - a3).max())
    print(np.abs(a2 - a4).max())
    print('======================')

    def c1():
        a1 = np.dot(a.conj().T, a)

    def c2():
        a2 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
        doth(a, a2)

    def c3():
        a2 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
        doth2(a, a2)

    def c4():
        a2 = np.zeros([a.shape[0]]*2, dtype=a.dtype)
        dot(a, a.conj().T, a2)


    print(timeit(c1, number=1000))
    print(timeit(c2, number=1000))
    print(timeit(c4, number=1000))