added suffix to main and datavector.py

src/smokescreen/__main__.py

@@ -42,6 +42,7 @@ def datavector_main(path_to_sacc: Path_fr,
                     reference_cosmology: Union[dict, CosmologyType] = ccl.CosmologyVanillaLCDM(),
                     path_to_output: Path_drw = None,
                     keep_original_sacc: bool = False,
+                    output_suffix: str = None,
                     ) -> None:
     r"""Main function to conceal a SACC file using a firecrown likelihood.
 
@@ -50,19 +51,21 @@ def datavector_main(path_to_sacc: Path_fr,
         likelihood_path (str): Path to the firecrown likelihood module file.
         shifts_dict (dict): Dictionary with fixed values for the firecrown shifts parameters.
             Example: {"Omega_c": (0.20, 0.39), "sigma8": (0.70, 0.90)}
-        shift_type (str): Type of shift to apply to the data vector. 
+        shift_type (str): Type of shift to apply to the data vector.
             Options are 'add' and 'mult'. Defaults to 'add'.
         systematics (dict): Dictionary with fixed values for the firecrown systematics parameters.
-        shift_distribution (str): Distribution type for the parameter shifts. 
+        shift_distribution (str): Distribution type for the parameter shifts.
             Options are 'flat' and 'gaussian'. Defaults to 'flat'.
         seed (int, str): Seed for the blinding process. Defaults to 2112.
-        reference_cosmology (Union[CosmologyType, dict]): 
-            Cosmology object or dictionary with cosmological 
+        reference_cosmology (Union[CosmologyType, dict]):
+            Cosmology object or dictionary with cosmological
             parameters you want different than the VanillaLCDM as reference cosmology.
             Defaults to ccl.CosmologyVanillaLCDM().
         path_to_output (str): Path to save the blinded sacc file. Defaults to None.
-        keep_original_sacc (bool): If True, keeps the original sacc file. 
+        keep_original_sacc (bool): If True, keeps the original sacc file.
             Defaults to False [keeps only the encrypted file].
+        output_suffix (str): Custom suffix for the output file name.
+            Defaults to None (uses 'concealed_data_vector').
     """
     print(banner)
     if isinstance(reference_cosmology, dict):
@@ -88,15 +91,18 @@ def datavector_main(path_to_sacc: Path_fr,
     # saves the blinded sacc file
     if path_to_output is not None:
         smoke.save_concealed_datavector(path_to_output, root_name,
-                                        output_format=input_format)
+                                        output_format=input_format,
+                                        suffix=output_suffix)
     else:
         # get the input file directory
         path_to_output = os.path.dirname(path_to_sacc)
         smoke.save_concealed_datavector(path_to_output, root_name,
-                                        output_format=input_format)
+                                        output_format=input_format,
+                                        suffix=output_suffix)
     # Determine extension based on format
     ext = '.hdf5' if input_format == 'hdf5' else '.fits'
-    outprintfile = f"{path_to_output}/{root_name}_concealed_data_vector{ext}"
+    _suffix = output_suffix if output_suffix is not None else "concealed_data_vector"
+    outprintfile = f"{path_to_output}/{root_name}_{_suffix}{ext}"
     print(f"\nConcealed sacc file saved as:\n\t{outprintfile}")
 
     print(f"\nEncrypting the original sacc file {path_to_sacc} ...", end="")

src/smokescreen/datavector.py

@@ -533,7 +533,8 @@ def apply_concealing_to_likelihood_datavec(self):
         return self.concealed_data_vector
 
     def save_concealed_datavector(self, path_to_save, file_root,
-                                  return_sacc=False, output_format=None):
+                                  return_sacc=False, output_format=None,
+                                  suffix=None):
         """
         Saves the concealed (blinded) data-vector to a file.
 
@@ -550,13 +551,17 @@ def save_concealed_datavector(self, path_to_save, file_root,
             If True, returns the sacc object with the blinded data-vector.
         output_format : str, optional
             Output format to use. If None, uses the detected input format.
+        suffix : str, optional
+            Suffix for the output file name. Defaults to 'concealed_data_vector'.
 
         Returns
         -------
         sacc.sacc.Sacc or None
             If `return_sacc` is True, returns the sacc object with
             the blinded data-vector. Otherwise, returns None.
         """
+        if suffix is None:
+            suffix = "concealed_data_vector"
         # Determine output format: use specified format or fall back to input format
         if output_format is None:
             output_format = getattr(self, '_input_format', 'fits')
@@ -582,7 +587,7 @@ def save_concealed_datavector(self, path_to_save, file_root,
             ext = '.fits'
             save_method = concealed_sacc.save_fits
 
-        output_path = f"{path_to_save}/{file_root}_concealed_data_vector{ext}"
+        output_path = f"{path_to_save}/{file_root}_{suffix}{ext}"
         save_method(output_path, overwrite=True)
         if return_sacc:
             return concealed_sacc

tests/test_datavector.py

@@ -908,3 +908,56 @@ def test_save_concealed_datavector_default_format_uses_input_format(mock_getuser
     np.testing.assert_array_equal(loaded_sacc.mean, blinded_dv)
 
     os.remove(expected_hdf5_name)
+
+
+@patch('src.smokescreen.datavector.getpass.getuser', return_value='test_user')
+def test_save_concealed_datavector_custom_suffix(mock_getuser):
+    cosmo = COSMO
+    likelihood = "./examples/cosmic_shear/cosmicshear_likelihood.py"
+    syst_dict = {
+        "trc1_delta_z": 0.1,
+        "trc0_delta_z": 0.1,
+    }
+    shift_dict = {"Omega_c": 0.34, "sigma8": 0.85}
+    sacc_data, _ = load_sacc_file("./examples/cosmic_shear/cosmicshear_sacc.fits")
+    sck = ConcealDataVector(cosmo, likelihood, shift_dict, sacc_data, syst_dict, seed=1234)
+    sck.calculate_concealing_factor()
+    sck.apply_concealing_to_likelihood_datavec()
+
+    temp_file_path = "./tests/"
+    temp_file_root = "temp_sacc_custom_suffix"
+
+    sck.save_concealed_datavector(temp_file_path, temp_file_root, suffix="my_blind")
+
+    expected_path = f"{temp_file_path}{temp_file_root}_my_blind.fits"
+    default_path = f"{temp_file_path}{temp_file_root}_concealed_data_vector.fits"
+
+    assert os.path.exists(expected_path)
+    assert not os.path.exists(default_path)
+
+    os.remove(expected_path)
+
+
+@patch('src.smokescreen.datavector.getpass.getuser', return_value='test_user')
+def test_save_concealed_datavector_default_suffix(mock_getuser):
+    cosmo = COSMO
+    likelihood = "./examples/cosmic_shear/cosmicshear_likelihood.py"
+    syst_dict = {
+        "trc1_delta_z": 0.1,
+        "trc0_delta_z": 0.1,
+    }
+    shift_dict = {"Omega_c": 0.34, "sigma8": 0.85}
+    sacc_data, _ = load_sacc_file("./examples/cosmic_shear/cosmicshear_sacc.fits")
+    sck = ConcealDataVector(cosmo, likelihood, shift_dict, sacc_data, syst_dict, seed=1234)
+    sck.calculate_concealing_factor()
+    sck.apply_concealing_to_likelihood_datavec()
+
+    temp_file_path = "./tests/"
+    temp_file_root = "temp_sacc_default_suffix"
+
+    sck.save_concealed_datavector(temp_file_path, temp_file_root)
+
+    expected_path = f"{temp_file_path}{temp_file_root}_concealed_data_vector.fits"
+    assert os.path.exists(expected_path)
+
+    os.remove(expected_path)

tests/test_main.py

@@ -43,7 +43,7 @@ def test_main(mock_load_sacc, mock_smokescreen, mock_print):
     mock_smokescreen_instance.calculate_concealing_factor.assert_called_once()
     mock_smokescreen_instance.apply_concealing_to_likelihood_datavec.assert_called_once()
     mock_smokescreen_instance.save_concealed_datavector.assert_called_once_with(
-        path_to_output, 'cosmicshear_sacc', output_format='fits'
+        path_to_output, 'cosmicshear_sacc', output_format='fits', suffix=None
     )
 
 
@@ -83,7 +83,7 @@ def test_main_loads_cosmology_from_dict(mock_load_sacc, mock_smokescreen, mock_p
     mock_smokescreen_instance.calculate_concealing_factor.assert_called_once()
     mock_smokescreen_instance.apply_concealing_to_likelihood_datavec.assert_called_once()
     mock_smokescreen_instance.save_concealed_datavector.assert_called_once_with(
-        path_to_output, 'cosmicshear_sacc', output_format='fits'
+        path_to_output, 'cosmicshear_sacc', output_format='fits', suffix=None
     )
 
 
@@ -124,7 +124,42 @@ def test_main_gaussian_shift(mock_load_sacc, mock_smokescreen, mock_print):
     mock_smokescreen_instance.calculate_concealing_factor.assert_called_once()
     mock_smokescreen_instance.apply_concealing_to_likelihood_datavec.assert_called_once()
     mock_smokescreen_instance.save_concealed_datavector.assert_called_once_with(
-        path_to_output, 'cosmicshear_sacc', output_format='fits'
+        path_to_output, 'cosmicshear_sacc', output_format='fits', suffix=None
+    )
+
+
+@patch('builtins.print')
+@patch('smokescreen.__main__.ConcealDataVector')
+@patch('smokescreen.__main__.load_sacc_file')
+def test_datavector_main_custom_suffix(mock_load_sacc, mock_smokescreen, mock_print):
+    # Arrange
+    path_to_sacc = "./examples/cosmic_shear/cosmicshear_sacc.fits"
+    likelihood_path = "./tests/test_data/mock_likelihood.py"
+    systematics = {}
+    shifts_dict = {"Omega_c": [-0.1, 0.2], "sigma8": [-0.1, 0.1]}
+    shift_type = 'add'
+    shift_distribution = 'flat'
+    seed = 2112
+    reference_cosmology = CosmologyVanillaLCDM()
+    path_to_output = "./tests/test_data/"
+    keep_original_sacc = True
+    output_suffix = "my_suffix"
+
+    mock_smokescreen_instance = MagicMock()
+    mock_smokescreen.return_value = mock_smokescreen_instance
+
+    sacc_file = MagicMock()
+    mock_load_sacc.return_value = (sacc_file, 'fits')
+
+    # Act
+    __main__.datavector_main(path_to_sacc, likelihood_path, shifts_dict, systematics,
+                             shift_type, shift_distribution, seed, reference_cosmology,
+                             path_to_output, keep_original_sacc,
+                             output_suffix=output_suffix)
+
+    # Assert
+    mock_smokescreen_instance.save_concealed_datavector.assert_called_once_with(
+        path_to_output, 'cosmicshear_sacc', output_format='fits', suffix="my_suffix"
     )