ABCMB 0.1.2__tar.gz → 0.2.2__tar.gz

{abcmb-0.1.2 → abcmb-0.2.2}/ABCMB.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ABCMB
-Version: 0.1.2
+Version: 0.2.2
 Summary: A fast, differentiable, and extensible CMB code
 Home-page: https://github.com/TonyZhou729/ABCMB
 Author: Zilu Zhou, Cara Giovanetti, Hongwan Liu
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: matplotlib
@@ -27,6 +28,8 @@ Requires-Dist: equinox==0.13.2
 Requires-Dist: interpax
 Requires-Dist: jax==0.8.1
 Requires-Dist: pytest
+Requires-Dist: optimistix==0.0.11
+Dynamic: license-file
 
 <h1 align="center">
 ABCMB<!-- omit from toc -->
@@ -34,6 +37,7 @@ ABCMB<!-- omit from toc -->
 <h4 align="center">
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-red.svg)](https://opensource.org/licenses/MIT)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.15104%20-green.svg)](https://arxiv.org/abs/2602.15104)
 [![Run Tests](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml/badge.svg)](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml)
 <!--[![arXiv](https://img.shields.io/badge/arXiv-2408.14538%20-green.svg)](https://arxiv.org/abs/2408.14538) -->
 
@@ -42,14 +46,19 @@ ABCMB<!-- omit from toc -->
 Autodifferentiable Boltzmann solver for the CMB (ABCMB) is a Python+JAX package for differentiable computation of the Cosmic Microwave Background.  ABCMB is **complete to linear order** in $\Lambda\rm{CDM}$ cosmology.  It computes the matter and CMB power spectra and includes effects like lensing, massive neutrinos, and a state-of-the-art treatment of the physics of recombination through the companion code [HyRex](https://github.com/TonyZhou729/HyRex).
 
 ## Installation
-We recommend installing ABCMB in a clean conda environment.  After downloading and unpacking the code, in the code directory run 
+ABCMB is pip installable!  Just run
 ```
-conda create -n ABCMB
-conda activate ABCMB
-pip install -U -r requirements.txt
+pip install ABCMB
+```
+We recommend always doing so in a conda environment, preferably even a clean one.
 
+If you'd like to clone the repo instead, after cloning you can run
+```
+pip install .
 ```
-optionally specifying your preferred python version after the environment name.  Note that this will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
+from the code directory. 
+
+Note that both methods of installing will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
 
 ## Examples
 We have included several pedagogical jupyter notebooks to walk you through how to get started with ABCMB in our [example_notebooks](https://github.com/TonyZhou729/ABCMB/tree/main/example_notebooks) folder.  We suggest you start with [ABCMB_basics](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_basics.ipynb) to get a sense of how to run the code.  If you'd like to add new physics to ABCMB, check out [ABCMB_Fluids](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_Fluids.ipynb).  If you'd like to run ABCMB with the Big Bang Nucleosynthesis (BBN) code [LINX](https://github.com/cgiovanetti/LINX/tree/main) to do BBN+CMB joint analyses, check out [ABCMB_with_LINX](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_with_LINX.ipynb).
@@ -57,5 +66,20 @@ We have included several pedagogical jupyter notebooks to walk you through how t
 ## Issues
 Please feel free to open an issue if something is amiss in ABCMB!
 
+## Citation
+
+If you use ABCMB to publish scientific research, we suggest you cite
+```
+@misc{abcmb,
+      title={{ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum}}, 
+      author={Zilu Zhou and Cara Giovanetti and Hongwan Liu},
+      year={2026},
+      eprint={2602.15104},
+      archivePrefix={arXiv},
+      primaryClass={astro-ph.CO},
+      url={https://arxiv.org/abs/2602.15104}, 
+}
+```
+
 
 

{abcmb-0.1.2 → abcmb-0.2.2}/ABCMB.egg-info/SOURCES.txt

@@ -1,3 +1,4 @@
+LICENSE
 README.md
 pyproject.toml
 setup.cfg

{abcmb-0.1.2 → abcmb-0.2.2}/ABCMB.egg-info/requires.txt

@@ -6,3 +6,4 @@ equinox==0.13.2
 interpax
 jax==0.8.1
 pytest
+optimistix==0.0.11

abcmb-0.2.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Zilu Zhou, Cara Giovanetti, Hongwan Liu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{abcmb-0.1.2 → abcmb-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ABCMB
-Version: 0.1.2
+Version: 0.2.2
 Summary: A fast, differentiable, and extensible CMB code
 Home-page: https://github.com/TonyZhou729/ABCMB
 Author: Zilu Zhou, Cara Giovanetti, Hongwan Liu
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: matplotlib
@@ -27,6 +28,8 @@ Requires-Dist: equinox==0.13.2
 Requires-Dist: interpax
 Requires-Dist: jax==0.8.1
 Requires-Dist: pytest
+Requires-Dist: optimistix==0.0.11
+Dynamic: license-file
 
 <h1 align="center">
 ABCMB<!-- omit from toc -->
@@ -34,6 +37,7 @@ ABCMB<!-- omit from toc -->
 <h4 align="center">
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-red.svg)](https://opensource.org/licenses/MIT)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.15104%20-green.svg)](https://arxiv.org/abs/2602.15104)
 [![Run Tests](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml/badge.svg)](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml)
 <!--[![arXiv](https://img.shields.io/badge/arXiv-2408.14538%20-green.svg)](https://arxiv.org/abs/2408.14538) -->
 
@@ -42,14 +46,19 @@ ABCMB<!-- omit from toc -->
 Autodifferentiable Boltzmann solver for the CMB (ABCMB) is a Python+JAX package for differentiable computation of the Cosmic Microwave Background.  ABCMB is **complete to linear order** in $\Lambda\rm{CDM}$ cosmology.  It computes the matter and CMB power spectra and includes effects like lensing, massive neutrinos, and a state-of-the-art treatment of the physics of recombination through the companion code [HyRex](https://github.com/TonyZhou729/HyRex).
 
 ## Installation
-We recommend installing ABCMB in a clean conda environment.  After downloading and unpacking the code, in the code directory run 
+ABCMB is pip installable!  Just run
 ```
-conda create -n ABCMB
-conda activate ABCMB
-pip install -U -r requirements.txt
+pip install ABCMB
+```
+We recommend always doing so in a conda environment, preferably even a clean one.
 
+If you'd like to clone the repo instead, after cloning you can run
+```
+pip install .
 ```
-optionally specifying your preferred python version after the environment name.  Note that this will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
+from the code directory. 
+
+Note that both methods of installing will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
 
 ## Examples
 We have included several pedagogical jupyter notebooks to walk you through how to get started with ABCMB in our [example_notebooks](https://github.com/TonyZhou729/ABCMB/tree/main/example_notebooks) folder.  We suggest you start with [ABCMB_basics](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_basics.ipynb) to get a sense of how to run the code.  If you'd like to add new physics to ABCMB, check out [ABCMB_Fluids](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_Fluids.ipynb).  If you'd like to run ABCMB with the Big Bang Nucleosynthesis (BBN) code [LINX](https://github.com/cgiovanetti/LINX/tree/main) to do BBN+CMB joint analyses, check out [ABCMB_with_LINX](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_with_LINX.ipynb).
@@ -57,5 +66,20 @@ We have included several pedagogical jupyter notebooks to walk you through how t
 ## Issues
 Please feel free to open an issue if something is amiss in ABCMB!
 
+## Citation
+
+If you use ABCMB to publish scientific research, we suggest you cite
+```
+@misc{abcmb,
+      title={{ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum}}, 
+      author={Zilu Zhou and Cara Giovanetti and Hongwan Liu},
+      year={2026},
+      eprint={2602.15104},
+      archivePrefix={arXiv},
+      primaryClass={astro-ph.CO},
+      url={https://arxiv.org/abs/2602.15104}, 
+}
+```
+
 
 

{abcmb-0.1.2 → abcmb-0.2.2}/README.md

@@ -4,6 +4,7 @@ ABCMB<!-- omit from toc -->
 <h4 align="center">
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-red.svg)](https://opensource.org/licenses/MIT)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.15104%20-green.svg)](https://arxiv.org/abs/2602.15104)
 [![Run Tests](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml/badge.svg)](https://github.com/TonyZhou729/ABCMB/actions/workflows/accuracy.yml)
 <!--[![arXiv](https://img.shields.io/badge/arXiv-2408.14538%20-green.svg)](https://arxiv.org/abs/2408.14538) -->
 
@@ -12,14 +13,19 @@ ABCMB<!-- omit from toc -->
 Autodifferentiable Boltzmann solver for the CMB (ABCMB) is a Python+JAX package for differentiable computation of the Cosmic Microwave Background.  ABCMB is **complete to linear order** in $\Lambda\rm{CDM}$ cosmology.  It computes the matter and CMB power spectra and includes effects like lensing, massive neutrinos, and a state-of-the-art treatment of the physics of recombination through the companion code [HyRex](https://github.com/TonyZhou729/HyRex).
 
 ## Installation
-We recommend installing ABCMB in a clean conda environment.  After downloading and unpacking the code, in the code directory run 
+ABCMB is pip installable!  Just run
 ```
-conda create -n ABCMB
-conda activate ABCMB
-pip install -U -r requirements.txt
+pip install ABCMB
+```
+We recommend always doing so in a conda environment, preferably even a clean one.
 
+If you'd like to clone the repo instead, after cloning you can run
+```
+pip install .
 ```
-optionally specifying your preferred python version after the environment name.  Note that this will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
+from the code directory. 
+
+Note that both methods of installing will automatically attempt to install JAX for CPU; to install for GPU, refer to the [JAX documentation](https://docs.jax.dev/en/latest/installation.html) for a quick JAX installation guide.
 
 ## Examples
 We have included several pedagogical jupyter notebooks to walk you through how to get started with ABCMB in our [example_notebooks](https://github.com/TonyZhou729/ABCMB/tree/main/example_notebooks) folder.  We suggest you start with [ABCMB_basics](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_basics.ipynb) to get a sense of how to run the code.  If you'd like to add new physics to ABCMB, check out [ABCMB_Fluids](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_Fluids.ipynb).  If you'd like to run ABCMB with the Big Bang Nucleosynthesis (BBN) code [LINX](https://github.com/cgiovanetti/LINX/tree/main) to do BBN+CMB joint analyses, check out [ABCMB_with_LINX](https://github.com/TonyZhou729/ABCMB/blob/main/example_notebooks/ABCMB_with_LINX.ipynb).
@@ -27,5 +33,20 @@ We have included several pedagogical jupyter notebooks to walk you through how t
 ## Issues
 Please feel free to open an issue if something is amiss in ABCMB!
 
+## Citation
+
+If you use ABCMB to publish scientific research, we suggest you cite
+```
+@misc{abcmb,
+      title={{ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum}}, 
+      author={Zilu Zhou and Cara Giovanetti and Hongwan Liu},
+      year={2026},
+      eprint={2602.15104},
+      archivePrefix={arXiv},
+      primaryClass={astro-ph.CO},
+      url={https://arxiv.org/abs/2602.15104}, 
+}
+```
+
 
 

{abcmb-0.1.2 → abcmb-0.2.2}/abcmb/ABCMBTools.py

@@ -27,9 +27,9 @@ def wigner_d_matrix(mu, ells, m, n):
     ells : array
         Multipole values [m, m+1, m+2, ..., ellmax]
     m : int
-        First index (must be positive and >= |n|)
+        First index (must be positive and >= abs(n))
     n : int
-        Second index (must satisfy |n| <= m)
+        Second index (must satisfy abs(n) <= m)
 
     Returns:
     --------
@@ -102,14 +102,14 @@ def d1n(mu, ells, n):
     ells : array
         Multipole values
     n : int
-        Second index (|n| <= 1)
+        Second index (abs(n) <= 1)
 
     Returns:
     --------
     array
         d^ell_{1n} elements
     """
-    # Wigner matrices where m=1, and |n|<=m.
+    # Wigner matrices where m=1, and abs(n)<=m.
     ells_patched = jnp.concatenate((jnp.array([1]), ells))
     res = wigner_d_matrix(mu, ells_patched, 1, n)
     return res[:, 1:]
@@ -125,14 +125,14 @@ def d2n(mu, ells, n):
     ells : array
         Multipole values
     n : int
-        Second index (|n| <= 2)
+        Second index (abs(n) <= 2)
 
     Returns:
     --------
     array
         d^ell_{2n} elements
     """
-    # Wigner matrices where m=2, and |n|<=m.
+    # Wigner matrices where m=2, and abs(n)<=m.
     res = wigner_d_matrix(mu, ells, 2, n)
     return res
 
@@ -147,14 +147,14 @@ def d3n(mu, ells, n):
     ells : array
         Multipole values
     n : int
-        Second index (|n| <= 3)
+        Second index (abs(n) <= 3)
 
     Returns:
     --------
     array
         d^ell_{3n} elements, zero-padded for ell < 3
     """
-    # Wigner matrices where m=3, and |n|<=m.
+    # Wigner matrices where m=3, and abs(n)<=m.
     ells_sliced = ells[1:] # Compute starting at ell=3
     res = wigner_d_matrix(mu, ells_sliced, 3, n)
     res_patched = jnp.concatenate((jnp.zeros((mu.size, 1)), res), axis=1) # Pad zeros for ell<3.
@@ -171,14 +171,14 @@ def d4n(mu, ells, n):
     ells : array
         Multipole values
     n : int
-        Second index (|n| <= 4)
+        Second index (abs(n) <= 4)
 
     Returns:
     --------
     array
         d^ell_{4n} elements, zero-padded for ell < 4
     """
-    # Wigner matrices where m=4, and |n|<=m.
+    # Wigner matrices where m=4, and abs(n)<=m.
     ells_sliced = ells[2:] # Compute starting at ell=4
     res = wigner_d_matrix(mu, ells_sliced, 4, n)
     res_patched = jnp.concatenate((jnp.zeros((mu.size, 2)), res), axis=1) # Pad zeros for ell<4.

abcmb-0.2.2/abcmb/__init__.py

@@ -0,0 +1,7 @@
+"""
+ABCMB.
+
+A fully differentiable Boltzmann solver for the CMB.
+"""
+
+__author__ = 'Zilu Zhou, Cara Giovanetti, and Hongwan Liu'

{abcmb-0.1.2 → abcmb-0.2.2}/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.
@@ -374,7 +372,7 @@ class Background(eqx.Module):
             saveat=saveat,
             stepsize_controller=controller,
             args=params,
-            adjoint=adjoint
+            adjoint=adjoint,
         )
 
         # Numerical jitter causes this interpolation to go out of bounds on
@@ -614,14 +612,14 @@ class Background(eqx.Module):
         adjoint=ForwardMode()
         sol = diffeqsolve(
             term,
-            solver=Kvaerno5(),            # Higher order integrator for more accuracy
+            solver=Kvaerno5(),           
             stepsize_controller=stepsize_controller,
-            t0=0.,                 # Initial x value (~0 in this case)
-            t1=-10.,                  # Final x value (smallest x value)
-            dt0=-1.e-3,                  # Initial step size
+            t0=0.,                
+            t1=-10.,                  
+            dt0=-1.e-3,                 
             max_steps=2048,
-            y0=0.0,                     # Initial value tau(x=0) = 0
-            saveat=SaveAt(dense=True), # Save at all points in x, reverse order since integrating backwards
+            y0=0.0,                     
+            saveat=SaveAt(dense=True),
             adjoint=adjoint
         )
         return sol
@@ -771,7 +769,7 @@ class Background(eqx.Module):
         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 @@ class Background(eqx.Module):
          # 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 @@ class Background(eqx.Module):
         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,16 +866,19 @@ class Background(eqx.Module):
         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):
     """
     Object for computing the reionization correction to the free electron fraction.
     Provides the base methods 
-        xe_reion : calculates the tanh electron fraction correction at redshifts lna, given z_reion and params
-        tau_reion_fn : calculates the optical depth to reionization.
+
+    xe_reion : calculates the tanh electron fraction correction at redshifts lna, given z_reion and params
+    tau_reion_fn : calculates the optical depth to reionization.
+    
     At the moment we only support the CAMB tanh parameterization, but we need different approaches
     based on whether the use inputs the optical depth tau_reion or the reionization redshift z_reion.
+
     """
 
     z_reion : jnp.float64