refactor foreground params

simonsobs · Aug 13, 2024 · a3b346e · a3b346e
1 parent 4b176d4
commit a3b346e
Show file tree

Hide file tree

Showing 7 changed files with 76 additions and 94 deletions.
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
@@ -26,7 +26,7 @@ jobs:
       if: matrix.os == 'macos-latest'
       run: |
         set -x
-        sudo ln -s /opt/homebrew/bin/gfortran-11 /usr/local/bin/gfortran
+        sudo ln -s /opt/homebrew/bin/gfortran-12 /usr/local/bin/gfortran
         gfortran --version
 
     - name: Install dependencies via pip

diff --git a/mflike/Foreground.yaml b/mflike/Foreground.yaml
@@ -46,5 +46,3 @@ components:
   ee:
     - radio
     - dust
-
-params: !defaults [params_common, params_TT, params_TE, params_EE]
diff --git a/mflike/__init__.py b/mflike/__init__.py
@@ -1,21 +1,10 @@
 from importlib.metadata import PackageNotFoundError, version
 
-from .foreground import BandpowerForeground, Foreground
+from .foreground import BandpowerForeground, Foreground, EEForeground, TEForeground, TTForeground
 from .mflike import TT, TE, EE, TTTEEE
 
 try:
     __version__ = version("mflike")
 except PackageNotFoundError:
     # package is not installed
     pass
-
-__all__ = [
-    # mflike
-    "TT",
-    "TE",
-    "EE",
-    "TTTEEE",
-    # foreground
-    "Foreground",
-    "BandpowerForeground",
-]
diff --git a/mflike/foreground.py b/mflike/foreground.py
@@ -43,37 +43,14 @@
 import numpy as np
 from cobaya.log import LoggedError
 from cobaya.theory import Provider, Theory
+from cobaya.yaml import yaml_load
 from scipy import constants
 
 try:
     from numpy import trapezoid
 except ImportError:
     from numpy import trapz as trapezoid
 
-nuis_params_defaults = {
-    "a_tSZ": 3.30,
-    "a_kSZ": 1.60,
-    "a_p": 6.90,
-    "beta_p": 2.20,
-    "a_c": 4.90,
-    "beta_c": 2.20,
-    "a_s": 3.10,
-    "a_gtt": 2.80,
-    "a_gte": 0.10,
-    "a_gee": 0.10,
-    "a_psee": 0.0,
-    "a_pste": 0.0,
-    "xi": 0.10,
-    "T_d": 9.60,
-    "beta_s": -2.5,
-    "alpha_s": 1.0,
-    "T_effd": 19.6,
-    "beta_d": 1.5,
-    "alpha_dT": -0.6,
-    "alpha_dE": -0.4,
-    "alpha_p": 1.0,
-    "alpha_tSZ": 0.0
-}
 
 # Converts from cmb temperature to differential source intensity
 # (see eq. 8 of https://arxiv.org/abs/1303.5070).
@@ -99,7 +76,22 @@ def _cmb2bb(nu):
     return np.exp(x) * (nu * x / np.expm1(x)) ** 2
 
 
-class Foreground(Theory):
+class ForegroundParamsTheory(Theory):
+
+    @classmethod
+    def get_class_options(cls, input_options={}):
+        # dynamically generate defaults for params based on nbins
+        options = super().get_class_options().copy()
+        param_requested_cls = input_options.get(
+            'requested_cls') or options.get('requested_cls', ['tt', 'te', 'ee'])
+        params = yaml_load(cls.get_text_file_content('params_common.yaml'))
+        for spec in param_requested_cls:
+            params |= yaml_load(cls.get_text_file_content('params_%s.yaml' % spec.upper()))
+        options["params"] = params
+        return options
+
+
+class Foreground(ForegroundParamsTheory):
     normalisation: dict
     components: dict
     experiments: list[str]
@@ -120,40 +112,41 @@ def initialize(self):
         from fgspectra import frequency as fgf
         from fgspectra import power as fgp
 
-        template_path = os.path.join(os.path.dirname(os.path.abspath(fgp.__file__)), "data")
-        tsz_file = os.path.join(template_path, "cl_tsz_150_bat.dat")
-        cibc_file = os.path.join(template_path, "cl_cib_Choi2020.dat")
-        cibxtsz_file = os.path.join(template_path, "cl_sz_x_cib.dat")
+        self.fg_component_list = {s: self.components[s] for s in self.requested_cls}
 
+        template_path = os.path.join(os.path.dirname(os.path.abspath(fgp.__file__)), "data")
         # set pivot freq and multipole
         self.fg_nu_0 = self.normalisation["nu_0"]
         self.fg_ell_0 = self.normalisation["ell_0"]
 
-        # We don't seem to be using this
-        # cirrus = fgc.FactorizedCrossSpectrum(fgf.PowerLaw(), fgp.PowerLaw())
-        self.ksz = fgc.FactorizedCrossSpectrum(fgf.ConstantSED(), fgp.kSZ_bat())
-        self.cibp = fgc.FactorizedCrossSpectrum(fgf.ModifiedBlackBody(), fgp.PowerLaw())
-        self.radio = fgc.FactorizedCrossSpectrum(fgf.PowerLaw(), fgp.PowerLaw())
-        self.tsz = fgc.FactorizedCrossSpectrum(fgf.ThermalSZ(), fgp.PowerLawRescaledTemplate(tsz_file))
-        self.cibc = fgc.FactorizedCrossSpectrum(fgf.CIB(), fgp.PowerSpectrumFromFile(cibc_file))
-        self.dust = fgc.FactorizedCrossSpectrum(fgf.ModifiedBlackBody(), fgp.PowerLaw())
+        if 'tt' in self.requested_cls:
+            tsz_file = os.path.join(template_path, "cl_tsz_150_bat.dat")
+            cibc_file = os.path.join(template_path, "cl_cib_Choi2020.dat")
+            cibxtsz_file = os.path.join(template_path, "cl_sz_x_cib.dat")
+
+            # We don't seem to be using this
+            # cirrus = fgc.FactorizedCrossSpectrum(fgf.PowerLaw(), fgp.PowerLaw())
+            self.ksz = fgc.FactorizedCrossSpectrum(fgf.ConstantSED(), fgp.kSZ_bat())
+            self.cibp = fgc.FactorizedCrossSpectrum(fgf.ModifiedBlackBody(), fgp.PowerLaw())
+            self.tsz = fgc.FactorizedCrossSpectrum(fgf.ThermalSZ(), fgp.PowerLawRescaledTemplate(tsz_file))
+            self.cibc = fgc.FactorizedCrossSpectrum(fgf.CIB(), fgp.PowerSpectrumFromFile(cibc_file))
+
+            tsz_cib_sed = fgf.Join(fgf.ThermalSZ(), fgf.CIB())
+            tsz_cib_power_spectra = [
+                fgp.PowerLawRescaledTemplate(tsz_file),
+                fgp.PowerSpectrumFromFile(cibc_file),
+                fgp.PowerSpectrumFromFile(cibxtsz_file)
+            ]
+            tsz_cib_cl = fgp.PowerSpectraAndCovariance(*tsz_cib_power_spectra)
 
-        tsz_cib_sed = fgf.Join(fgf.ThermalSZ(), fgf.CIB())
-        tsz_cib_power_spectra = [
-            fgp.PowerLawRescaledTemplate(tsz_file),
-            fgp.PowerSpectrumFromFile(cibc_file),
-            fgp.PowerSpectrumFromFile(cibxtsz_file)
-        ]
-        tsz_cib_cl = fgp.PowerSpectraAndCovariance(*tsz_cib_power_spectra)
+            self.tSZ_and_CIB = fgc.CorrelatedFactorizedCrossSpectrum(tsz_cib_sed, tsz_cib_cl)
 
-        self.tSZ_and_CIB = fgc.CorrelatedFactorizedCrossSpectrum(tsz_cib_sed, tsz_cib_cl)
+        self.radio = fgc.FactorizedCrossSpectrum(fgf.PowerLaw(), fgp.PowerLaw())
+        self.dust = fgc.FactorizedCrossSpectrum(fgf.ModifiedBlackBody(), fgp.PowerLaw())
 
         if self.ells is None:
             self.ells = np.arange(self.lmin, self.lmax + 1)
 
-    def initialize_with_provider(self, provider: Provider):
-        self.fg_component_list = {s: self.components[s] for s in self.requested_cls}
-
     # Gets the actual power spectrum of foregrounds given the passed parameters
     def _get_foreground_model_arrays(self, fg_params, ell=None):
         r"""
@@ -358,7 +351,10 @@ def get_fg_totals(self):
 
     def must_provide(self, **requirements):
         if (req := requirements.get("fg_totals")) is not None:
-            self.requested_cls = req.get("requested_cls", self.requested_cls)
+            if set(self.requested_cls) != set(req.get("requested_cls",
+                                                      ['tt', 'te', 'ee'])):
+                raise ValueError(
+                    "requested_cls must be the same in Foreground and MFLike")
             self.ells = req.get("ells", self.ells)
             self.experiments = req.get("experiments", self.experiments)
 
@@ -371,7 +367,6 @@ class BandpowerForeground(Foreground):
 
     def initialize(self):
         super().initialize()
-        super().initialize_with_provider(self)
         if self.bands is None:
             self.bands = {
                 f"{exp}_s0": {"nu": [self.bandint_freqs[iexp]], "bandpass": [1.0]}
@@ -510,3 +505,15 @@ def _bandpass_construction(self, _initialize=False, **params):
             if self._initialized:
                 self.log.info("bandpass is delta function, no band integration performed")
         self._initialized = True
+
+
+class TTForeground(BandpowerForeground):
+    requested_cls = ['tt']
+
+
+class EEForeground(BandpowerForeground):
+    requested_cls = ['ee']
+
+
+class TEForeground(BandpowerForeground):
+    requested_cls = ['te']
diff --git a/mflike/mflike.py b/mflike/mflike.py
@@ -43,11 +43,8 @@
 from numbers import Real
 import sacc
 from cobaya.conventions import data_path, packages_path_input
-from cobaya.likelihoods.base_classes import InstallableLikelihood, _fast_chi_square
+from cobaya.likelihoods.base_classes import InstallableLikelihood
 from cobaya.log import LoggedError
-from cobaya.parameterization import expand_info_param
-
-from mflike.foreground import nuis_params_defaults
 
 
 class _MFLike(InstallableLikelihood):
@@ -68,8 +65,6 @@ class _MFLike(InstallableLikelihood):
     supported_params: dict
     lmax_theory: Optional[int]
 
-    _fast_chi_squared = _fast_chi_square()
-
     def initialize(self):
         # Set default values to data member not initialized via yaml file
         self.l_bpws = None
@@ -118,20 +113,11 @@ def initialize(self):
         self.lmax_theory = self.lmax_theory or 9000
         self.log.debug(f"Maximum multipole value: {self.lmax_theory}")
 
-        self.use_systematics_template = bool(self.systematics_template)
-
-        if self.use_systematics_template:
-            self.systematics_template = self.systematics_template
+        if self.systematics_template:
             # Initialize template for marginalization, if needed
             self._init_template_from_file()
 
         self._constant_nuisance: Optional[dict] = None
-
-        for p in nuis_params_defaults:
-            if p not in self.supported_params:
-                if isinstance(expand_info_param(self.params[p]).get("value"), Real):
-                    continue
-                self.params[p] = nuis_params_defaults[p]
         self.log.info("Initialized!")
 
     def get_fg_requirements(self):
@@ -513,7 +499,7 @@ def get_modified_theory(self, Dls, fg_totals, **nuis_params):
         cmbfg_dict = self._get_rotated_spectra(cmbfg_dict, **nuis_params)
 
         # Introduce templates of systematics from file, if needed
-        if self.use_systematics_template:
+        if self.systematics_template:
             cmbfg_dict = self._get_template_from_file(cmbfg_dict, **nuis_params)
 
         # Built theory
@@ -665,14 +651,14 @@ def _init_template_from_file(self):
         the ``syslibrary.syslib.ReadTemplateFromFile``
         function.
         """
-        if not self.systematics_template.get("rootname"):
+        if root := self.systematics_template.get("rootname"):
             raise LoggedError(self.log, "Missing 'rootname' for systematics template!")
 
         from syslibrary import syslib
 
         # decide where to store systematics template.
         # Currently stored inside syslibrary package
-        templ_from_file = syslib.ReadTemplateFromFile(rootname=self.systematics_template["rootname"])
+        templ_from_file = syslib.ReadTemplateFromFile(rootname=root)
         self.dltempl_from_file = templ_from_file(ell=self.l_bpws)
 
     def _get_template_from_file(self, dls_dict, **nuis_params):
@@ -715,5 +701,6 @@ class TT(_MFLike):
 class TE(_MFLike):
     ...
 
+
 class EE(_MFLike):
     ...
diff --git a/mflike/tests/test_mflike.py b/mflike/tests/test_mflike.py
@@ -82,7 +82,6 @@ class MFLikeTest(unittest.TestCase):
     def setUp(self):
         install({"likelihood": {"mflike.TTTEEE": None}}, path=packages_path)
 
-
     def test_mflike(self):
         from mflike import TTTEEE, BandpowerForeground
         import camb
@@ -129,7 +128,7 @@ def test_cobaya_TT(self):
                 },
             },
             "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}},
-                       "mflike.BandpowerForeground": {}},
+                       "mflike.BandpowerForeground": {'requested_cls': ['tt']}},
             "params": cosmo_params | nuis_params,
             "packages_path": packages_path,
         }
@@ -149,7 +148,7 @@ def test_cobaya_TE(self):
                 },
             },
             "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}},
-                       "mflike.BandpowerForeground": {}},
+                       "mflike.BandpowerForeground": {'requested_cls': ['te']}},
             "params": cosmo_params | nuis_params,
             "packages_path": packages_path,
         }
@@ -170,16 +169,18 @@ def test_cobaya_EE(self):
                 },
             },
             "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}},
-                       "mflike.BandpowerForeground": {}},
+                       "mflike.BandpowerForeground": {'requested_cls': ['ee']}},
             "params": cosmo_params | nuis_params,
             "packages_path": packages_path,
             "debug": True,
         }
-
-        model = get_model(info)
-        my_mflike = model.likelihood["mflike.EE"]
-        chi2 = -2 * (model.loglike(nuis_params, return_derived=False) - my_mflike.logp_const)
-        self.assertAlmostEqual(chi2, chi2s["ee"], 2)
+        for _ in (False, True):
+            model = get_model(info)
+            my_mflike = model.likelihood["mflike.EE"]
+            chi2 = -2 * (model.loglike(nuis_params, return_derived=False) - my_mflike.logp_const)
+            self.assertAlmostEqual(chi2, chi2s["ee"], 2)
+            info["theory"].pop("mflike.BandpowerForeground", None)
+            info["theory"]["mflike.EEForeground"] = None
 
     def test_cobaya_TTTEEE(self):
         nuis_params = common_nuis_params | TT_nuis_params | TE_nuis_params | EE_nuis_params

diff --git a/pyproject.toml b/pyproject.toml
@@ -26,7 +26,7 @@ requires-python = ">=3.9.0"
 dependencies = [
   "fgspectra>=1.2.1",
   "syslibrary>=0.2.0",
-  "cobaya>=3.4.1",
+  "cobaya>=3.5.2",
   "sacc>=0.9.0",
 ]