Refactor reionization models, add sample ExpReionization alternative …

…to Tanh (cmbant#147)

* add ExpReionization

* fix ReadParams

* fix packaging

* fix validate

* notebook

* python 12

* readthedocs

* readthedocs

* readthedocs

* update sample notebook

* update sample notebook

.readthedocs.yml

@@ -7,9 +7,10 @@ build:
 
 python:
   install:
-    - requirements: docs/requirements.txt
     - method: pip
       path: .
+      extra_requirements:
+        - docs
 
 sphinx:
   configuration: docs/source/conf.py

.travis.yml

@@ -11,6 +11,8 @@ branches:
     - CAMB_sources
     - rayleigh
 
+if: (type != pull_request) OR (type = pull_request AND TRAVIS_PULL_REQUEST_SLUG != 'cmbant/camb')
+
 jobs:
   include:
     - name: "Ubuntu 20.04 + Python 3.7"

camb/__init__.py

@@ -7,7 +7,7 @@
 __author__ = "Antony Lewis"
 __contact__ = "antony at cosmologist dot info"
 __url__ = "https://camb.readthedocs.io"
-__version__ = "1.5.0"
+__version__ = "1.5.1"
 
 from . import baseconfig
 
@@ -23,7 +23,7 @@
 from . import nonlinear
 from .model import CAMBparams, TransferParams
 from .results import CAMBdata, MatterTransferData, ClTransferData
-from .reionization import TanhReionization
+from .reionization import TanhReionization, ExpReionization
 from .nonlinear import Halofit
 from .dark_energy import DarkEnergyFluid, DarkEnergyPPF
 from .initialpower import InitialPowerLaw, SplinedInitialPower

camb/_compilers.py

@@ -3,7 +3,6 @@
 import struct
 import platform
 import re
-from pkg_resources import parse_version
 
 is_windows = platform.system() == "Windows"
 
@@ -38,6 +37,7 @@ def check_ifort():
 
 
 def check_gfortran(version=gfortran_min, msg=False, retry=False):
+    from packaging.version import parse as parse_version
     global compiler_environ
     gfortran_version = get_gfortran_version()
     version = str(version)

camb/camb.py

@@ -154,6 +154,7 @@ def do_set(setter):
     do_set(cp.set_accuracy)
     do_set(cp.set_classes)
     do_set(cp.DarkEnergy.set_params)
+    do_set(cp.Reion.set_extra_params)
     do_set(cp.set_cosmology)
     do_set(cp.set_matter_power)
     do_set(cp.set_for_lmax)
@@ -203,6 +204,7 @@ def extract_params(set_func):
                     params.add(arg)
 
     extract_params(cp.DarkEnergy.set_params)
+    extract_params(cp.Reion.set_extra_params)
     extract_params(cp.set_cosmology)
     if not transfer_only:
         extract_params(cp.InitPower.set_params)
@@ -232,14 +234,16 @@ def set_params_cosmomc(p, num_massive_neutrinos=1, neutrino_hierarchy='degenerat
     pars = inpars or model.CAMBparams()
     if p.get('alpha1', 0) or p.get('Aphiphi', 1) != 1:
         raise ValueError('Parameter not currrently supported by set_params_cosmomc')
+
+    pars.set_dark_energy(w=p.get('w', -1), wa=p.get('wa', 0), dark_energy_model=dark_energy_model)
+    pars.Reion.set_extra_params(deltazrei=p.get('deltazrei', None))
     pars.set_cosmology(H0=p['H0'], ombh2=p['omegabh2'], omch2=p['omegach2'], mnu=p.get('mnu', 0.06),
-                       omk=p.get('omegak', 0), tau=p['tau'], deltazrei=p.get('deltazrei', None),
+                       omk=p.get('omegak', 0), tau=p['tau'],
                        nnu=p.get('nnu', constants.default_nnu), Alens=p.get('Alens', 1.0),
                        YHe=p.get('yheused', None), meffsterile=p.get('meffsterile', 0),
                        num_massive_neutrinos=num_massive_neutrinos, neutrino_hierarchy=neutrino_hierarchy)
     pars.InitPower.set_params(ns=p['ns'], r=p.get('r', 0), As=p['A'] * 1e-9, nrun=p.get('nrun', 0),
                               nrunrun=p.get('nrunrun', 0))
-    pars.set_dark_energy(w=p.get('w', -1), wa=p.get('wa', 0), dark_energy_model=dark_energy_model)
     pars.set_for_lmax(lmax, lens_potential_accuracy=lens_potential_accuracy)
     pars.NonLinearModel.set_params(halofit_version=halofit_version)
     pars.WantTensors = pars.InitPower.has_tensors()

camb/model.py

@@ -10,6 +10,7 @@
 from .nonlinear import NonLinearModel
 from .dark_energy import DarkEnergyModel, DarkEnergyEqnOfState
 from .recombination import RecombinationModel
+from .reionization import ReionizationModel
 from .sources import SourceWindow
 from . import bbn
 import logging
@@ -418,7 +419,7 @@ def set_cosmology(self, H0: Optional[float] = None, ombh2=0.022, omch2=0.12, omk
                       neutrino_hierarchy: Union[str, int] = 'degenerate', num_massive_neutrinos=1,
                       mnu=0.06, nnu=constants.default_nnu, YHe: Optional[float] = None, meffsterile=0.0,
                       standard_neutrino_neff=constants.default_nnu, TCMB=constants.COBE_CMBTemp,
-                      tau: Optional[float] = None, zrei: Optional[float] = None, deltazrei: Optional[float] = None,
+                      tau: Optional[float] = None, zrei: Optional[float] = None,
                       Alens=1.0, bbn_predictor: Union[None, str, bbn.BBNPredictor] = None, theta_H0_range=(10, 100)):
         r"""
         Sets cosmological parameters in terms of physical densities and parameters (e.g. as used in Planck analyses).
@@ -461,7 +462,6 @@ def set_cosmology(self, H0: Optional[float] = None, ombh2=0.022, omch2=0.12, omk
         :param TCMB: CMB temperature (in Kelvin)
         :param tau: optical depth; if None and zrei is None, current Reion settings are not changed
         :param zrei: reionization mid-point optical depth (set tau=None to use this)
-        :param deltazrei: redshift width of reionization; if None, uses default
         :param Alens: (non-physical) scaling of the lensing potential compared to prediction
         :param bbn_predictor: :class:`.bbn.BBNPredictor` instance used to get YHe from BBN consistency if YHe is None,
          or name of a BBN predictor class, or file name of an interpolation table
@@ -527,11 +527,9 @@ def set_cosmology(self, H0: Optional[float] = None, ombh2=0.022, omch2=0.12, omk
         if tau is not None:
             if zrei is not None:
                 raise CAMBError('Cannot set both tau and zrei')
-            self.Reion.set_tau(tau, delta_redshift=deltazrei)
+            self.Reion.set_tau(tau)
         elif zrei is not None:
-            self.Reion.set_zrei(zrei, delta_redshift=deltazrei)
-        elif deltazrei:
-            raise CAMBError('must set tau if setting deltazrei')
+            self.Reion.set_zrei(zrei)
 
         return self
 
@@ -570,14 +568,16 @@ def N_eff(self):
             return sum(self.nu_mass_degeneracies[:self.nu_mass_eigenstates]) + self.num_nu_massless
 
     def set_classes(self, dark_energy_model=None, initial_power_model=None,
-                    non_linear_model=None, recombination_model=None):
+                    non_linear_model=None, recombination_model=None,
+                    reionization_model=None):
         """
         Change the classes used to implement parts of the model.
 
         :param dark_energy_model: 'fluid', 'ppf', or name of a DarkEnergyModel class
         :param initial_power_model: name of an InitialPower class
         :param non_linear_model: name of a NonLinearModel class
-        :param recombination_model: name of recombination_model class
+        :param recombination_model: name of RecombinationModel class
+        :param reionization_model: name of a ReionizationModel class
         """
         if dark_energy_model:
             self.DarkEnergy = self.make_class_named(dark_energy_model, DarkEnergyModel)
@@ -587,6 +587,8 @@ def set_classes(self, dark_energy_model=None, initial_power_model=None,
             self.NonLinear = self.make_class_named(non_linear_model, NonLinearModel)
         if recombination_model:
             self.Recomb = self.make_class_named(recombination_model, RecombinationModel)
+        if reionization_model:
+            self.Reion = self.make_class_named(reionization_model, ReionizationModel)
 
     def set_dark_energy(self, w=-1.0, cs2=1.0, wa=0, dark_energy_model='fluid'):
         r"""

camb/reionization.py

@@ -1,4 +1,4 @@
-from .baseconfig import F2003Class, fortran_class
+from .baseconfig import F2003Class, fortran_class, f_pointer
 from ctypes import c_bool, c_double, POINTER, byref, c_void_p
 
 
@@ -10,17 +10,14 @@ class ReionizationModel(F2003Class):
         ("Reionization", c_bool, "Is there reionization? (can be off for matter power, which is independent of it)")]
 
 
-@fortran_class
-class TanhReionization(ReionizationModel):
+class BaseTauWithHeReionization(ReionizationModel):
     """
-    This default (unphysical) tanh x_e parameterization is described in
-    Appendix B of `arXiv:0804.3865 <https://arxiv.org/abs/0804.3865>`_
+    Abstract class for models that map z_re to tau, and include second reionization of Helium
     """
     _fields_ = [
         ("use_optical_depth", c_bool, "Whether to use the optical depth or redshift paramters"),
-        ("redshift", c_double, "Reionization redshift if use_optical_depth=False"),
+        ("redshift", c_double, "Reionization redshift (xe=0.5) if use_optical_depth=False"),
         ("optical_depth", c_double, "Optical depth if use_optical_depth=True"),
-        ("delta_redshift", c_double, "Duration of reionization"),
         ("fraction", c_double,
          "Reionization fraction when complete, or -1 for full ionization of hydrogen and first ionization of helium."),
         ("include_helium_fullreion", c_bool, "Whether to include second reionization of helium"),
@@ -30,28 +27,45 @@ class TanhReionization(ReionizationModel):
         ("tau_solve_accuracy_boost", c_double, "Accuracy boosting parameter for solving for z_re from tau"),
         ("timestep_boost", c_double,
          "Accuracy boosting parameter for the minimum number of time sampling steps through reionization"),
-        ("max_redshift", c_double, "Maxmimum redshift allowed when mapping tau into reionization redshift")]
+        ("max_redshift", c_double, "Maxmimum redshift allowed when mapping tau into reionization redshift"),
+        ("__min_redshift", c_double, "Minimim redshift allowed when mapping tau into reionization redshift"),
+        ("__fHe", c_double, "Helium fraction"),
+        ("__state", f_pointer)
+    ]
 
     _fortran_class_module_ = 'Reionization'
-    _fortran_class_name_ = 'TTanhReionization'
+    _fortran_class_name_ = 'TBaseTauWithHeReionization'
 
     _methods_ = [('GetZreFromTau', [c_void_p, POINTER(c_double)], c_double, {"nopass": True})]
 
-    def set_zrei(self, zrei, delta_redshift=None):
+    def get_zre(self, params, tau=None):
+        """
+        Get the midpoint redshift of reionization.
+
+        :param params: :class:`.model.CAMBparams` instance with cosmological parameters
+        :param tau: if set, calculate the redshift for optical depth tau, otherwise uses curently set parameters
+        :return: reionization mid-point redshift
+        """
+        if self.use_optical_depth or tau:
+            from .camb import CAMBparams
+            assert isinstance(params, CAMBparams)
+            return self.f_GetZreFromTau(byref(params), c_double(tau or self.optical_depth))
+        else:
+            return self.redshift
+
+    def set_zrei(self, zrei):
         """
         Set the mid-point reionization redshift
 
         :param zrei: mid-point redshift
         :param delta_redshift:  delta z for reionization
-        :return:  self
+        :return: self
         """
         self.use_optical_depth = False
         self.redshift = zrei
-        if delta_redshift is not None:
-            self.delta_redshift = delta_redshift
         return self
 
-    def set_tau(self, tau, delta_redshift=None):
+    def set_tau(self, tau):
         """
         Set the optical depth
 
@@ -61,21 +75,56 @@ def set_tau(self, tau, delta_redshift=None):
         """
         self.use_optical_depth = True
         self.optical_depth = tau
-        if delta_redshift is not None:
-            self.delta_redshift = delta_redshift
         return self
 
-    def get_zre(self, params, tau=None):
+
+@fortran_class
+class TanhReionization(BaseTauWithHeReionization):
+    """
+    This default (unphysical) tanh x_e parameterization is described in
+    Appendix B of `arXiv:0804.3865 <https://arxiv.org/abs/0804.3865>`_
+    """
+    _fields_ = [
+        ("delta_redshift", c_double, "Duration of reionization")]
+
+    _fortran_class_name_ = 'TTanhReionization'
+
+    def set_extra_params(self, deltazrei=None):
         """
-        Get the midpoint redshift of reionization.
+        Set extra parameters (not tau, or zrei)
 
-        :param params: :class:`.model.CAMBparams` instance with cosmological parameters
-        :param tau: if set, calculate the redshift for optical depth tau, otherwise uses curently set parameters
-        :return: reionization mid-point redshift
+        :param delta_redshift: delta z for reionization
         """
-        if self.use_optical_depth or tau:
-            from .camb import CAMBparams
-            assert isinstance(params, CAMBparams)
-            return self.f_GetZreFromTau(byref(params), c_double(tau or self.optical_depth))
-        else:
-            return self.redshift
+        if deltazrei is not None:
+            self.delta_redshift = deltazrei
+
+
+@fortran_class
+class ExpReionization(BaseTauWithHeReionization):
+    """
+        An ionization fraction that decreases exponentially at high z, saturating to fully inionized at fixed redshift.
+        This model has a minimum non-zero tau around 0.04 for reion_redshift_complete=6.1.
+        Similar to e.g. arXiv:1509.02785, arXiv:2006.16828, but not attempting to fit shape near x_e~1 at z<6.1
+    """
+    _fields_ = [
+        ("reion_redshift_complete", c_double, "end of reionization"),
+        ("reion_exp_smooth_width", c_double, "redshift scale to smooth exponential"),
+        ("reion_exp_power", c_double, "power in exponential, exp(-lambda(z-redshift_complete)^exp_power)")]
+
+    _fortran_class_name_ = 'TExpReionization'
+
+    def set_extra_params(self, reion_redshift_complete=None, reion_exp_power=None, reion_exp_smooth_width=None):
+        """
+        Set extra parameters (not tau, or zrei)
+
+        :param reion_redshift_complete: redshift at which reionization complete (e.g. around 6)
+        :param reion_exp_power: power in exponential decay with redshift
+        :param reion_exp_smooth_width: smoothing parameter to keep derivative smooth
+        """
+
+        if reion_redshift_complete is not None:
+            self.reion_redshift_complete = reion_redshift_complete
+        if reion_exp_power is not None:
+            self.reion_exp_power = reion_exp_power
+        if reion_exp_smooth_width is not None:
+            self.reion_exp_smooth_width = reion_exp_smooth_width