added script to compute lensed waveforms

scripts/lensed_waveform_exec.py

+#!/home1/uddeepta.deka/softwares/miniconda3/envs/lensing/bin/python3.10
+__author__ =  "U.Deka"
+"""
+This HTCondor compatible code is used to generate microlensed waveforms
+and amplification factor in the time and frequency domain for the 
+charged lens model using contour integration method.
+"""
+
+from sys import argv
+import sys
+sys.path.append("../utils/")
+import numpy as np
+from lenses import ChargedLens
+from generate_ampfac_charged_lens import Ftilde_contour, ReconstructPeak, ampFac, Fw_geom
+import time
+import os
+from scipy.interpolate import interp1d
+from astropy.cosmology import FlatLambdaCDM, Planck18
+import lal
+import pycbc.psd as psd
+from pycbc import detector
+from pycbc.waveform import get_fd_waveform
+
+
+#===========================================================================
+#============================ lensed waveform ==============================
+#===========================================================================
+class MicrolensedWaveform:
+    """
+    Class that generates lensed (and unlensed) waveforms.
+
+    Parmeters:
+    ----------
+    y : float, optional
+        Dimensionless impact parameter [default: 0.1]
+    q : float, optional
+        Dimensionless effective charge [default: 0.0]
+    Ml : float, optional
+        Lens mass in units of solar mass [default: 100.]
+    zl : float, optional
+        Lens redshift [default: 0.5]
+    dect : str, optional
+        Detector config. `Aplus` or `ET` [default: `ET`]
+    f_hi : float, optional
+        Maximum frequency (in Hz) in the detector band [default: 1024]
+    zs : float, optional
+        Lens redshift [default: 2.0]
+    ra : float, optional
+        Right ascension of the source [default: 1.7]
+    dec : float, optional
+        Declination of the source [default: 0.6]
+    pol : float, optional
+        Polarization angle of the source [default: 4.8]
+    m1 : float, optional
+        Mass of component 1 binary in units of solar mass [default: 20]
+    m2 : float, optional
+        Mass of component 2 binary in units of solar mass [default: 20]
+    chi1 : float, optional
+        Spin of component 1 BH along z-direction [default:0.]
+    chi2 : float, optional
+        Spin of component 2 BH along z-direction [default:0.]
+    iota : float, optional
+        Inclination angle in radians [default: PI/8]
+    phi0 : float, optional
+        Coalescence phase in radians [default: PI/4]
+    approx : str, optional
+        Waveform approximant [default: `IMRPhenomXP`]
+    delta_f : float, optional
+        Frequency step (in Hz) used to generate the waveform [default: 1/32]
+    t_max : float, optional
+        Maximum dimensionless time of integration for F(t) [default: 30]
+    dt_fac : float, optional
+        Factor to control the time step for F(t) [default: 7.5]
+    """
+    
+    def __init__(self, y=0.1, q=0., Ml=100., zl=0.5, dect='ET', f_hi=1024,
+                zs=2.0, m1=20, m2=20, chi1=0., chi2=0., iota=np.pi/8., phi0=np.pi/4., 
+                ra=1.7, dec=0.6, pol=4.8, approx='IMRPhenomXP', delta_f=1./32,
+                t_max = 30, dt_fac = 7.5):
+        """
+        Class Initializer
+        """
+        # Cosmology
+        cosmo = FlatLambdaCDM(H0=Planck18.H0.value, Om0=Planck18.Om0)
+        # Lens parameters
+        self.lens = ChargedLens(y1=y, qeff=q)
+        self.Ml = Ml
+        self.zl = zl
+        self.Mlz = self.Ml * (1 + self.zl)
+        # Detector
+        self.dect = dect
+        if self.dect=='Aplus':
+            psd_fname = '../data/PSD/AplusDesign_O5.txt'
+            self.f_lo = 18.
+        elif self.dect=='ET':
+            psd_fname = '../data/PSD/et_d.txt'
+            self.f_lo = 10.
+        else:
+            print("Detector unknown... Reverting to ET...")
+            psd_fname = '../data/PSD/et_d.txt'
+            self.dect = 'ET'
+            self.f_lo = 10.
+        self.f_Sh, self.dect_asd = np.loadtxt(psd_fname, unpack=True)
+        self.f_hi = f_hi
+        self.ra = ra
+        self.dec = dec
+        self.pol = pol
+        # Waveform
+        self.m1 = m1
+        self.m2 = m2
+        self.zs = zs
+        self.ldist_s = cosmo.luminosity_distance(self.zs).value
+        self.chi1 = chi1
+        self.chi2 = chi2
+        self.iota = iota
+        self.phi0 = phi0
+        self.approx = approx
+        self.delta_f = delta_f
+        self.f_final = 0.2 / (self.m1 + self.m2) / lal.MTSUN_SI
+        # Unlensed waveform
+        self.f_UL, self.hp_UL, self.hc_UL = self.get_unlensed_fd_waveform()
+        self.f_plus, self.f_cross, self.hh_UL, self.dect_psd = self.detector_setup()
+        # ampfac params
+        self.t_max = t_max
+        self.dt_fac = dt_fac
+        self.ts = 8 * np.pi * lal.MTSUN_SI * self.Mlz
+        self.t, self.Ft, self.Ft_recons, self.w, self.Ff, self.Ff_recons = self.get_lensed_ampfac()
+        self.Ff_geom = Fw_geom(self.lens, self.w)
+        # Lensed waveform 
+        self.hh_L, self.hh_L_recons = self.get_lensed_fd_waveform()
+
+    def get_unlensed_fd_waveform(self):
+        """
+        Returns unlensed plus- and cross-polarized waveforms
+        """
+        hp, hc = get_fd_waveform(
+            approximant=self.approx, mass1=self.m1, mass2=self.m2,
+            distance=self.ldist_s, spin1z=self.chi1, spin2z=self.chi2,
+            inclination=self.iota, coa_phase=self.phi0, delta_f=self.delta_f,
+            f_lower=self.f_lo, f_final=self.f_final
+        )
+        hp = hp[: int(self.f_final / self.delta_f) + 1]
+        hc = hc[: int(self.f_final / self.delta_f) + 1]
+        f = hp.get_sample_frequencies().data
+        return f, hp, hc
+
+    def detector_setup(self):
+        """
+        Returns frequency domain unlensed waveform and detector PSD
+        """
+        if self.dect == 'Aplus':
+            inst = detector.Detector('H1') # change this to relevant detector
+        else:
+            inst = detector.Detector('V1')
+        f_plus, f_cross = inst.antenna_pattern(self.ra, self.dec, self.pol, 10.)
+        hh_u = f_plus * self.hp_UL + f_cross * self.hc_UL
+        psd_arr = self.dect_asd**2
+        dect_psd = psd.read.from_numpy_arrays(freq_data=self.f_Sh, noise_data=psd_arr,
+                                              length=len(self.hp_UL.data.data), delta_f=self.delta_f,
+                                              low_freq_cutoff=self.f_lo)
+        return f_plus, f_cross, hh_u, dect_psd
+
+
+    def get_lensed_ampfac(self):
+        """
+        Computes and returns the time and frequency domain amplification factor
+        """
+        wmin = self.ts * min(self.f_UL)
+        wmax = self.ts * max(self.f_UL)
+        dt = np.pi / wmax / self.dt_fac        
+
+        t_vals, Ft_vals = Ftilde_contour(self.lens, self.t_max, dt, 
+                                         dx1=0.025, dx2=0.025, verbose=False)
+        
+        # interpolating
+        dt_ = np.pi / max(wmax, 200) / 50
+        num = int(2**np.ceil(np.log2((max(t_vals) - min(t_vals))/dt_)))
+        t_dense = np.linspace(min(t_vals), max(t_vals), num)
+        Ft_interped = np.interp(t_dense, t_vals, Ft_vals)
+        # applying peak correction
+        rp = ReconstructPeak(self.lens, t_dense, Ft_interped, peak_window_frac=0.05)
+        Ft_recons = rp.get_peak_reconstructed_Ft()
+        # doing the FFT
+        w_vals, Ff_vals = ampFac(t_dense, Ft_interped)
+        idxs = (w_vals > wmin) & (w_vals <= wmax)
+        w_vals = w_vals[idxs]
+        Ff_vals = Ff_vals[idxs]
+        _, Ff_recons = ampFac(t_dense, Ft_recons)
+        Ff_recons = Ff_recons[idxs]
+        f_vals = w_vals / self.ts
+        Ff_abs_interp = interp1d(f_vals, np.abs(Ff_vals), fill_value='extrapolate', kind='linear')
+        Ff_phase_interp = interp1d(f_vals, np.angle(Ff_vals), fill_value='extrapolate', kind='linear')
+        Ff_recons_abs_interp = interp1d(f_vals, np.abs(Ff_recons), fill_value='extrapolate', kind='linear')
+        Ff_recons_phase_interp = interp1d(f_vals, np.angle(Ff_recons), fill_value='extrapolate', kind='linear')
+        Ff_interped = Ff_abs_interp(self.f_UL) * np.exp(1j * Ff_phase_interp(self.f_UL))
+        Ff_recons_interped = Ff_recons_abs_interp(self.f_UL) * np.exp(1j * Ff_recons_phase_interp(self.f_UL))
+        return t_dense, Ft_interped, Ft_recons, self.ts*self.f_UL, Ff_interped, Ff_recons_interped
+    
+    def get_lensed_fd_waveform(self):
+        """ Returns the lensed waveform """
+        lensed_hp, lensed_hc = self.Ff * self.hp_UL, self.Ff * self.hc_UL
+        lensed_hp = lensed_hp.cyclic_time_shift(3)
+        lensed_hc = lensed_hc.cyclic_time_shift(3)
+
+        lensed_hp_recons, lensed_hc_recons = self.Ff_recons * self.hp_UL, self.Ff_recons * self.hc_UL
+        lensed_hp_recons = lensed_hp_recons.cyclic_time_shift(3)
+        lensed_hc_recons = lensed_hc_recons.cyclic_time_shift(3)
+        
+        hh_l = self.f_plus * lensed_hp + self.f_cross * lensed_hc
+        hh_l_recons = self.f_plus * lensed_hp_recons + self.f_cross * lensed_hc_recons
+        return hh_l, hh_l_recons
+
+#===========================================================================
+#=============================== Driver code ===============================
+#===========================================================================
+
+def usage():
+    print(
+        "Computes the lensed waveform due to charged lens."
+    )
+    
+    print("\nIt uses the class 'MicrolensedWaveform' to generate results.")
+    print("The following options can be used to overwrite parameters.")
+    print("For example, to change the impact parameter to 0.5, run:")
+    print("python lensed_waveform_condor_exec.py y 0.5")
+    print(MicrolensedWaveform.__doc__)
+    print("\t folder_path: relative path to the output directory from cwd")
+    
+def main():
+    print(" ---------------------------------------------------------------- ")
+    print(" :: lensed_waveform_exec.py :: use flag -h for list of options :: ")
+    print(" ---------------------------------------------------------------- ")
+    kwargs = {"y":0.1, "q":0.0, "Ml":100.0, "zl":0.5, "dect":"ET",
+              "f_hi":1024.0, "zs":2.0, "m1":20.0, "m2":20.0, "chi1":0.0, "chi2":0.0,
+              "iota":np.pi/8., "phi0":np.pi/4., "ra":1.7, "dec":0.6, "pol":4.8,
+              "approx":"IMRPhenomXP", "delta_f":1./32, "t_max":30., "dt_fac":7.5,}
+    folder_path = os.getcwd() + "/lensed_waveform_data/"
+    
+    for i in range(len(sys.argv)):
+        if sys.argv[i] == "-h":
+            usage()
+            return
+        if sys.argv[i] == "folder_path":
+            folder_path = os.getcwd() + "/" + str(sys.argv[i+1]) + "/"
+        if sys.argv[i] in kwargs:
+            if (sys.argv[i]=="dect") or (sys.argv[i]=="approx"):
+                kwargs[sys.argv[i]] = str(sys.argv[i+1])
+            else:
+                kwargs[sys.argv[i]] = float(sys.argv[i+1])
+    
+    tStart = time.perf_counter()
+    mw = MicrolensedWaveform(**kwargs)
+    
+    runtag = 'y_%.3f_q_%.3f_Ml_%.1f_zl_%.1f_dect_%s_fhi_%.1f_zs_%.1f_m1_%.1f_m2_%.1f_approx_%s_deltaf_%.5f_tmax_%.1f_dtfac_%.1f'%(
+        mw.lens.y1, mw.lens.qeff, mw.Ml, mw.zl, mw.dect, mw.f_hi, mw.zs, mw.m1, mw.m2, mw.approx, mw.delta_f, mw.t_max, mw.dt_fac)
+    outdir = folder_path + runtag + "/"
+    isExist = os.path.exists(outdir)
+    if not isExist:
+        os.makedirs(outdir)
+        print("Created new folder to store data.")
+    ## Saving data to files ...
+    Ft_datafile = outdir + 'Ft_data.txt'
+    waveform_datafile = outdir + 'h_lensed_data.npz'
+    header = f"Created using lensed_waveform_exec.py.\nUsage: \nimport numpy as np\ndata = np.loadtxt('{Ft_datafile}')\nt = data[:,0]\nFt = data[:,1]\nFt_recons = data[:,2]\n"
+    with open(Ft_datafile, 'w') as f:
+        np.savetxt(f, np.array([mw.t, mw.Ft, mw.Ft_recons]).T, header=header)
+    np.savez(waveform_datafile, w=mw.w, f=mw.f_UL, Ff=mw.Ff, Ff_recons=mw.Ff_recons, 
+             h_ul=mw.hh_UL, h_l=mw.hh_L, h_l_recons=mw.hh_L_recons)
+    tEnd = time.perf_counter()
+    print("Data stored in : ", outdir)
+    print("Time taken to generate data : %.2f s."%(tEnd-tStart))    
+    return
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file