diff --git a/bilby/core/result.py b/bilby/core/result.py index 47a9707cc..98a8c8bd1 100644 --- a/bilby/core/result.py +++ b/bilby/core/result.py @@ -82,7 +82,7 @@ def _determine_file_name(filename, outdir, label, extension, gzip): return result_file_name(outdir, label, extension, gzip) -def read_in_result(filename=None, outdir=None, label=None, extension='json', gzip=False, result_class=None): +def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip=False, result_class=None): """ Reads in a stored bilby result object Parameters @@ -92,12 +92,18 @@ def read_in_result(filename=None, outdir=None, label=None, extension='json', gzi outdir, label, extension: str Name of the output directory, label and extension used for the default naming scheme. + extension: str, optional + The file extension to use. If not given, the extension is inferred from + the filename if provided. If the filename is not given, defaults to + 'json'. result_class: bilby.core.result.Result, or child of The result class to use. By default, `bilby.core.result.Result` is used, but objects which inherit from this class can be given providing additional methods. """ - filename = _determine_file_name(filename, outdir, label, extension, gzip) + filename = _determine_file_name( + filename, outdir, label, extension or "json", gzip + ) if result_class is None: result_class = Result @@ -105,20 +111,48 @@ def read_in_result(filename=None, outdir=None, label=None, extension='json', gzi raise ValueError(f"Input result_class={result_class} not understood") # Get the actual extension (may differ from the default extension if the filename is given) - extension = os.path.splitext(filename)[1].lstrip('.') + if extension is None: + ext = os.path.splitext(filename)[1][1:] + extension = ext if ext else None if extension == 'gz': # gzipped file extension = os.path.splitext(os.path.splitext(filename)[0])[1].lstrip('.') - if 'json' in extension: - result = result_class.from_json(filename=filename) - elif ('hdf5' in extension) or ('h5' in extension): - result = result_class.from_hdf5(filename=filename) - elif ("pkl" in extension) or ("pickle" in extension): - result = result_class.from_pickle(filename=filename) - elif extension is None: - raise ValueError("No filetype extension provided") - else: - raise ValueError("Filetype {} not understood".format(extension)) + if extension is None: + raise ValueError("No filetype extension provided and could not be inferred from filename") + + extension = extension.lower() + read_functions = { + 'json': result_class.from_json, + 'hdf5': result_class.from_hdf5, + 'h5': result_class.from_hdf5, + 'pkl': result_class.from_pickle, + 'pickle': result_class.from_pickle, + } + if extension not in read_functions: + raise ValueError( + f"Filetype {extension} not understood, known types are {list(read_functions.keys())}" + ) + + func = read_functions[extension] + + # Raise IO errors since this is what the caching relies on + # Catch all other exceptions and raise a FileLoadError + try: + result = func(filename=filename) + except IOError as e: + raise IOError( + f"Failed to read in file {filename} using " + f"`{result_class.__name__}.{func.__name__}` " + f"(extension={extension}). " + "This is likely because the file does not exist or is not a valid bilby result." + f"The error was: {e}" + ) from e + except Exception as e: + raise FileLoadError( + f"Failed to read in file {filename} using " + f"`{result_class.__name__}.{func.__name__}` " + f"(extension={extension}). The error was: {e}" + ) from e return result @@ -777,7 +811,8 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, Parameters ========== filename: optional, - Filename to write to (overwrites the default) + Filename to write to (overwrites the default). Assumed to include the + file extension. overwrite: bool, optional Whether or not to overwrite an existing result file. default=False @@ -794,24 +829,39 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, _outdir = None if filename is not None: - _outdir, base_filename = os.path.split(filename) + # Overwrite the full filename with the base filename + # Well append the outdir later on + _outdir, filename = os.path.split(filename) _outdir = None if _outdir == "" else _outdir - base, ext = os.path.splitext(base_filename) - if extension is None: - extension = ext[1:] if ext else "json" - filename = base_filename - elif extension is True: - message = "Result.save_to_file called with extension=True. " - if len(ext) > 0: - message += f"Overwriting extension to json from {ext}, this" - else: - message += "This" - message += " behaviour is deprecated and will be removed." - logger.warning(message) - extension = "json" - filename = base_filename - filename = f"{base}.{extension}" - if extension is None or extension is True: + _, ext = os.path.splitext(filename) + ext = ext[1:] if ext else None + # If the extension has not been set, try to infer it from the filename + # if not, it will fall back to the default + if ext in EXTENSIONS: + if extension is None: + logger.debug( + f"Inferred extension '{ext}' from filename '{filename}'. " + "Using this extension for saving." + ) + extension = ext + elif ext != extension: + message = ( + f"The specified extension '{ext}' " + f"does not match the provided extension '{extension}'. " + ) + logger.warning(message) + + if extension is None: + logger.info("No extension given, defaulting to JSON.") + extension = 'json' + + if extension is True: + message = ( + "Result.save_to_file called with extension=True. " + "This will default to json, and ignore the extension from the filename. " + "This behaviour is deprecated and will be removed. " + ) + logger.warning(message) extension = 'json' outdir = _outdir if outdir is None else outdir @@ -845,13 +895,13 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, else: with open(output_path, 'w') as file: json.dump(dictionary, file, indent=2, cls=BilbyJsonEncoder) - elif extension == 'hdf5': + elif extension in ['hdf5', 'h5']: import h5py dictionary["__module__"] = self.__module__ dictionary["__name__"] = self.__class__.__name__ with h5py.File(output_path, 'w') as h5file: recursively_save_dict_contents_to_group(h5file, '/', dictionary) - elif extension == 'pkl': + elif extension in ['pkl', 'pickle']: safe_file_dump(self, output_path, "dill") else: raise ValueError(f"Extension type {extension} not understood") @@ -2291,3 +2341,7 @@ class ResultListError(ResultError): class FileMovedError(ResultError): """ Exceptions that occur when files have been moved """ + + +class FileLoadError(ResultError): + """ Exceptions that occur when files cannot be loaded """ diff --git a/test/core/result_test.py b/test/core/result_test.py index cbe6ec6d0..23ba8e6b5 100644 --- a/test/core/result_test.py +++ b/test/core/result_test.py @@ -4,11 +4,13 @@ import shutil import os import json +import parameterized import pytest from unittest.mock import patch import bilby -from bilby.core.result import ResultError +from bilby.core.result import ResultError, FileLoadError +from bilby.core.utils import logger class TestJson(unittest.TestCase): @@ -57,6 +59,9 @@ class TestResult(unittest.TestCase): def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory self.outdir = str(tmp_path / "test") + # Outdirs not used in the result object + self.other_outdir = str(tmp_path / "test2") + self.other_outdir2 = str(tmp_path / "test3") def setUp(self): np.random.seed(7) @@ -135,16 +140,25 @@ def test_result_file_name_pickle(self): "{}/{}_result.pkl".format(outdir, label), ) - def test_fail_save_and_load(self): - with self.assertRaises(ValueError): + def test_fail_save_and_load_missing_inputs(self): + with self.assertRaises( + ValueError, msg="No information given to load file" + ): bilby.core.result.read_in_result() - with self.assertRaises(ValueError): + def test_fail_save_and_load_no_extension(self): + with self.assertRaises(ValueError, msg="No file extension provided."): bilby.core.result.read_in_result(filename="no_file_extension") + def test_fail_save_and_load_invalid_extension(self): + with self.assertRaises(ValueError, msg="Filetype .invalid not understood."): + bilby.core.result.read_in_result(filename="file.invalid") + + def test_fail_save_and_load_invalid_file(self): with self.assertRaises(IOError): bilby.core.result.read_in_result(filename="not/a/file.json") + def test_fail_save_and_load_incomplete_json(self): with self.assertRaises(IOError): incomplete_json = """ { @@ -285,14 +299,35 @@ def _save_with_outdir_and_filename(self, filename, outdir, template): self.result.save_to_file(filename=filename, outdir=outdir, extension="json", gzip=False) self.assertTrue(os.path.isfile(template)) - def test_save_with_outdir_and_filename(self): - self._save_with_outdir_and_filename("out/result", "out2", "out2/result.json") - self._save_with_outdir_and_filename("out/result", None, "out/result.json") - self._save_with_outdir_and_filename("result", "out", "out/result.json") + def test_save_with_outdir_and_filename_different_outdir(self): + self._save_with_outdir_and_filename( + f"{self.other_outdir}/result", self.other_outdir2, f"{self.other_outdir2}/result" + ) + + def test_save_with_outdir_and_filename_same_outdir(self): + self._save_with_outdir_and_filename( + f"{self.other_outdir}/result", None, f"{self.other_outdir}/result" + ) + + def test_save_with_outdir_and_filename_no_outdir_in_filename(self): + self._save_with_outdir_and_filename( + "result", self.other_outdir, f"{self.other_outdir}/result" + ) + + def test_save_with_filename_only(self): self._save_with_outdir_and_filename( - "result", None, os.path.join(self.result.outdir, "result.json")) + "result", None, os.path.join(self.result.outdir, "result") + ) + + def test_save_with_outdir_no_filename(self): self._save_with_outdir_and_filename( - None, "out", os.path.join("out", f"{self.result.label}_result.json")) + None, self.other_outdir, os.path.join(self.other_outdir, f"{self.result.label}_result.json") + ) + + def test_save_no_filename_or_outdir(self): + self._save_with_outdir_and_filename( + None, None, os.path.join(self.result.outdir, f"{self.result.label}_result.json") + ) def test_save_and_overwrite_json(self): self._save_and_overwrite_test(extension='json') @@ -564,21 +599,25 @@ class NotAResult(bilby.core.result.Result): check_point_plot=False, result_class=NotAResult ) - # result should be specified result_class assert isinstance(result, NotAResult) - cached_result = bilby.run_sampler( - likelihood, - priors, - sampler='bilby_mcmc', - outdir=self.outdir, - nsamples=10, - L1steps=1, - proposal_cycle="default_noGMnoKD", - printdt=1, - check_point_plot=False, - result_class=NotAResult - ) + # Due to a quirk with how the logger is configured (propagate=False) + # It's easier to just patch the logger and check the call exists + with patch("bilby.core.sampler.logger") as mock_logger: + cached_result = bilby.run_sampler( + likelihood, + priors, + sampler='bilby_mcmc', + outdir=self.outdir, + nsamples=10, + L1steps=1, + proposal_cycle="default_noGMnoKD", + printdt=1, + check_point_plot=False, + result_class=NotAResult, + clean=False, + ) + mock_logger.warning.assert_any_call("Using cached result") # so should a result loaded from cache assert isinstance(cached_result, NotAResult) @@ -777,6 +816,11 @@ def test_sanity_check_labels(self): class TestPPPlots(unittest.TestCase): + @pytest.fixture(autouse=True) + def init_outdir(self, tmp_path): + # Use pytest's tmp_path fixture to create a temporary directory + self.outdir = str(tmp_path / "test_pp_plots") + def setUp(self): priors = bilby.core.prior.PriorDict(dict( a=bilby.core.prior.Uniform(0, 1, latex_label="$a$"), @@ -785,7 +829,7 @@ def setUp(self): self.results = [ bilby.core.result.Result( label=str(ii), - outdir='.', + outdir=self.outdir, search_parameter_keys=list(priors.keys()), priors=priors, injection_parameters=priors.sample(), @@ -874,38 +918,86 @@ def test_reweight_different_likelihood_weights_correct(self): self.assertLess(abs(new.log_evidence - self.result.log_evidence), 0.05) self.assertNotEqual(new.log_evidence, self.result.log_evidence) - def test_save_to_file_filename_with_extension_and_extension_none(self): + +class TestResultSaveAndRead(unittest.TestCase): + + @pytest.fixture(autouse=True) + def init_outdir(self, tmp_path): + # Use pytest's tmp_path fixture to create a temporary directory + self.outdir = str(tmp_path / "test_result_save_and_read") + + def setUp(self): + np.random.seed(7) + bilby.utils.command_line_args.bilby_test_mode = False + priors = bilby.prior.PriorDict( + dict( + x=bilby.prior.Uniform(0, 1, "x", latex_label="$x$", unit="s"), + y=bilby.prior.Uniform(0, 1, "y", latex_label="$y$", unit="m"), + c=1, + d=2, + ) + ) + result = bilby.core.result.Result( + label="label", + outdir=self.outdir, + sampler="nestle", + search_parameter_keys=["x", "y"], + fixed_parameter_keys=["c", "d"], + priors=priors, + sampler_kwargs=dict(test="test", func=lambda x: x), + injection_parameters=dict(x=0.5, y=0.5), + meta_data=dict(test="test"), + sampling_time=100.0, + ) + + n = 100 + posterior = pd.DataFrame( + dict(x=np.random.normal(0, 1, n), y=np.random.normal(0, 1, n)) + ) + result.posterior = posterior + result.log_evidence = 10 + result.log_evidence_err = 11 + result.log_bayes_factor = 12 + result.log_noise_evidence = 13 + self.result = result + + @parameterized.parameterized.expand([ + ".h5", ".hdf5", ".json", ".pkl", ".pickle", + ]) + def test_save_and_read_filename_with_extension_and_extension_none(self, ext): # Should use the extension from filename - filename = os.path.join(self.result.outdir, "custom_name.hdf5") + filename = os.path.join(self.result.outdir, f"custom_name.{ext}") self.result.save_to_file(filename=filename, extension=None) self.assertTrue(os.path.isfile(filename)) + bilby.core.result.read_in_result(filename=filename) os.remove(filename) - def test_save_to_file_filename_with_extension_and_extension_true(self): - """This is a strange default, that we should remove, but is here for consistency""" - filename = os.path.join(self.result.outdir, "custom_name.hdf5") - expected = os.path.join(self.result.outdir, "custom_name.json") - self.result.save_to_file(filename=filename, extension=True) - self.assertTrue(os.path.isfile(expected)) - self.assertFalse(os.path.isfile(filename)) - os.remove(expected) - - def test_save_to_file_filename_with_extension_and_extension_set(self): - # Should override the extension in filename with the one provided in extension + @parameterized.parameterized.expand([ + ("json",), + ("pkl",), + ("pickle",), + (True,), + ]) + def test_save_and_read_filename_with_extension_and_extension(self, extension): + """Test all the extensions that are support when the filename is provided""" filename = os.path.join(self.result.outdir, "custom_name.hdf5") - expected = os.path.join(self.result.outdir, "custom_name.json") - self.result.save_to_file(filename=filename, extension="json") + expected = filename + with self.assertLogs(logger, level='WARNING') as cm: + self.result.save_to_file(filename=filename, extension=extension) + self.assertIn("does not match the provided extension", cm.output[0]) self.assertTrue(os.path.isfile(expected)) - self.assertFalse(os.path.isfile(filename)) + if extension is True: + extension = "json" + bilby.core.result.read_in_result(filename=expected, extension=extension) os.remove(expected) - def test_save_to_file_filename_without_extension_and_extension_none(self): + def test_save_and_read_filename_without_extension_and_extension_none(self): # Should use the default extension (json) filename = os.path.join(self.result.outdir, "custom_name_noext") - expected = filename + ".json" + expected = filename self.result.save_to_file(filename=filename, extension=None) self.assertTrue(os.path.isfile(expected)) - self.assertFalse(os.path.isfile(filename)) + bilby.core.result.read_in_result(filename=expected, extension="json") os.remove(expected) def test_save_to_file_defaults_to_pickle_with_incorrect_extension(self): @@ -915,8 +1007,27 @@ def test_save_to_file_defaults_to_pickle_with_incorrect_extension(self): self.result.save_to_file(filename=filename, extension="bar") self.assertTrue(os.path.isfile(expected)) self.assertFalse(os.path.isfile(filename)) + bilby.core.result.read_in_result(filename=expected) os.remove(expected) + @parameterized.parameterized.expand([ + ("json", "hdf5"), + ("json", "pkl"), + ("hdf5", "json"), + ("pkl", "json"), + ("json", "pkl"), + ("hdf5", "pkl"), + ]) + def test_save_and_read_incorrect_extension(self, save_extension, read_extension): + """Test that an incorrect extension raises a somewhat helpful error""" + filename = os.path.join(self.result.outdir, "my_result") + self.result.save_to_file(filename=filename, extension=save_extension) + with self.assertRaises( + (FileLoadError, IOError), msg=f"Failed to read in file {filename}" + ): + bilby.core.result.read_in_result(filename=filename, extension=read_extension) + os.remove(filename) + if __name__ == "__main__": unittest.main()