Compute square root of a large s

In the studies of large scale inverse problems, we often meet a problem that is how to decompose a sparse positive definite matrix into its square root. In mathematical notations, we need to find a matrix such that

For example , taking SVD according to the definition of $A^{1/2} is too time consuming.

We can employ the following iterative method proposed by Rice (1982)

Here will converge to .

For implementing this iterative algorithm in Python, we give one in the following.

def sqrtMatrix(A, p, max_iter=10, approTol=1e-5):
    nx, ny = A.shape
    Xk = sps.eye(nx)*0.0
    
    for k in range(max_iter):
        temp = Xk
        for j in range(p-1):
            temp = temp@Xk
        Xk = Xk + (1/p)*(A - temp)                     
        ## set small elements to zeros to save memory
        row, col, vals = sps.find(Xk)
        index = (vals < approTol)
        rr, cc = row[index], col[index]
        Xk[rr, cc] = 0
        ## set small elements to zeros to save memory       
        print("Iterate number is ", k)
    
    return Xk

For large scale matrix, during the iteration, there are many small numbers appeared in , which leads to not enough memory problem. So, we add four line of codes to set small elements to zero to save memory. By numerical experiments, the performance is acceptable.