From 0b7955da9c80309f4f94e8c98fe838cfe7243b5e Mon Sep 17 00:00:00 2001 From: Arthur Loureiro Date: Fri, 13 Mar 2026 10:26:54 +0100 Subject: [PATCH] added suffix to main and datavector.py --- src/smokescreen/__main__.py | 22 +++++++++------ src/smokescreen/datavector.py | 9 ++++-- tests/test_datavector.py | 53 +++++++++++++++++++++++++++++++++++ tests/test_main.py | 41 +++++++++++++++++++++++++-- 4 files changed, 112 insertions(+), 13 deletions(-) diff --git a/src/smokescreen/__main__.py b/src/smokescreen/__main__.py index a3ee13d..5810275 100644 --- a/src/smokescreen/__main__.py +++ b/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="") diff --git a/src/smokescreen/datavector.py b/src/smokescreen/datavector.py index 99200b2..ec8e97a 100644 --- a/src/smokescreen/datavector.py +++ b/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,6 +551,8 @@ 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 ------- @@ -557,6 +560,8 @@ def save_concealed_datavector(self, path_to_save, file_root, 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 diff --git a/tests/test_datavector.py b/tests/test_datavector.py index 1b72345..60b1676 100644 --- a/tests/test_datavector.py +++ b/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) diff --git a/tests/test_main.py b/tests/test_main.py index a4bf65f..952600a 100644 --- a/tests/test_main.py +++ b/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" )