Download data function (#12)

* added theory return for cosmosis

* add function to download data

* add instructions in case of data not loading

* version bump

act_dr6_lenslike/__init__.py

@@ -1,5 +1,5 @@
 """Likelihood for the Atacama Cosmology Telescope DR6 CMB lensing data."""
 __author__ = ["Mathew Madhavacheril", "Ian Harrison"]
-__version__ = "1.0.1"
+__version__ = "1.0.3"
 
 from .act_dr6_lenslike import *

act_dr6_lenslike/act_dr6_lenslike.py

@@ -24,6 +24,54 @@
 # HELPER FUNCTIONS
 # ================
 
+def download(url, filename):
+    # thanks to https://stackoverflow.com/a/63831344
+    # this function can be considered CC-BY-SA 4.0
+    import functools
+    import pathlib
+    import shutil
+    import requests
+    from tqdm.auto import tqdm
+
+    r = requests.get(url, stream=True, allow_redirects=True)
+    if r.status_code != 200:
+        r.raise_for_status()  # Will only raise for 4xx codes, so...
+        raise RuntimeError(f"Request to {url} returned status code {r.status_code}")
+    file_size = int(r.headers.get('Content-Length', 0))
+
+    path = pathlib.Path(filename).expanduser().resolve()
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    desc = "(Unknown total file size)" if file_size == 0 else ""
+    r.raw.read = functools.partial(r.raw.read, decode_content=True)  # Decompress if needed
+    with tqdm.wrapattr(r.raw, "read", total=file_size, desc=desc) as r_raw:
+        with path.open("wb") as f:
+            shutil.copyfileobj(r_raw, f)
+
+    return path
+
+def get_data(data_url="https://lambda.gsfc.nasa.gov/data/suborbital/ACT/ACT_dr6/likelihood/data/",
+             data_filename="ACT_dr6_likelihood_v1.1.tgz"):
+
+    if os.path.exists(os.path.join(file_dir, data_dir)):
+        print('Data already exists at {}, not downloading again.'.format(os.path.join(file_dir, data_dir)))
+    else:
+        import tarfile
+
+        orig_cwd = os.getcwd()
+        os.mkdir(os.path.join(file_dir, data_dir))
+        os.chdir(os.path.join(file_dir, data_dir))
+
+        print('Downloading data {} and placing it in likelihood folder.'.format(data_filename))
+        download(data_url+data_filename, data_filename)
+
+        tar = tarfile.open(data_filename)
+        tar.extractall()
+        tar.close()
+
+        os.remove(data_filename)
+        os.chdir(orig_cwd)
+
 def pp_to_kk(clpp,ell):
     return clpp * (ell*(ell+1.))**2. / 4.
 
@@ -175,6 +223,12 @@ def load_data(variant, ddir=data_dir,
     """
     # TODO: review defaults
 
+    if not os.path.exists(ddir):
+        raise FileNotFoundError("Requested data directory {} does not exist.\
+                                Please place the data there. Default data can \
+                                be downloaded to the default location \
+                                with the act_dr6_lenslike.get_data() function.".format(ddir))
+
 
     v,baseline,include_planck = parse_variant(variant)
 
@@ -345,7 +399,7 @@ def generic_lnlike(data_dict,ell_kk,cl_kk,ell_cmb,cl_tt,cl_ee,cl_te,cl_bb,trim_l
     delta = d['data_binned_clkk'] - bclkk
     lnlike = -0.5 * np.dot(delta,np.dot(cinv,delta))
     if return_theory:
-        lnlike, bclkk
+        return lnlike, bclkk
     else:
         return lnlike
 

act_dr6_lenslike/tests/test_act.py

@@ -7,7 +7,7 @@
 
 class ACTLikeTest(unittest.TestCase):
 
-    def generic_call(self,variant,lens_only,exp_chisq=None):
+    def generic_call(self,variant,lens_only,exp_chisq=None,return_theory=False):
         data_file = data_dir+'like_corrs/cosmo2017_10K_acc3_lenspotentialCls.dat'
         try:
             ell, cl_tt, cl_ee, cl_bb, cl_te, cl_pp, cl_tp, cl_ep= np.loadtxt(data_file, unpack=True)
@@ -26,13 +26,20 @@ def generic_call(self,variant,lens_only,exp_chisq=None):
         data_dict = alike.load_data(variant,lens_only=lens_only,like_corrections=not(lens_only))
         ell_kk = ell
         ell_cmb=ell
-        chisq=-2*alike.generic_lnlike(data_dict,ell_kk,cl_kk,ell_cmb,cl_tt,cl_ee,cl_te,cl_bb,trim_lmax = 2998)
-        self.assertAlmostEqual(chisq,  exp_chisq, 1)
+
+        if return_theory:
+            chisq,bclkk=alike.generic_lnlike(data_dict,ell_kk,cl_kk,ell_cmb,cl_tt,cl_ee,cl_te,cl_bb,trim_lmax = 2998,return_theory=True)
+            self.assertAlmostEqual(-2*chisq,  exp_chisq, 1)
+        else:
+            chisq=-2*alike.generic_lnlike(data_dict,ell_kk,cl_kk,ell_cmb,cl_tt,cl_ee,cl_te,cl_bb,trim_lmax = 2998)
+            self.assertAlmostEqual(chisq,  exp_chisq, 1)
 
     def test_act_baseline_lensonly(self):
         self.generic_call('act_baseline',True,14.06)
     def test_act_baseline(self):
         self.generic_call('act_baseline',False,14.71)
+    def test_act_baseline_return_theory(self):
+        self.generic_call('act_baseline',False,14.71,return_theory=True)
     def test_actplanck_baseline_lensonly(self):
         self.generic_call('actplanck_baseline',True,21.07)
     def test_actplanck_baseline(self):
@@ -55,6 +62,7 @@ def test_act_cinpaint_lensonly(self):
 if __name__ == '__main__':
     ACTLikeTest().test_act_baseline_lensonly()
     ACTLikeTest().test_act_baseline()
+    ACTLikeTest().test_act_baseline_return_theory()
     ACTLikeTest().test_actplanck_baseline_lensonly()
     ACTLikeTest().test_actplanck_baseline()
     ACTLikeTest().test_act_extended_lensonly()

requirements.txt

@@ -1,3 +1,5 @@
 pytest
 numpy
-scipy
+scipy
+requests
+tqdm