TypeError: 'JVPTracer' object is not callable

[harshalc03]

I'm not really familiar with jax and not sure how is the function chi1z_intermsof_Mcqchiachis causing the issue. Any help would be appreciated. The code is as follows:

import jax
import numpy as np
import pandas as pd
import gwfast.gwfastUtils as util
import gwfast.fisherTools as tool
import gwfast.waveforms as wave

fish_m1m2 = np.load('/home/harshalc03/Desktop/first_10fisherm1m2.npy')
params = wave.IMRPhenomHM().ParNums
fish_final = np.empty(fish_m1m2.shape)

M_chirp = event_10['Mc']
sym_mass_ratio = event_10['eta']
spin1z = event_10['chi1z']
spin2z = event_10['chi2z']
mass_1, mass_2 = util.m1m2_from_Mceta(M_chirp, sym_mass_ratio)
mass_ratio = mass_1/mass_2
chi_sym = (mass_1*spin1z+mass_2*spin2z)/(mass_1+mass_2)
chi_antisym = (mass_1*spin1z-mass_2*spin2z)/(mass_1+mass_2)

def m1_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a):
    return (q**(2/5))*((q+1)**(1/5))*Mc

def m2_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a):
    return ((q+1)**(1/5))*Mc/(q**(3/5))

def chi1z_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a):
    return ((m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) + m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))(chi_s + chi_a))/(2*m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))

def chi2z_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a):
    return ((m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) + m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))(chi_s - chi_a))/(2*m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))

m1_deri = np.array([])
m2_deri = np.array([])
chi1z_deri = np.array([])
chi2z_deri = np.array([])
for i in range(4):
    m1_deri = np.append(m1_deri, jax.grad(m1_intermsof_Mcqchiachis, argnums=i))
    m2_deri = np.append(m2_deri, jax.grad(m2_intermsof_Mcqchiachis, argnums=i))
    chi1z_deri = np.append(chi1z_deri, jax.grad(chi1z_intermsof_Mcqchiachis, argnums=i))
    chi2z_deri = np.append(chi2z_deri, jax.grad(chi2z_intermsof_Mcqchiachis, argnums=i))


def Jac_m1m2chi1zchi2z_Mcqchiachis(Mc, q, chi_s, chi_a):
    Jac_matrix = np.array([m1_deri, m2_deri, chi1z_deri, chi2z_deri])
    final_jac = np.empty((4,4))
    for i in range(4):
        for j in range(4):
            final_jac[i,j] = Jac_matrix[i,j](Mc, q, chi_s, chi_a)
    
    return final_jac
    
    
for i in range(fish_m1m2.shape[2]):
    
    event = {}
    for par in params:
        event[par] = event_10[par][i]

    # Creating the Jacobian matrix
    rotMatrix = np.identity(11)
    rotMatrix[np.ix_([0,1,9,10], [0,1,9,10])] = Jac_m1m2chi1zchi2z_Mcqchiachis(M_chirp[i], mass_ratio[i], chi_sym[i], chi_antisym[i])
    
    # Getting the new covariance matrix
    fish_final[:,:,i] = rotMatrix.T@fish_m1m2[:,:,i]@rotMatrix

cov_Mc_q_chis_chia, inv_cov_Mc_q_chis_chia = tool.CovMatr(fish_final)

The error is as follows:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[31], line 54
     52 # Creating the Jacobian matrix
     53 rotMatrix = np.identity(11)
---> 54 rotMatrix[np.ix_([0,1,9,10], [0,1,9,10])] = Jac_m1m2chi1zchi2z_Mcqchiachis(M_chirp[i], mass_ratio[i], chi_sym[i], chi_antisym[i])
     56 # Getting the new covariance matrix
     57 fish_final[:,:,i] = rotMatrix.T@fish_m1m2[:,:,i]@rotMatrix

Cell In[31], line 41, in Jac_m1m2chi1zchi2z_Mcqchiachis(Mc, q, chi_s, chi_a)
     39 for i in range(4):
     40     for j in range(4):
---> 41         final_jac[i,j] = Jac_matrix[i,j](Mc, q, chi_s, chi_a)
     43 return final_jac

    [... skipping hidden 10 frame]

Cell In[31], line 20, in chi1z_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a)
     19 def chi1z_intermsof_Mcqchiachis(Mc, q, chi_s, chi_a):
---> 20     return ((m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) + m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))(chi_s + chi_a))/(2*m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))

TypeError: 'JVPTracer' object is not callable

[jakevdp]

Thanks for the clear reproduction!

A JVPTracer is basically an array within the context of automatic differentiation. Your code is attempting to call this array as if it were a function, which leads to an error because arrays are not callable. In this line:

return ((m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) + m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))(chi_s + chi_a))/(2*m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))

The result of (m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) + m2_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a) is an array; let's call it arr for simplicity. Then your expression looks like this:

return (arr(chi_s + chi_a))/(2*m1_intermsof_Mcqchiachis(Mc,q, chi_s, chi_a))

The problem is that arr(chi_s + chi_a) is not a valid operation, because arrays are not callable functions.

[harshalc03]

Ahh shoot, it was a simple error which I missed. I had forgot to add the * sign for multiplication. Thanks a lot @jakevdp