Aug. 18 Friday 2023

These plots includes the three types of computing Ncal

Origonal Method

Ncal = np.matmul(G.T,np.matmul(psr.N,G)) #N_TOA-N_par x N_TOA-N_par
NcalInv = np.linalg.inv(Ncal) #N_TOA-N_par x N_TOA-N_par

Origonal Method with Jax

Ncal = jnp.matmul(G.T,np.matmul(psr.N,G)) #N_TOA-N_par x N_TOA-N_par
NcalInv = jnp.linalg.inv(Ncal) #N_TOA-N_par x N_TOA-N_par

Modified Method with Jax

L = jsc.linalg.cholesky(psr.N)            
A = jnp.matmul(L,G)
del L
Ncal = jnp.matmul(A.T,A)
del A
NcalInv = jnp.linalg.inv(Ncal)

Here is a text file containing information on memory usage per line of code when computing Ncal: NcalInv_mem.txt

Black is the original matrix multiplication,red is original matrix multiplication with Jax, and blue is Cholesky Decomposition with Jax