matrix_sqrt.py

import numpy as np

def matrix_sqrt(A):
    '''Solve for B in A = B*B'''

    # B can be decomposed into V * S * V.I where S is diagonal and V is a column
    # matrix of the eigenvectors (provided it meets certain criteria, but that
    # probably won't happen).
    V = np.matrix(np.linalg.eig(A)[1])  # Returns eigenvalues and eigenvectors
    VI = V.I
    #  Then
    #   (V * S * V.I) * (V * S * V.I) = V * S * S * V.I = A
    #   S * S = V.I * A * V
    # And S * S can be solved by taking the square root of each element, since
    # S is diagonal.  Thus
    SS = VI * A * V
    # Explicitly converting to diagonal vector, then converting back.  Helps with
    # floating point rounding errors.  If it makes you uncomfortable then uncomment
    # assert np.allclose(SS, np.diagflat(np.diag(SS)))
    S = np.matrix(np.diagflat(np.sqrt(np.diag(SS))))
    
    # Now we can calculate B as V * S * V.I
    B = V * S * VI
    
    # And in case you don't believe me, here's another assert that all this worked
    # assert np.allclose(A, B * B)
    return B

A = np.matrix(np.random.random((4, 4)) + 1j * np.random.random((4, 4)))
print(matrix_sqrt(A))