diff --git a/abcmb/background.py b/abcmb/background.py index dc6dcfd..4b38553 100644 --- a/abcmb/background.py +++ b/abcmb/background.py @@ -102,8 +102,6 @@ class Background(eqx.Module): tau_reion : float lna_rec : float rA_rec : float # Comoving angular diameter distance at recombination. - rs_d : float # Sound horizon at baryon decoupling - z_d : float # redshift of baryon decoupling # Transfer related lna_transfer_start : float # Time where transfer functions start integrating. @@ -771,7 +769,7 @@ def _tabulate_kappa_d(self, params): array Tabulated baryon optical depth values (units: dimensionless) """ - integrand = lambda lna, y, args: jnp.float64(-1./self.tau_c(lna, params)/self.aH(lna, params)/(self.R_ratio_lna(lna))) + integrand = lambda lna, y, args: jnp.float64(-1./self.tau_c(lna, params)/self.aH(lna, params)/(self.R_ratio_lna(lna, params))) term = ODETerm(integrand) stepsize_controller = PIDController(pcoeff=0.4, icoeff=0.3, dcoeff=0, rtol=1.e-3, atol=1.e-6) adjoint=ForwardMode() @@ -812,7 +810,7 @@ def _tabulate_rs(self, params): # initial condition assuming cs**2 = 1/3 at early times rs0 = 1./jnp.sqrt(3) / (self.aH( self.lna_tau_tab[0], params )) - integrand = lambda lna, y, args: 1./jnp.sqrt(3*(1+self.R_ratio_lna(lna))) / (self.aH(lna, params)) + integrand = lambda lna, y, args: 1./jnp.sqrt(3*(1+self.R_ratio_lna(lna, params))) / (self.aH(lna, params)) term = ODETerm(integrand) stepsize_controller = PIDController(pcoeff=0.4, icoeff=0.3, dcoeff=0, rtol=1.e-3, atol=1.e-6) adjoint=ForwardMode() @@ -850,7 +848,7 @@ def z_d(self, params): float Decoupling redshift (units: dimensionless) """ - return self.find_z_at_kappad_equals_one(1/jnp.exp(self.lna_tau_tab) - 1, self._tabulate_kappa_d()) + return self.find_z_at_kappad_equals_one(1/jnp.exp(self.lna_tau_tab) - 1, self._tabulate_kappa_d(params)) def rs_d(self, params): """ @@ -868,7 +866,7 @@ def rs_d(self, params): float Sound horizon at decoupling (units: Mpc) """ - return self.interp_rs_at_z(1/jnp.exp(self.lna_tau_tab) - 1, self._tabulate_rs(), self.z_d()) + return self.interp_rs_at_z(1/jnp.exp(self.lna_tau_tab) - 1, self._tabulate_rs(params), self.z_d(params)) class ReionizationModel(eqx.Module): """ diff --git a/abcmb/species.py b/abcmb/species.py index bb2c14e..95a178c 100644 --- a/abcmb/species.py +++ b/abcmb/species.py @@ -43,11 +43,13 @@ class Fluid(eqx.Module): delta_idx : int = eqx.field(default=0) num_moments : int = eqx.field(default=0, static=True) - name : str = eqx.field(default="") + name : str = eqx.field(default="", static=True) is_matter : bool = eqx.field(default=False) # Does the fluid contribute towards matter overdensity today. def __init__(self, delta_idx, specs): self.delta_idx = delta_idx + self.name = "" + self.is_matter = False def rho(self, lna, args): """ @@ -361,6 +363,7 @@ class DarkEnergy(BackgroundFluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "DarkEnergy" def rho(self, lna, args): """ @@ -425,6 +428,8 @@ class ColdDarkMatter(StandardFluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "ColdDarkMatter" + self.is_matter = True def rho(self, lna, args): """ @@ -535,6 +540,7 @@ class MasslessNeutrino(StandardFluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "MasslessNeutrino" self.num_moments = specs["l_max_massless_nu"] + 1 def rho(self, lna, args): @@ -699,6 +705,8 @@ class MassiveNeutrino(Fluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "MassiveNeutrino" + self.is_matter = True self.num_ells_per_bin = specs["l_max_massive_nu"] + 1 self.num_moments = 3 * self.num_ells_per_bin @@ -1003,6 +1011,8 @@ class Baryon(StandardFluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "Baryon" + self.is_matter = True def rho(self, lna, args): """ @@ -1199,6 +1209,7 @@ class Photon(StandardFluid): def __init__(self, delta_idx, specs): super().__init__(delta_idx, specs) + self.name = "Photon" self.num_F_ell_modes = specs["l_max_g"] + 1 self.num_G_ell_modes = specs["l_max_pol_g"] + 1 self.num_moments = self.num_F_ell_modes + self.num_G_ell_modes diff --git a/abcmb/version.py b/abcmb/version.py index 77648b6..984fc57 100644 --- a/abcmb/version.py +++ b/abcmb/version.py @@ -1 +1 @@ -__version__ = "0.2.1" \ No newline at end of file +__version__ = "0.2.2" \ No newline at end of file