Anisotropic xi_eta corrections

requirements.txt

@@ -1,6 +1,7 @@
 pytest
 numpy>=2.0
 scipy
+matplotlib
 mcfit
 numexpr
 astropy

zeus21/correlations.py

@@ -770,12 +770,15 @@ def get_all_corrs_IIxIII(self, Cosmo_Parameters, T21_coefficients):
         return 1
 
 
-    def get_xi_Sum_2ExpEta(self, xiEta, etaCoeff1, etaCoeff2):
+    def get_xi_Sum_2ExpEta(self, xiEtaPara, xiEtaPerp, etaCoeff1, etaCoeff2):
         """
         Computes the correlation function of the VCB portion of the SFRD, expressed using sums of two exponentials
         if rho(z1, x1) / rhobar = Ae^-b tilde(eta) + Ce^-d tilde(eta) and rho(z2, x2) / rhobar = Fe^-g tilde(eta) + He^-k tilde(eta)
         Then this computes <rho(z1, x1) * rho(z2, x2)> - <rho(z1, x1)> <rho(z2, x2)>
         Refer to eq. A12 in 2407.18294 for more details
+
+        Computes velocity correlation with xiEta anisotropy, i.e., with xiEtePara and xiEtaPerp separated
+        At call site, if isotropic version is desired, isotropic xiEta is passed in for xiEtaPara and xiEtaPerp is set to None
         """
 
         aa, bb, cc, dd = etaCoeff1
@@ -791,8 +794,18 @@ def get_xi_Sum_2ExpEta(self, xiEta, etaCoeff1, etaCoeff2):
         cfDG = ne.evaluate('cc * ff / normDD / normGG')
         chDK = ne.evaluate('cc * hh / normDD / normKK')
 
-        #The below involves horribly long writing, but breaking this into pieces makes for slightly longer computation time
-        xiNumerator  = ne.evaluate('afBG * (1 / (1 - 6*bb * gg * xiEta / ((1+2*bb)*(1+2*gg)))**(3/2) - 1) + ahBK * (1 / (1 - 6*bb * kk * xiEta / ((1+2*bb)*(1+2*kk)))**(3/2) - 1) + cfDG * (1 / (1 - 6*dd * gg * xiEta / ((1+2*dd)*(1+2*gg)))**(3/2) - 1) + chDK * (1 / (1 - 6*dd * kk * xiEta / ((1+2*dd)*(1+2*kk)))**(3/2) - 1)')
+        if xiEtaPerp is None:
+            xiEta = xiEtaPara
+            #The below involves horribly long writing, but breaking this into pieces makes for slightly longer computation time
+            xiNumerator  = ne.evaluate('afBG * (1 / (1 - 6*bb * gg * xiEta / ((1+2*bb)*(1+2*gg)))**(3/2) - 1) + ahBK * (1 / (1 - 6*bb * kk * xiEta / ((1+2*bb)*(1+2*kk)))**(3/2) - 1) + cfDG * (1 / (1 - 6*dd * gg * xiEta / ((1+2*dd)*(1+2*gg)))**(3/2) - 1) + chDK * (1 / (1 - 6*dd * kk * xiEta / ((1+2*dd)*(1+2*kk)))**(3/2) - 1)')
+        else:
+            xiNumerator  = ne.evaluate(
+                'afBG * (1 / ((1 - 6*bb*gg*xiEtaPara/((1+2*bb)*(1+2*gg))) * (1 - 6*bb*gg*xiEtaPerp/((1+2*bb)*(1+2*gg)))**2)**(1/2) - 1)'
+                ' + ahBK * (1 / ((1 - 6*bb*kk*xiEtaPara/((1+2*bb)*(1+2*kk))) * (1 - 6*bb*kk*xiEtaPerp/((1+2*bb)*(1+2*kk)))**2)**(1/2) - 1)'
+                ' + cfDG * (1 / ((1 - 6*dd*gg*xiEtaPara/((1+2*dd)*(1+2*gg))) * (1 - 6*dd*gg*xiEtaPerp/((1+2*dd)*(1+2*gg)))**2)**(1/2) - 1)'
+                ' + chDK * (1 / ((1 - 6*dd*kk*xiEtaPara/((1+2*dd)*(1+2*kk))) * (1 - 6*dd*kk*xiEtaPerp/((1+2*dd)*(1+2*kk)))**2)**(1/2) - 1)'
+            )
+
         xiDenominator  = ne.evaluate('afBG + ahBK + cfDG + chDK')
 
         xiTotal = ne.evaluate('xiNumerator / xiDenominator')
@@ -807,10 +820,18 @@ def get_all_corrs_III(self, User_Parameters, Cosmo_Parameters, T21_coefficients)
 
         corrdNL = self._corrdNL
 
+        if Cosmo_Parameters.USE_ANISO_XI_ETA and Cosmo_Parameters.USE_RELATIVE_VELOCITIES:
+            corrEtaParaNL = Cosmo_Parameters.xiEtaPara_RR_CF[np.ix_(self._iRnonlinear,self._iRnonlinear)]
+            corrEtaParaNL[0:Cosmo_Parameters.indexminNL,0:Cosmo_Parameters.indexminNL] = corrEtaParaNL[Cosmo_Parameters.indexminNL,Cosmo_Parameters.indexminNL]
+            corrEtaParaNL = corrEtaParaNL.reshape(1, *corrEtaParaNL.shape)
 
-        corrEtaNL = Cosmo_Parameters.xiEta_RR_CF[np.ix_(self._iRnonlinear,self._iRnonlinear)]
-        corrEtaNL[0:Cosmo_Parameters.indexminNL,0:Cosmo_Parameters.indexminNL] = corrEtaNL[Cosmo_Parameters.indexminNL,Cosmo_Parameters.indexminNL]
-        corrEtaNL = corrEtaNL.reshape(1, *corrEtaNL.shape)
+            corrEtaPerpNL = Cosmo_Parameters.xiEtaPerp_RR_CF[np.ix_(self._iRnonlinear,self._iRnonlinear)]
+            corrEtaPerpNL[0:Cosmo_Parameters.indexminNL,0:Cosmo_Parameters.indexminNL] = corrEtaPerpNL[Cosmo_Parameters.indexminNL,Cosmo_Parameters.indexminNL]
+            corrEtaPerpNL = corrEtaPerpNL.reshape(1, *corrEtaPerpNL.shape)
+        else:
+            corrEtaNL = Cosmo_Parameters.xiEta_RR_CF[np.ix_(self._iRnonlinear,self._iRnonlinear)]
+            corrEtaNL[0:Cosmo_Parameters.indexminNL,0:Cosmo_Parameters.indexminNL] = corrEtaNL[Cosmo_Parameters.indexminNL,Cosmo_Parameters.indexminNL]
+            corrEtaNL = corrEtaNL.reshape(1, *corrEtaNL.shape)
 
 
         _coeffTx_units = T21_coefficients.coeff_Gammah_Tx_III #includes -10^40 erg/s/SFR normalizaiton and erg/K conversion factor
@@ -836,7 +857,10 @@ def get_all_corrs_III(self, User_Parameters, Cosmo_Parameters, T21_coefficients)
         expGammaCorr = ne.evaluate('exp(gammaCorrdNL) - 1') # equivalent to np.exp(gammaTimesCorrdNL)-1.0
 
         if Cosmo_Parameters.USE_RELATIVE_VELOCITIES == True:
-            etaCorr_xa = self.get_xi_Sum_2ExpEta(corrEtaNL, vcbCoeffsR1, vcbCoeffsR2)
+            if Cosmo_Parameters.USE_ANISO_XI_ETA:
+                etaCorr_xa = self.get_xi_Sum_2ExpEta(corrEtaParaNL, corrEtaPerpNL, vcbCoeffsR1, vcbCoeffsR2)
+            else:
+                etaCorr_xa = self.get_xi_Sum_2ExpEta(corrEtaNL, None, vcbCoeffsR1, vcbCoeffsR2)
             totalCorr = ne.evaluate('expGammaCorr * etaCorr_xa + expGammaCorr + etaCorr_xa - gammaCorrdNL') ###TO DO (linearized VCB flucts): - etaCorr_xa_lin #note that the Taylor expansion of the cross-term is 0 to linear order
         else:
             totalCorr = ne.evaluate('expGammaCorr - gammaCorrdNL') ###TO DO (linearized VCB flucts): - etaCorr_xa_lin #note that the Taylor expansion of the cross-term is 0 to linear order
@@ -867,7 +891,11 @@ def get_all_corrs_III(self, User_Parameters, Cosmo_Parameters, T21_coefficients)
         coeffsXaTxALL = coeffzp2Tx * coeffmatrixxaTx
 
         corrdNLBIG = corrdNL[:,:, np.newaxis, :, :]
-        corrEtaNLBIG = corrEtaNL[:,:, np.newaxis, :, :]
+        if Cosmo_Parameters.USE_ANISO_XI_ETA and Cosmo_Parameters.USE_RELATIVE_VELOCITIES:
+            corrEtaParaNLBIG  = corrEtaParaNL[:,:, np.newaxis, :, :]
+            corrEtaPerpNLBIG = corrEtaPerpNL[:,:, np.newaxis, :, :]
+        else:
+            corrEtaNLBIG = corrEtaNL[:,:, np.newaxis, :, :]
 
         vcbCoeffsR1 = vcbCoeffsR1[:,:,:,:,:]
         vcbCoeffsR2 = np.transpose(vcbCoeffsR1, (0,3,4,1,2))
@@ -878,13 +906,20 @@ def get_all_corrs_III(self, User_Parameters, Cosmo_Parameters, T21_coefficients)
 
         for ir in range(len(Cosmo_Parameters._Rtabsmoo)):
             corrdNL = corrdNLBIG[:,:,:,:,ir]
-            corrEtaNL = corrEtaNLBIG[:,:,:,:,ir]
+            if Cosmo_Parameters.USE_ANISO_XI_ETA and Cosmo_Parameters.USE_RELATIVE_VELOCITIES:
+                corrEtaParaNL  = corrEtaParaNLBIG[:,:,:,:,ir]
+                corrEtaPerpNL = corrEtaPerpNLBIG[:,:,:,:,ir]
+            else:
+                corrEtaNL = corrEtaNLBIG[:,:,:,:,ir]
 
             gammaCorrdNL = ne.evaluate('gammamatrixR1R2 * corrdNL')
             expGammaCorrdNL = ne.evaluate('exp(gammaCorrdNL) - 1')
 
             if Cosmo_Parameters.USE_RELATIVE_VELOCITIES == True:
-                etaCorr_Tx = self.get_xi_Sum_2ExpEta(corrEtaNL, vcbCoeffsR1, vcbCoeffsR2)
+                if Cosmo_Parameters.USE_ANISO_XI_ETA:
+                    etaCorr_Tx = self.get_xi_Sum_2ExpEta(corrEtaParaNL, corrEtaPerpNL, vcbCoeffsR1, vcbCoeffsR2)
+                else:
+                    etaCorr_Tx = self.get_xi_Sum_2ExpEta(corrEtaNL, None, vcbCoeffsR1, vcbCoeffsR2)
                 totalCorr = ne.evaluate('expGammaCorrdNL * etaCorr_Tx + expGammaCorrdNL + etaCorr_Tx - gammaCorrdNL') ###TO DO (linearized VCB flucts): - etaCorr_xa_lin #note that the Taylor expansion of the cross-term is 0 to linear order
             else:
                 totalCorr = ne.evaluate('expGammaCorrdNL - gammaCorrdNL') ###TO DO (linearized VCB flucts): - etaCorr_xa_lin #note that the Taylor expansion of the cross-term is 0 to linear order
@@ -943,4 +978,4 @@ def get_Pk_from_xi(self, rsinput, xiinput):
 
         kPf, Pf = mcfit.xi2P(rsinput, l=0, lowring=True)(xiinput, extrap=False)
 
-        return kPf, Pf
+        return kPf, Pf

zeus21/cosmology.py

@@ -14,20 +14,38 @@
 
 import numpy as np
 from scipy.interpolate import RegularGridInterpolator
+from scipy.interpolate import interp1d
 
 from . import constants
 from .inputs import Cosmo_Parameters
 
-def cosmo_wrapper(User_Parameters):
+from dataclasses import fields
+
+Cosmo_Param_Fields = tuple(
+    f.name for f in fields(Cosmo_Parameters)
+    if f.init and f.name != "UserParams"
+)
+
+def cosmo_wrapper(User_Parameters, **kwargs):
     """
-    Wrapper function for all the cosmology. It takes Cosmo_Parameters_Input and returns:
-    Cosmo_Parameters, Class_Cosmo, Correlations, HMF_interpolator
+    Wrapper function for all the cosmology. 
+    It takes User_Parameters plus keyword arguments and returns:
+    Cosmo_Parameters, Class_Cosmo, HMF_interpolator
     """
-
-    CosmoParams = Cosmo_Parameters(User_Parameters) 
+    cosmo_kwargs = {
+        name: value
+        for name, value in kwargs.items()
+        if name in Cosmo_Param_Fields
+    }
+    for name in kwargs:
+        if name not in Cosmo_Param_Fields:
+            print(f"cosmo_wrapper: ignoring unsupported kwarg '{name}'")
+
+    CosmoParams = Cosmo_Parameters(UserParams=User_Parameters, **cosmo_kwargs)
+    ClassCosmo = CosmoParams.ClassCosmo
     HMFintclass = HMF_interpolator(User_Parameters,CosmoParams)
 
-    return CosmoParams, HMFintclass
+    return CosmoParams, ClassCosmo, HMFintclass
 
 
 

zeus21/inputs.py

@@ -225,6 +225,7 @@ class Cosmo_Parameters:
     # Flags
     Flag_emulate_21cmfast: bool = False
     USE_RELATIVE_VELOCITIES: bool = False
+    USE_ANISO_XI_ETA: bool = False
     HMF_CHOICE: str = "ST"
 
 
@@ -292,6 +293,7 @@ def __post_init__(self, UserParams):
         schema = {
             "Flag_emulate_21cmfast": (bool, None),
             "USE_RELATIVE_VELOCITIES": (bool, None),
+            "USE_ANISO_XI_ETA": (bool, None),
             "HMF_CHOICE": (str, {'ST','Yung'}),
         }
         validate_fields(self, schema)
@@ -450,10 +452,19 @@ def runclass(self):
             psi0 = 1 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapezoid(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j0bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
             psi2 = -2 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapezoid(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j2bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
 
-            k_eta, P_eta = mcfit.xi2P(rVelIntp, l=0, lowring = True)((6 * psi0**2 + 3 * psi2**2), extrap = False)
+            if self.USE_ANISO_XI_ETA:
+                k_eta_para,  P_eta_para  = mcfit.xi2P(rVelIntp, l=0, lowring = True)(6 * (psi0 + psi2)**2, extrap = False)
+                k_eta_perp, P_eta_perp = mcfit.xi2P(rVelIntp, l=0, lowring = True)(6 * (psi0 - psi2/2.0)**2, extrap = False)
+
+                ClassCosmo.pars['k_eta_para']  = k_eta_para[P_eta_para > 0]
+                ClassCosmo.pars['P_eta_para']  = P_eta_para[P_eta_para > 0]
+                ClassCosmo.pars['k_eta_perp'] = k_eta_perp[P_eta_perp > 0]
+                ClassCosmo.pars['P_eta_perp'] = P_eta_perp[P_eta_perp > 0]
+            else:
+                k_eta, P_eta = mcfit.xi2P(rVelIntp, l=0, lowring = True)((6 * psi0**2 + 3 * psi2**2), extrap = False)
 
-            ClassCosmo.pars['k_eta'] = k_eta[P_eta > 0]
-            ClassCosmo.pars['P_eta'] = P_eta[P_eta > 0]
+                ClassCosmo.pars['k_eta'] = k_eta[P_eta > 0]
+                ClassCosmo.pars['P_eta'] = P_eta[P_eta > 0]
 
     #        print("HAC: Finished running CLASS a second time to get velocity transfer functions")
 
@@ -483,9 +494,17 @@ def run_correlations(self):
 
         ###HAC: Interpolated object for eta power spectrum
         if self.USE_RELATIVE_VELOCITIES == True:
-            P_eta_interp = interp1d(self.ClassCosmo.pars['k_eta'], self.ClassCosmo.pars['P_eta'], bounds_error = False, fill_value = 0)
-            self._PkEtaCF = P_eta_interp(self._klistCF)
-            self.xiEta_RR_CF = self.get_xi_R1R2(field = 'vcb')
+            if self.USE_ANISO_XI_ETA:
+                P_eta_para_interp = interp1d(self.ClassCosmo.pars['k_eta_para'], self.ClassCosmo.pars['P_eta_para'], bounds_error = False, fill_value = 0)
+                P_eta_perp_interp = interp1d(self.ClassCosmo.pars['k_eta_perp'], self.ClassCosmo.pars['P_eta_perp'], bounds_error = False, fill_value = 0)
+                self._PkEtaParaCF = P_eta_para_interp(self._klistCF)
+                self._PkEtaPerpCF = P_eta_perp_interp(self._klistCF)
+                self.xiEtaPara_RR_CF = self.get_xi_R1R2(field = 'vcb_para')
+                self.xiEtaPerp_RR_CF = self.get_xi_R1R2(field = 'vcb_perp')
+            else:
+                P_eta_interp = interp1d(self.ClassCosmo.pars['k_eta'], self.ClassCosmo.pars['P_eta'], bounds_error = False, fill_value = 0)
+                self._PkEtaCF = P_eta_interp(self._klistCF)
+                self.xiEta_RR_CF = self.get_xi_R1R2(field = 'vcb')
         else:
             self._PkEtaCF = np.zeros_like(self._PklinCF)
             self.xiEta_RR_CF = np.zeros_like(self.xi_RR_CF)
@@ -503,7 +522,13 @@ def get_xi_R1R2 (self, field = None):
             _PkRR = np.array([[self._PklinCF]]) * windowR1 * windowR2
         elif field == 'vcb':
             _PkRR = np.array([[self._PkEtaCF]]) * windowR1 * windowR2
+        elif field == 'vcb_para':
+            _PkRR = np.array([[self._PkEtaParaCF]]) * windowR1 * windowR2
+        elif field == 'vcb_perp':
+            _PkRR = np.array([[self._PkEtaPerpCF]]) * windowR1 * windowR2
         else:
+            if self.USE_ANISO_XI_ETA:
+                raise ValueError('field has to be either delta, vcb_para or vcb_perp in get_xi_R1R2')
             raise ValueError('field has to be either delta or vcb in get_xi_R1R2')
 
         self.rlist_CF, xi_RR_CF = self._xif(_PkRR, extrap = False)