Merge pull request #55 from LSSTDESC/fix-seed_sacc

Ensure the same seed leads to same cosmology shifts and added Metadata to sacc

src/smokescreen/datavector.py

@@ -21,6 +21,8 @@
 import os
 import types
 import inspect
+import datetime
+import getpass
 from copy import deepcopy
 import pyccl as ccl
 import sacc
@@ -81,6 +83,8 @@ def __init__(self, cosmo, systm_dict, likelihood, shifts_dict, sacc_data,
         self.systematics_dict = systm_dict
         # save the data-vector
         self.sacc_data = sacc_data
+        # seed for the random number generator
+        self.seed = seed
         # checks if the sacc_data is in the correct format:
         assert isinstance(self.sacc_data, sacc.sacc.Sacc), "sacc_data must be a sacc object"
         # load the likelihood
@@ -227,7 +231,7 @@ def _create_concealed_cosmo(self):
         # sometimes we have this extra paramters that can cause problems:
         try:
             del concealed_cosmo_dict['extra_parameters']
-        except KeyError:
+        except KeyError:  # pragma: no cover
             pass
         for k in self.__shifts.keys():
             concealed_cosmo_dict[k] = self.__shifts[k]
@@ -327,6 +331,14 @@ def save_concealed_datavector(self, path_to_save, file_root,
         concealed_sacc = save_to_sacc(self.sacc_data,
                                       self.concealed_data_vector,
                                       idx)
+        # copies the metadata from the original sacc file:
+        concealed_sacc.metadata = self.sacc_data.metadata
+        # adds metadata to the sacc file:
+        concealed_sacc.metadata['concealed'] = True
+        concealed_sacc.metadata['creator'] = getpass.getuser()
+        concealed_sacc.metadata['creation'] = datetime.datetime.now().isoformat()
+        concealed_sacc.metadata['info'] = 'Concealed (blinded) data-vector, created by Smokescreen.'
+        concealed_sacc.metadata['seed_smokescreen'] = self.seed
         concealed_sacc.save_fits(f"{path_to_save}/{file_root}_concealed_data_vector.fits",
                                  overwrite=True)
         if return_sacc:

src/smokescreen/param_shifts.py

@@ -88,13 +88,17 @@ def draw_flat_or_deterministic_param_shifts(cosmo, shifts_dict, seed):
             # failed_keys.append(key)
             raise ValueError(f"[{error}]Key {key} not in cosmology parameters")
     shifts = {}
-    for key, value in shifts_dict.items():
-        if isinstance(value, tuple):
-            # check if the tuple is of length 2
-            if len(value) == 2:
-                shifts[key] = np.random.uniform(value[0], value[1])
+    # we loop over the ccl dict keys to ensure params are drawn in the same order!
+    for key in cosmo.to_dict().keys():
+        if key in shifts_dict.keys():
+            if isinstance(shifts_dict[key], tuple):
+                # check if the tuple is of length 2
+                if len(shifts_dict[key]) == 2:
+                    shifts[key] = np.random.uniform(shifts_dict[key][0], shifts_dict[key][1])
+                else:
+                    raise ValueError(f"Tuple {shifts_dict[key]} has to be of length 2")
             else:
-                raise ValueError(f"Tuple {value} has to be of length 2")
+                shifts[key] = shifts_dict[key]
         else:
-            shifts[key] = value
+            pass
     return shifts

tests/test_datavector.py

@@ -1,6 +1,8 @@
-import pytest
+import pytest  # noqa: F401
 import types
 import os
+import datetime
+from unittest.mock import patch, MagicMock  # noqa: F401
 import numpy as np
 import sacc
 import pyccl as ccl
@@ -306,7 +308,8 @@ def test_load_likelihood():
         smokescreen._load_likelihood(invalid_likelihood, sacc_data)
 
 
-def test_save_concealed_datavector():
+@patch('src.smokescreen.datavector.getpass.getuser', return_value='test_user')
+def test_save_concealed_datavector(mock_getuser):
     # Create mock inputs
     cosmo = COSMO
     likelihood = "./examples/cosmic_shear/cosmicshear_likelihood.py"
@@ -316,7 +319,8 @@ def test_save_concealed_datavector():
     }
     shift_dict = {"Omega_c": 0.34, "sigma8": 0.85}
     sacc_data = sacc.Sacc.load_fits("./examples/cosmic_shear/cosmicshear_sacc.fits")
-    sck = ConcealDataVector(cosmo, syst_dict, likelihood, shift_dict, sacc_data)
+    sck = ConcealDataVector(cosmo, syst_dict, likelihood,
+                            shift_dict, sacc_data, seed=1234)
 
     # Calculate the concealing factor and apply it to the likelihood data vector
     sck.calculate_concealing_factor()
@@ -326,7 +330,18 @@ def test_save_concealed_datavector():
     temp_file_path = "./tests/"
     temp_file_root = "temp_sacc"
     temp_file_name = f"{temp_file_path}{temp_file_root}_concealed_data_vector.fits"
-    sck.save_concealed_datavector(temp_file_path, temp_file_root)
+    returned_sacc = sck.save_concealed_datavector(temp_file_path,
+                                                  temp_file_root,
+                                                  return_sacc=True)
+    # checks if the return is a sacc object
+    assert isinstance(returned_sacc, sacc.Sacc)
+
+    returned_sacc = sck.save_concealed_datavector(temp_file_path,
+                                                  temp_file_root,
+                                                  return_sacc=False)
+
+    # Check that the return is None
+    assert returned_sacc is None
 
     # Check that the file was created
     assert os.path.exists(temp_file_name)
@@ -335,5 +350,11 @@ def test_save_concealed_datavector():
     loaded_sacc = sacc.Sacc.load_fits(temp_file_name)
     np.testing.assert_array_equal(loaded_sacc.mean, blinded_dv)
 
+    info_str = 'Concealed (blinded) data-vector, created by Smokescreen.'
+    assert loaded_sacc.metadata['concealed'] is True
+    assert loaded_sacc.metadata['creator'] == mock_getuser.return_value
+    assert loaded_sacc.metadata['creation'][:10] == datetime.date.today().isoformat()
+    assert loaded_sacc.metadata['info'] == info_str
+    assert loaded_sacc.metadata['seed_smokescreen'] == 1234
     # Clean up the temporary file
     os.remove(temp_file_name)