Add differentiable likelihood (#4)

* Added jax differentiable likelihood.

* Flake8 compliance.

* Added jax tests to workflow.

* Cobaya compliance, fixed import error in jax like import.

* Jax optional reqs changes.

* Added jax tests.

* Minor fixes & code style.

* Path for jax like test.

* Add ell cut for jax test.

* Fix indexing off-by-one in data load.

* Added missing test to default file run.

* Added HMC example.

* Updated readme.

* Typo.

* Typo.

Author: HTJense

.github/workflows/testing.yml

@@ -18,15 +18,15 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
 
       - name: Install Dependencies
         shell: bash -el {0}
         run: |
           python -m pip install --upgrade pip
           pip install pytest
           pip install -r requirements.txt
-          pip install -e .
+          pip install -e .[jax]
 
       - name: Cache testing Dataset
         shell: bash -el {0}
@@ -39,4 +39,4 @@ jobs:
       - name: Run Tests
         shell: bash -el {0}
         run: |
-          pytest -v act_dr6_cmbonly/tests/test_act.py
+          pytest -v

README.md

@@ -28,7 +28,7 @@ You can now run the tests with
 ```
 pytest -v --pyargs act_dr6_cmbonly
 ```
-If the tests return without any error (i.e. with only warnings), then the code is probably correctly installed. You can then attempt to run chains with
+If the tests return without any error (i.e. with only warnings), then the code is probably correctly installed. You may get some tests which get skipped if you do not install the differentiable likelihood (see below) - you do not need to worry about this. You can then attempt to run chains with
 ```
 cobaya-run yamls/run_act.yaml
 ```
@@ -46,3 +46,18 @@ If you ever get an error that the likelihood cannot locate the data, then you ca
 By default, the likelihood will look for the data in either
 - `<pip directory>/act_dr6_cmbonly/data/` if no cobaya packages path is given, or
 - `<cobaya packages path>/data/ACTDR6CMBonly/` if a cobaya packages path is given.
+
+## The differentiable likelihood
+
+If you are for whatever reason interested, I created a differentiable likelihood as well. You can install this by installing the package with
+```
+pip install -e .[jax]
+```
+Which will also install the `jax` and `cosmopower-jax` prerequisites. The differentiable likelihood can then be imported with
+```
+import act_dr6_cmbonly
+like = act_dr6_cmbonly.ACTDR6jax()
+```
+I provide an example of how to run a chain with the differentiable likelihood, see the `examples/run_hmc.py` file.
+
+For the most part, I do not expect that a differentiable likelihood adds much to ACT DR6 on its own. However, should people be interested in running joint analyses with other probes that provide differentiable likelihoods, then this simple likelihood should suffice. For the most part, it is simply 100 lines of python that does the same as the cobaya likelihood, but with JAX instead of numpy (as a result, some of the data is stored a bit differently internally to make use of JAX optimizations).

act_dr6_cmbonly/__init__.py

@@ -1,9 +1,14 @@
 __author__ = "Hidde T. Jense"
 __url__ = "https://github.com/ACTCollaboration/dr6-cmbonly"
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 try:
     from .act_dr6_cmbonly import ACTDR6CMBonly  # noqa: F401
     from .PlanckActCut import PlanckActCut  # noqa: F401
 except ImportError:
     pass
+
+try:
+    from .act_dr6_jaxlike import ACTDR6jax  # noqa: F401
+except ImportError:
+    ACTDR6jax = None  # noqa: F401

act_dr6_cmbonly/act_dr6_cmbonly.py

@@ -131,5 +131,5 @@ def loglike(self, cl):
         return self.logp_const + logp
 
     def logp(self, **param_values):
-        cl = self.theory.get_Cl(ell_factor=True)
+        cl = self.provider.get_Cl(ell_factor=True)
         return self.loglike(cl)

act_dr6_cmbonly/act_dr6_jaxlike.py

@@ -0,0 +1,136 @@
+import jax.numpy as np
+import os
+import sacc
+from typing import Sequence
+
+
+class ACTDR6jax:
+    """
+    A differentiable Likelihood implementation for ACT DR6
+    foreground-marginalized (cmb-only) data.
+
+    To make use of this module, make sure to install the package with
+        pip install -e .[jax]
+    to include any JAX dependencies.
+
+    I do not expect this class to get used a lot, but for the sake of
+    "it is easy and doable to write this in 100 lines of code", I wrote this.
+
+    Author: Hidde T. Jense
+    """
+    data_folder: str = "ACTDR6CMBonly"
+    input_filename: str = "act_dr6_cmb_sacc.fits"
+    polarizations: Sequence[str] = ["tt", "te", "ee"]
+    tt_lmax: int = 9000
+    ell_cuts: dict = {
+        "TT": [600, 6500],
+        "TE": [600, 6500],
+        "EE": [500, 6500]
+    }
+
+    def __init__(self, verbose: bool = False) -> None:
+        self.__verbose = verbose
+
+    def load_data(self) -> None:
+        # load the data
+        if self.verbose:
+            print(f"Loading data from {self.input_filename}.")
+
+        saccfile = sacc.Sacc.load_fits(
+            os.path.join(self.data_folder, self.input_filename)
+        )
+
+        idx_max = 0
+        pol_dt = {"t": "0", "e": "e", "b": "b"}
+        self.spec_meta = []
+        self.cull = []
+
+        for pol in self.polarizations:
+            p1, p2 = pol.lower()
+            t1, t2 = pol_dt[p1], pol_dt[p2]
+            dt = f"cl_{t1}{t2}"
+
+            tracers = saccfile.get_tracer_combinations(dt)
+
+            for tr1, tr2 in tracers:
+                if self.verbose:
+                    print(f"{tr1}x{tr2}")
+
+                lmin, lmax = self.ell_cuts.get(pol.upper(), (-np.inf, np.inf))
+
+                ls, mu, ind = saccfile.get_ell_cl(dt, tr1, tr2,
+                                                  return_ind=True)
+                mask = np.logical_and(ls >= lmin, ls <= lmax)
+                if not np.all(mask):
+                    if self.verbose:
+                        print(f"Cutting {pol} data to the \
+                            range [{lmin}-{lmax}].")
+                    self.cull.append(ind[~mask])
+                window = saccfile.get_bandpower_windows(ind[mask])
+
+                self.spec_meta.append({
+                    "data_type": dt,
+                    "tracer1": tr1,
+                    "tracer2": tr2,
+                    "pol": pol.lower(),
+                    "ell": ls[mask],
+                    "spec": mu[mask],
+                    "idx": ind[mask],
+                    "window": window
+                })
+
+                idx_max = max(idx_max, max(ind))
+
+        self.data_vec = np.zeros((idx_max + 1,))
+        self.spec_picker = np.zeros((idx_max + 1, len(self.spec_meta)))
+        self.win_func = np.zeros((idx_max + 1, self.tt_lmax - 1))
+
+        """
+        Most of the magic happens here - we create some binning and stacking
+        functions to ensure that everything can be quickly jit-ed by JAX.
+
+        (This is not a super clean way to do this, but it is the most
+        efficient way to do this easily.)
+        """
+        for i, m in enumerate(self.spec_meta):
+            self.data_vec = self.data_vec.at[m["idx"]].set(m["spec"])
+            self.spec_picker = self.spec_picker.at[m["idx"], i].set(1)
+
+            j1, j2 = m["window"].values.min()-2, m["window"].values.max()-2
+            if j2 >= self.tt_lmax-1:
+                j2 = self.tt_lmax-2
+            imax = j2-j1+1
+            self.win_func = self.win_func.at[m["idx"], j1:j2+1].set(
+                m["window"].weight[:imax].astype(float).T
+            )
+
+        self.covmat = np.array(saccfile.covariance.covmat)
+
+        for culls in self.cull:
+            self.covmat = self.covmat.at[culls, :].set(0.0)
+            self.covmat = self.covmat.at[:, culls].set(0.0)
+            self.covmat = self.covmat.at[culls, culls].set(1e10)
+
+        self.inv_cov = np.linalg.inv(self.covmat)
+        self.logp_const = -0.5 * np.log(2.0 * np.pi) * len(self.data_vec)
+        self.logp_const -= 0.5 * np.linalg.slogdet(self.covmat)[1]
+
+    def logp(self, dell: np.ndarray) -> float:
+        """
+        Compute the log-likelihood of some cosmology.
+        It expects an Lx3 array, where L is the number of ell-modes
+        (where dell[:,0] = TT[2..lmax], dell[:,1] = TE[2..lmax], and
+        dell[:,2] = EE[2..lmax]).
+        """
+        ps_vec = np.dot(self.win_func, dell[:self.tt_lmax-1])
+        ps_vec = np.sum(ps_vec * self.spec_picker, axis=1)
+
+        self.ps_vec = ps_vec
+
+        delta = self.data_vec - ps_vec
+        logp = -0.5 * (delta @ self.inv_cov @ delta)
+        return self.logp_const + logp
+
+    @property
+    def verbose(self) -> bool:
+        return self.__verbose

act_dr6_cmbonly/tests/test_jax.py

@@ -0,0 +1,71 @@
+"""
+This test checks that the differentiable DR6
+likelihood is working as intended.
+"""
+import pytest  # noqa F401
+import sys
+
+try:
+    import jax  # noqa F401
+    import jax.numpy as np
+    from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ
+except ImportError:
+    pass
+
+
+@pytest.mark.skipif("jax" not in sys.modules, reason="JAX is not installed.")
+def test_import_jaxlike():
+    from act_dr6_cmbonly import ACTDR6jax  # noqa F401
+
+
+@pytest.mark.skipif("jax" not in sys.modules, reason="JAX is not installed.")
+def test_jaxlike():
+    import act_dr6_cmbonly
+    like = act_dr6_cmbonly.ACTDR6jax()  # noqa F401
+
+
+@pytest.mark.skipif("jax" not in sys.modules, reason="JAX is not installed.")
+def test_jax_load_data():
+    import act_dr6_cmbonly
+    like = act_dr6_cmbonly.ACTDR6jax()
+    like.data_folder = "act_dr6_cmbonly/data"
+    like.load_data()
+
+
+@pytest.mark.skipif("cosmopower_jax" not in sys.modules,
+                    reason="Cosmopower-JAX is not installed.")
+def test_jax_loglike():
+    import act_dr6_cmbonly
+    T_CMB = 2.7255e6
+    cosmo_params = np.array([0.025, 0.12, 0.68, 0.054, 0.97, 3.05])
+
+    emu_tt = CPJ(probe='cmb_tt')
+    emu_te = CPJ(probe='cmb_te')
+    emu_ee = CPJ(probe='cmb_ee')
+    ellfac = emu_tt.modes * (emu_tt.modes + 1.0) / (2.0 * np.pi)
+    cl_tt = (T_CMB) ** 2.0 * emu_tt.predict(cosmo_params) * ellfac
+    cl_te = (T_CMB) ** 2.0 * emu_te.predict(cosmo_params) * ellfac
+    cl_ee = (T_CMB) ** 2.0 * emu_ee.predict(cosmo_params) * ellfac
+
+    cell = np.stack([cl_tt, cl_te, cl_ee], axis=1)
+
+    like = act_dr6_cmbonly.ACTDR6jax()
+    like.data_folder = "act_dr6_cmbonly/data"
+    like.ell_cuts = {
+        "TT": [600, 2508],
+        "TE": [600, 2508],
+        "EE": [500, 2508]
+    }
+    like.tt_lmax = 2508
+    like.load_data()
+
+    logp = like.logp(cell)
+
+    assert np.isclose(logp, -1236.189)
+
+
+if __name__ == "__main__":
+    test_import_jaxlike()
+    test_jax_load_data()
+    test_jaxlike()
+    test_jax_loglike()