Cholesky decomposition from a random input matrix

cholesky.py

import numpy as np


def clean_cholesky(chi):
    """Converts a random complex-valued matrix which has a shape 
    (N, N) to a Hermitian matrix.

    Args:
        chi (array (complex)): A matrix of shape (N, N).

    Returns
        rho (array (complex)): A Hermitian matrix using the Cholesky decomposition.
    """
    diagonal_elements = np.diag(np.diag(chi).real)
    indices = np.tril_indices(chi.shape[0], 0)
    chi[indices] = 0.
    T = chi + diagonal_elements
    return np.conj(T.T)@T


def test_clean_cholesky():
    """Tests the clean Cholesky function"""
    N = 4
    rand = np.random.random(size=(N, N)) + 1j*np.random.random(size=(N, N))
    rho = clean_cholesky(rand)
    
    eigvals, eigvecs = np.linalg.eig(rho)
    # test if all eigenvalues are greater than 0
    assert(np.all(eigvals) >= 0.)

test_clean_cholesky()

N = 4
rand = np.random.random(size=(N, N)) + 1j*np.random.random(size=(N, N))
rho = clean_cholesky(rand)
print(rho)