From dcda53e1e19de1a29480b415939ace6eb63ca0e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A5vard=20Berland?= Date: Wed, 7 Oct 2020 10:02:50 +0200 Subject: [PATCH] Catch KeyError when parameters is requested as class property (#169) --- src/fmu/ensemble/ensemble.py | 16 ++++++++++--- src/fmu/ensemble/ensemblecombination.py | 5 ++++- src/fmu/ensemble/ensembleset.py | 30 +++++++++++++++++++++++-- tests/test_ensemble.py | 24 ++++++++++++++++++++ tests/test_ensembleset.py | 29 ++++++++++++++++++++++++ 5 files changed, 98 insertions(+), 6 deletions(-) diff --git a/src/fmu/ensemble/ensemble.py b/src/fmu/ensemble/ensemble.py index ebd55e8a..8f627a76 100644 --- a/src/fmu/ensemble/ensemble.py +++ b/src/fmu/ensemble/ensemble.py @@ -292,7 +292,7 @@ def remove_data(self, localpaths): on each realization in the ensemble. Args: - localpath (string): Full localpath to + localpaths (string): Full localpaths to the data, or list of strings. """ if isinstance(localpaths, str): @@ -412,8 +412,18 @@ def manifest(self, manifest): @property def parameters(self): - """Getter for get_parameters(convert_numeric=True)""" - return self.load_txt("parameters.txt") + """Build a dataframe of the information in each + realizations parameters.txt. + + If no realizations have the file, an empty dataframe is returned. + + Returns: + pd.DataFrame + """ + try: + return self.load_txt("parameters.txt") + except KeyError: + return pd.DataFrame() def load_scalar(self, localpath, convert_numeric=False, force_reread=False): """Parse a single value from a file for each realization. diff --git a/src/fmu/ensemble/ensemblecombination.py b/src/fmu/ensemble/ensemblecombination.py index 7e82498d..2a371934 100644 --- a/src/fmu/ensemble/ensemblecombination.py +++ b/src/fmu/ensemble/ensemblecombination.py @@ -331,7 +331,10 @@ def get_volumetric_rates( @property def parameters(self): """Return parameters from the ensemble as a class property""" - return self.get_df("parameters.txt") + try: + return self.get_df("parameters.txt") + except KeyError: + return pd.DataFrame() def __len__(self): """Estimate the number of realizations in this diff --git a/src/fmu/ensemble/ensembleset.py b/src/fmu/ensemble/ensembleset.py index a34e679d..9ba99176 100644 --- a/src/fmu/ensemble/ensembleset.py +++ b/src/fmu/ensemble/ensembleset.py @@ -336,8 +336,18 @@ def add_ensemble(self, ensembleobject): @property def parameters(self): - """Getter for ensemble.parameters(convert_numeric=True)""" - return self.get_df("parameters.txt") + """Build a dataframe of the information in each + realizations parameters.txt. + + If no realizations have the file, an empty dataframe is returned. + + Returns: + pd.DataFrame + """ + try: + return self.get_df("parameters.txt") + except KeyError: + return pd.DataFrame() def load_scalar(self, localpath, convert_numeric=False, force_reread=False): """Parse a single value from a file @@ -438,6 +448,22 @@ def drop(self, localpath, **kwargs): except ValueError: pass # Allow localpath to be missing in some ensembles. + def remove_data(self, localpaths): + """Remove certain datatypes from each ensembles/realizations + datastores. This modifies the underlying realization + objects, and is equivalent to + + >>> del realization[localpath] + + on each realization in each ensemble. + + Args: + localpaths (string): Full localpath to + the data, or list of strings. + """ + for _, ensemble in self._ensembles.items(): + ensemble.remove_data(localpaths) + def process_batch(self, batch=None): """Process a list of functions to run/apply diff --git a/tests/test_ensemble.py b/tests/test_ensemble.py index 723034cd..6252b917 100644 --- a/tests/test_ensemble.py +++ b/tests/test_ensemble.py @@ -180,6 +180,30 @@ def test_reek001(tmpdir): assert len(reekensemble.keys()) == keycount - 1 +def test_noparameters(): + """Test what happens when parameters.txt is missing""" + + testdir = os.path.dirname(os.path.abspath(__file__)) + reekensemble = ScratchEnsemble( + "reektest", testdir + "/data/testensemble-reek001/" + "realization-*/iter-0" + ) + # Parameters.txt exist on disk, so it is loaded: + assert not reekensemble.parameters.empty + # Remove it each realization: + reekensemble.remove_data("parameters.txt") + assert reekensemble.parameters.empty + + # However, when parameters.txt is excplicitly asked for, + # an exception should be raised: + with pytest.raises(KeyError): + reekensemble.get_df("parameters.txt") + + reekensemble.load_smry(time_index="yearly", column_keys="FOPT") + assert not reekensemble.get_df("unsmry--yearly").empty + with pytest.raises(KeyError): + reekensemble.get_df("unsmry--yearly", merge="parameters.txt") + + def test_emptyens(): """Check that we can initialize an empty ensemble""" ens = ScratchEnsemble("emptyens") diff --git a/tests/test_ensembleset.py b/tests/test_ensembleset.py index d2632236..b352f68c 100644 --- a/tests/test_ensembleset.py +++ b/tests/test_ensembleset.py @@ -281,6 +281,35 @@ def rms_vol2df(kwargs): assert len(ensset5.get_df("unsmry--daily")) == 10980 +def test_noparameters(tmpdir): + testdir = os.path.dirname(os.path.abspath(__file__)) + ensdir = os.path.join(testdir, "data/testensemble-reek001/") + + tmpdir.chdir() + symlink_iter(ensdir, "iter-0") + symlink_iter(ensdir, "iter-1") + + iter0 = ScratchEnsemble("iter-0", str(tmpdir.join("realization-*/iter-0"))) + iter1 = ScratchEnsemble("iter-1", str(tmpdir.join("realization-*/iter-1"))) + + ensset = EnsembleSet("reek001", [iter0, iter1]) + assert not ensset.parameters.empty + + # Remove it each realization: + ensset.remove_data("parameters.txt") + assert ensset.parameters.empty + + # However, when parameters.txt is excplicitly asked for, + # an exception should be raised: + with pytest.raises(KeyError): + ensset.get_df("parameters.txt") + + ensset.load_smry(time_index="yearly", column_keys="FOPT") + assert not ensset.get_df("unsmry--yearly").empty + with pytest.raises(KeyError): + ensset.get_df("unsmry--yearly", merge="parameters.txt") + + def test_pred_dir(tmpdir): """Test import of a stripped 5 realization ensemble, manually doubled to two identical ensembles,