From 88b1d4bac2e7771ce3163dcf6b6c7535463311fc Mon Sep 17 00:00:00 2001 From: Anna Kvashchuk Date: Wed, 13 Dec 2023 10:23:45 +0100 Subject: [PATCH] Change load_parameters function to return only Dataset --- src/ert/analysis/_es_update.py | 20 ++++++----- src/ert/config/ext_param_config.py | 2 +- src/ert/config/field.py | 2 +- src/ert/config/gen_kw_config.py | 2 +- src/ert/config/surface_config.py | 2 +- src/ert/libres_facade.py | 6 ++-- src/ert/storage/local_ensemble.py | 12 ++----- .../analysis/test_adaptive_localization.py | 4 +-- tests/unit_tests/analysis/test_es_update.py | 33 ++++++++++------- tests/unit_tests/cli/test_integration_cli.py | 8 +++-- .../migration/test_block_fs_simple_case.py | 8 ++--- .../storage/test_field_parameter.py | 34 ++++++++++-------- .../unit_tests/storage/test_local_ensemble.py | 4 +-- .../storage/test_parameter_sample_types.py | 36 +++++++++++-------- 14 files changed, 95 insertions(+), 78 deletions(-) diff --git a/src/ert/analysis/_es_update.py b/src/ert/analysis/_es_update.py index 8e07c3fc8c7..6d5b1965690 100644 --- a/src/ert/analysis/_es_update.py +++ b/src/ert/analysis/_es_update.py @@ -268,15 +268,17 @@ def _create_temporary_parameter_storage( matrix: Union[npt.NDArray[np.double], xr.DataArray] if isinstance(config_node, GenKwConfig): t = time.perf_counter() - matrix = source_fs.load_parameters(param_group, iens_active_index).values.T # type: ignore + matrix = source_fs.load_parameters(param_group, iens_active_index)[ + "values" + ].values.T t_genkw += time.perf_counter() - t elif isinstance(config_node, SurfaceConfig): t = time.perf_counter() - matrix = source_fs.load_parameters(param_group, iens_active_index) # type: ignore + matrix = source_fs.load_parameters(param_group, iens_active_index)["values"] t_surface += time.perf_counter() - t elif isinstance(config_node, Field): t = time.perf_counter() - matrix = source_fs.load_parameters(param_group, iens_active_index) # type: ignore + matrix = source_fs.load_parameters(param_group, iens_active_index)["values"] t_field += time.perf_counter() - t else: raise NotImplementedError(f"{type(config_node)} is not supported") @@ -679,9 +681,7 @@ def analysis_ES( ) for parameter_group in not_updated_parameter_groups: for realization in iens_active_index: - ds = source_fs.load_parameters( - parameter_group, int(realization), var=None - ) + ds = source_fs.load_parameters(parameter_group, int(realization)) assert isinstance(ds, xr.Dataset) target_fs.save_parameters( parameter_group, @@ -813,7 +813,9 @@ def analysis_IES( updated_parameter_groups.append(param_group.name) source: Union[EnsembleReader, EnsembleAccessor] = target_fs try: - target_fs.load_parameters(group=param_group.name, realizations=0) + target_fs.load_parameters(group=param_group.name, realizations=0)[ + "values" + ] except Exception: source = source_fs temp_storage = _create_temporary_parameter_storage( @@ -843,8 +845,8 @@ def analysis_IES( for parameter_group in not_updated_parameter_groups: for realization in iens_active_index: prior_dataset = source_fs.load_parameters( - parameter_group, int(realization), var=None - ) + parameter_group, int(realization) + )["values"] assert isinstance(prior_dataset, xr.Dataset) target_fs.save_parameters( parameter_group, diff --git a/src/ert/config/ext_param_config.py b/src/ert/config/ext_param_config.py index 07d1bd58ed5..54d72d22ea3 100644 --- a/src/ert/config/ext_param_config.py +++ b/src/ert/config/ext_param_config.py @@ -74,7 +74,7 @@ def write_to_runpath( Path.mkdir(file_path.parent, exist_ok=True, parents=True) data: MutableDataType = {} - for da in ensemble.load_parameters(self.name, real_nr): + for da in ensemble.load_parameters(self.name, real_nr)["values"]: assert isinstance(da, xr.DataArray) name = str(da.names.values) try: diff --git a/src/ert/config/field.py b/src/ert/config/field.py index bd775acc698..815fa563dc2 100644 --- a/src/ert/config/field.py +++ b/src/ert/config/field.py @@ -180,7 +180,7 @@ def write_to_runpath( def _fetch_from_ensemble( self, real_nr: int, ensemble: EnsembleReader ) -> xr.DataArray: - da = ensemble.load_parameters(self.name, real_nr) + da = ensemble.load_parameters(self.name, real_nr)["values"] assert isinstance(da, xr.DataArray) return da diff --git a/src/ert/config/gen_kw_config.py b/src/ert/config/gen_kw_config.py index cea35a3d002..4a2c55445ab 100644 --- a/src/ert/config/gen_kw_config.py +++ b/src/ert/config/gen_kw_config.py @@ -225,7 +225,7 @@ def read_from_runpath( def write_to_runpath( self, run_path: Path, real_nr: int, ensemble: EnsembleReader ) -> Dict[str, Dict[str, float]]: - array = ensemble.load_parameters(self.name, real_nr, var="transformed_values") + array = ensemble.load_parameters(self.name, real_nr)["transformed_values"] assert isinstance(array, xr.DataArray) if not array.size == len(self.transfer_functions): raise ValueError( diff --git a/src/ert/config/surface_config.py b/src/ert/config/surface_config.py index fdd6d4db4c8..bd1a929e382 100644 --- a/src/ert/config/surface_config.py +++ b/src/ert/config/surface_config.py @@ -119,7 +119,7 @@ def read_from_runpath(self, run_path: Path, real_nr: int) -> xr.Dataset: def write_to_runpath( self, run_path: Path, real_nr: int, ensemble: EnsembleReader ) -> None: - data = ensemble.load_parameters(self.name, real_nr) + data = ensemble.load_parameters(self.name, real_nr)["values"] surf = xtgeo.RegularSurface( ncol=self.ncol, diff --git a/src/ert/libres_facade.py b/src/ert/libres_facade.py index 1f07d744c97..4c5a4927afd 100644 --- a/src/ert/libres_facade.py +++ b/src/ert/libres_facade.py @@ -298,9 +298,9 @@ def load_all_gen_kw_data( gen_kws = [config for config in gen_kws if config.name == group] for key in gen_kws: try: - ds = ensemble.load_parameters( - key.name, realizations, var="transformed_values" - ) + ds = ensemble.load_parameters(key.name, realizations)[ + "transformed_values" + ] assert isinstance(ds, xr.DataArray) ds["names"] = np.char.add(f"{key.name}:", ds["names"].astype(np.str_)) df = ds.to_dataframe().unstack(level="names") diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index a1b8e8ce500..c4423536afc 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -203,15 +203,9 @@ def has_parameter_group(self, group: str) -> bool: return param_group_file.exists() def load_parameters( - self, - group: str, - realizations: Union[int, npt.NDArray[np.int_], None] = None, - *, - var: Optional[str] = "values", - ) -> Union[xr.DataArray, xr.Dataset]: - if var is None: - return self._load_dataset(group, realizations) - return self._load_dataset(group, realizations)[var] + self, group: str, realizations: Union[int, npt.NDArray[np.int_], None] = None + ) -> xr.Dataset: + return self._load_dataset(group, realizations) @lru_cache # noqa: B019 def load_responses( diff --git a/tests/unit_tests/analysis/test_adaptive_localization.py b/tests/unit_tests/analysis/test_adaptive_localization.py index 62dbcde8151..726f0875f68 100644 --- a/tests/unit_tests/analysis/test_adaptive_localization.py +++ b/tests/unit_tests/analysis/test_adaptive_localization.py @@ -37,9 +37,9 @@ def run_cli_ES_with_case(poly_config): storage_path = ErtConfig.from_file(poly_config).ens_path with open_storage(storage_path) as storage: prior_ensemble = storage.get_ensemble_by_name(prior_sample_name) - prior_sample = prior_ensemble.load_parameters("COEFFS") + prior_sample = prior_ensemble.load_parameters("COEFFS")["values"] posterior_ensemble = storage.get_ensemble_by_name(posterior_sample_name) - posterior_sample = posterior_ensemble.load_parameters("COEFFS") + posterior_sample = posterior_ensemble.load_parameters("COEFFS")["values"] return prior_sample, posterior_sample diff --git a/tests/unit_tests/analysis/test_es_update.py b/tests/unit_tests/analysis/test_es_update.py index 19c59f624a4..54a5f917535 100644 --- a/tests/unit_tests/analysis/test_es_update.py +++ b/tests/unit_tests/analysis/test_es_update.py @@ -272,10 +272,12 @@ def test_update_snapshot( rng=rng, ) - sim_gen_kw = list(prior_ens.load_parameters("SNAKE_OIL_PARAM", 0).values.flatten()) + sim_gen_kw = list( + prior_ens.load_parameters("SNAKE_OIL_PARAM", 0)["values"].values.flatten() + ) target_gen_kw = list( - posterior_ens.load_parameters("SNAKE_OIL_PARAM", 0).values.flatten() + posterior_ens.load_parameters("SNAKE_OIL_PARAM", 0)["values"].values.flatten() ) # Check that prior is not equal to posterior after updationg @@ -414,10 +416,12 @@ def test_localization( rng=np.random.default_rng(42), ) - sim_gen_kw = list(prior_ens.load_parameters("SNAKE_OIL_PARAM", 0).values.flatten()) + sim_gen_kw = list( + prior_ens.load_parameters("SNAKE_OIL_PARAM", 0)["values"].values.flatten() + ) target_gen_kw = list( - posterior_ens.load_parameters("SNAKE_OIL_PARAM", 0).values.flatten() + posterior_ens.load_parameters("SNAKE_OIL_PARAM", 0)["values"].values.flatten() ) # Test that the localized values has been updated @@ -589,14 +593,16 @@ def sample_prior(nx, ny): ens_posterior = storage.get_ensemble_by_name("es_udpate") # Check that surfaces defined in INIT_FILES are not changed by ERT - surf_prior = ens_prior.load_parameters("TOP", list(range(ensemble_size))) + surf_prior = ens_prior.load_parameters("TOP", list(range(ensemble_size)))["values"] for i in range(ensemble_size): _prior_init = xtgeo.surface_from_file( f"surface/surf_init_{i}.irap", fformat="irap_ascii", dtype=np.float32 ) np.testing.assert_array_equal(surf_prior[i], _prior_init.values.data) - surf_posterior = ens_posterior.load_parameters("TOP", list(range(ensemble_size))) + surf_posterior = ens_posterior.load_parameters("TOP", list(range(ensemble_size)))[ + "values" + ] assert surf_prior.shape == surf_posterior.shape @@ -643,7 +649,7 @@ def _load_parameters(source_ens, iens_active_index, param_groups): temp_storage[param_group] = _temp_storage[param_group] return temp_storage - sim_fs.load_parameters("SNAKE_OIL_PARAM_BPR") + sim_fs.load_parameters("SNAKE_OIL_PARAM_BPR")["values"] param_groups = list(sim_fs.experiment.parameter_configuration.keys()) prior = _load_parameters(sim_fs, list(range(10)), param_groups) posterior = _load_parameters(posterior_fs, list(range(10)), param_groups) @@ -799,11 +805,12 @@ def g(X): ) benchmark(smoother_update_run) - prior_da = prior.load_parameters(param_group, range(num_ensemble)) - posterior_da = posterior_ens.load_parameters(param_group, range(num_ensemble)) + prior_da = prior.load_parameters(param_group, range(num_ensemble))["values"] + posterior_da = posterior_ens.load_parameters(param_group, range(num_ensemble))["values"] # Make sure some, but not all parameters were updated. assert not np.allclose(prior_da, posterior_da) # All parameters would be updated with a global update so this would fail. + assert np.isclose(prior_da, posterior_da).sum() > 0 @@ -1046,9 +1053,9 @@ def test_update_subset_parameters(storage, uniform_parameter, obs): smoother_update( prior, posterior_ens, "id", update_config, UpdateSettings(), ESSettings() ) - assert prior.load_parameters("EXTRA_PARAMETER", 0).equals( - posterior_ens.load_parameters("EXTRA_PARAMETER", 0) + assert prior.load_parameters("EXTRA_PARAMETER", 0)["values"].equals( + posterior_ens.load_parameters("EXTRA_PARAMETER", 0)["values"] ) - assert not prior.load_parameters("PARAMETER", 0).equals( - posterior_ens.load_parameters("PARAMETER", 0) + assert not prior.load_parameters("PARAMETER", 0)["values"].equals( + posterior_ens.load_parameters("PARAMETER", 0)["values"] ) diff --git a/tests/unit_tests/cli/test_integration_cli.py b/tests/unit_tests/cli/test_integration_cli.py index b3b28efe6f0..b535bd282ff 100644 --- a/tests/unit_tests/cli/test_integration_cli.py +++ b/tests/unit_tests/cli/test_integration_cli.py @@ -409,7 +409,9 @@ def test_that_prior_is_not_overwritten_in_ensemble_experiment( experiment_id, name="iter-0", ensemble_size=num_realizations ) sample_prior(ensemble, prior_mask) - prior_values = storage.get_ensemble(ensemble.id).load_parameters("COEFFS") + prior_values = storage.get_ensemble(ensemble.id).load_parameters("COEFFS")[ + "values" + ] storage.close() parser = ArgumentParser(prog="test_main") @@ -427,7 +429,9 @@ def test_that_prior_is_not_overwritten_in_ensemble_experiment( FeatureToggling.update_from_args(parsed) run_cli(parsed) storage = open_storage(ert_config.ens_path, mode="w") - parameter_values = storage.get_ensemble(ensemble.id).load_parameters("COEFFS") + parameter_values = storage.get_ensemble(ensemble.id).load_parameters("COEFFS")[ + "values" + ] if should_resample: with pytest.raises(AssertionError): diff --git a/tests/unit_tests/storage/migration/test_block_fs_simple_case.py b/tests/unit_tests/storage/migration/test_block_fs_simple_case.py index eeac40e5687..9251c0e3210 100644 --- a/tests/unit_tests/storage/migration/test_block_fs_simple_case.py +++ b/tests/unit_tests/storage/migration/test_block_fs_simple_case.py @@ -71,7 +71,7 @@ def test_migrate_surface(data, storage, parameter, ens_config): for key, var in data["/REAL_0/SURFACE"].groups.items(): expect = sorted_surface(var) - actual = ensemble.load_parameters(key, 0).values.ravel() + actual = ensemble.load_parameters(key, 0)["values"].values.ravel() assert list(expect) == list(actual), key @@ -84,7 +84,7 @@ def test_migrate_field(data, storage, parameter, ens_config): for key, var in data["/REAL_0/FIELD"].groups.items(): expect = np.array(var["VALUES"]).ravel() - actual = ensemble.load_parameters(key, [0]).values.ravel() + actual = ensemble.load_parameters(key, [0])["values"].values.ravel() assert list(expect) == list(actual), key @@ -101,13 +101,13 @@ def test_migrate_case(data, storage, enspath): # Compare FIELDs for key, data in real_group["FIELD"].groups.items(): expect = np.array(data["VALUES"]).ravel() - actual = ensemble.load_parameters(key, [real_index]) + actual = ensemble.load_parameters(key, [real_index])["values"] assert list(expect) == list(actual.values.ravel()), f"FIELD {key}" # Compare SURFACEs for key, data in real_group["SURFACE"].groups.items(): expect = sorted_surface(data) - actual = ensemble.load_parameters(key, real_index).values.ravel() + actual = ensemble.load_parameters(key, real_index)["values"].values.ravel() assert list(expect) == list(actual), f"SURFACE {key}" diff --git a/tests/unit_tests/storage/test_field_parameter.py b/tests/unit_tests/storage/test_field_parameter.py index accea22cddf..12b10e3ef8a 100644 --- a/tests/unit_tests/storage/test_field_parameter.py +++ b/tests/unit_tests/storage/test_field_parameter.py @@ -116,12 +116,12 @@ def test_load_two_parameters_forward_init(storage, tmpdir): with pytest.raises( KeyError, match="No dataset 'PARAM_A' in storage for realization 0" ): - _ = fs.load_parameters("PARAM_A", [0]) + _ = fs.load_parameters("PARAM_A", [0])["values"] with pytest.raises( KeyError, match="No dataset 'PARAM_B' in storage for realization 0" ): - _ = fs.load_parameters("PARAM_B", [0]) + _ = fs.load_parameters("PARAM_B", [0])["values"] assert load_from_forward_model("config.ert", fs, 0) == 1 @@ -142,10 +142,10 @@ def test_load_two_parameters_forward_init(storage, tmpdir): numpy.testing.assert_equal(prop_b.values.data, param_b) # should be loaded now - loaded_a = fs.load_parameters("PARAM_A", [0]) + loaded_a = fs.load_parameters("PARAM_A", [0])["values"] assert (loaded_a.values == 22).all() - loaded_b = fs.load_parameters("PARAM_B", [0]) + loaded_b = fs.load_parameters("PARAM_B", [0])["values"] assert (loaded_b.values == 77).all() @@ -176,10 +176,10 @@ def test_load_two_parameters_roff(storage, tmpdir): assert not ensemble_config["PARAM_A"].forward_init assert not ensemble_config["PARAM_B"].forward_init - loaded_a = fs.load_parameters("PARAM_A", [0]) + loaded_a = fs.load_parameters("PARAM_A", [0])["values"] assert (loaded_a.values == 22).all() - loaded_b = fs.load_parameters("PARAM_B", [0]) + loaded_b = fs.load_parameters("PARAM_B", [0])["values"] assert (loaded_b.values == 77).all() prop_a = xtgeo.gridproperty_from_file( @@ -232,10 +232,10 @@ def test_load_two_parameters(storage, tmpdir): assert not ensemble_config["PARAM_A"].forward_init assert not ensemble_config["PARAM_B"].forward_init - loaded_a = fs.load_parameters("PARAM_A", [0]) + loaded_a = fs.load_parameters("PARAM_A", [0])["values"] assert (loaded_a.values == 22).all() - loaded_b = fs.load_parameters("PARAM_B", [0]) + loaded_b = fs.load_parameters("PARAM_B", [0])["values"] assert (loaded_b.values == 77).all() prop_a = xtgeo.gridproperty_from_file( @@ -350,7 +350,7 @@ def test_transformation(storage, tmpdir): _, fs = create_runpath(storage, "config.ert", [True, True]) # stored internally as 2.5, 1.5 - loaded_a = fs.load_parameters("PARAM_A", [0, 1]) + loaded_a = fs.load_parameters("PARAM_A", [0, 1])["values"] assert np.isclose(loaded_a.values[0], 2.5).all() assert np.isclose(loaded_a.values[1], 1.5).all() @@ -419,7 +419,7 @@ def test_forward_init(storage, tmpdir, config_str, expect_forward_init): with pytest.raises( KeyError, match="No dataset 'MY_PARAM' in storage for realization 0" ): - fs.load_parameters("MY_PARAM", [0]) + fs.load_parameters("MY_PARAM", [0])["values"] # We try to load the parameters from the forward model, this would fail if # forward init was not set correctly @@ -437,7 +437,7 @@ def test_forward_init(storage, tmpdir, config_str, expect_forward_init): numpy.testing.assert_equal(prop.values.data, expect_param) if expect_forward_init: - arr = fs.load_parameters("MY_PARAM", [0]) + arr = fs.load_parameters("MY_PARAM", [0])["values"] assert len(arr.values.ravel()) == 16 @@ -547,12 +547,14 @@ def test_field_param_update(tmpdir): prior = storage.get_ensemble_by_name("prior") posterior = storage.get_ensemble_by_name("smoother_update") - prior_result = prior.load_parameters("MY_PARAM", list(range(5))) + prior_result = prior.load_parameters("MY_PARAM", list(range(5)))["values"] assert len(prior_result.x) == NCOL assert len(prior_result.y) == NROW assert len(prior_result.z) == NLAY - posterior_result = posterior.load_parameters("MY_PARAM", list(range(5))) + posterior_result = posterior.load_parameters("MY_PARAM", list(range(5)))[ + "values" + ] # Only assert on the first three rows, as there are only three parameters, # a, b and c, the rest have no correlation to the results. assert np.linalg.det( @@ -674,8 +676,10 @@ def test_parameter_update_with_inactive_cells_xtgeo_grdecl(tmpdir): prior = storage.get_ensemble_by_name("prior") posterior = storage.get_ensemble_by_name("smoother_update") - prior_result = prior.load_parameters("MY_PARAM", list(range(5))) - posterior_result = posterior.load_parameters("MY_PARAM", list(range(5))) + prior_result = prior.load_parameters("MY_PARAM", list(range(5)))["values"] + posterior_result = posterior.load_parameters("MY_PARAM", list(range(5)))[ + "values" + ] # check the shape of internal data used in the update assert prior_result.shape == (5, NCOL, NROW, NLAY) diff --git a/tests/unit_tests/storage/test_local_ensemble.py b/tests/unit_tests/storage/test_local_ensemble.py index 415c2f5029e..31270915b12 100644 --- a/tests/unit_tests/storage/test_local_ensemble.py +++ b/tests/unit_tests/storage/test_local_ensemble.py @@ -51,7 +51,7 @@ def test_that_egrid_files_are_saved_and_loaded_correctly(tmp_path): ds = da.to_dataset() ensemble.save_parameters("MY_PARAM", 1, ds) assert (ensemble_dir / "realization-1" / "MY_PARAM.nc").exists() - loaded_data = ensemble.load_parameters("MY_PARAM", 1) + loaded_data = ensemble.load_parameters("MY_PARAM", 1)["values"] np.testing.assert_array_equal(loaded_data.values, data.reshape((4, 5, 1))) @@ -95,7 +95,7 @@ def test_that_grid_files_are_saved_and_loaded_correctly(tmp_path): saved_file = ensemble_dir / "realization-1" / f"{param_group}.nc" assert saved_file.exists() - loaded_data = ensemble.load_parameters(param_group, 1) + loaded_data = ensemble.load_parameters(param_group, 1)["values"] np.testing.assert_array_equal( loaded_data.values, data.reshape((grid.nx, grid.ny, grid.nz)) ) diff --git a/tests/unit_tests/storage/test_parameter_sample_types.py b/tests/unit_tests/storage/test_parameter_sample_types.py index 22b0d7f5748..51a35892712 100644 --- a/tests/unit_tests/storage/test_parameter_sample_types.py +++ b/tests/unit_tests/storage/test_parameter_sample_types.py @@ -362,13 +362,13 @@ def test_surface_param( # Assert that the data has been internalised to storage if expect_num_loaded > 0: - arr = fs.load_parameters("MY_PARAM", 0).values.T + arr = fs.load_parameters("MY_PARAM", 0)["values"].values.T assert arr.flatten().tolist() == [0.0, 1.0, 2.0, 3.0] else: with pytest.raises( KeyError, match="No dataset 'MY_PARAM' in storage for realization 0" ): - fs.load_parameters("MY_PARAM", 0) + fs.load_parameters("MY_PARAM", 0)["values"] @pytest.mark.integration_test @@ -405,7 +405,11 @@ def test_gen_kw_forward_init(tmpdir, storage, load_forward_init): else: _, fs = create_runpath(storage, "config.ert") assert Path("simulations/realization-0/iter-0/kw.txt").exists() - value = fs.load_parameters("KW_NAME", 0).sel(names="MY_KEYWORD").values + value = ( + fs.load_parameters("KW_NAME", 0) + .sel(names="MY_KEYWORD")["values"] + .values + ) assert value == 1.31 @@ -493,7 +497,7 @@ def test_that_first_three_parameters_sampled_snapshot(tmpdir, storage): with open("prior.txt", mode="w", encoding="utf-8") as fh: fh.writelines("MY_KEYWORD NORMAL 0 1") _, fs = create_runpath(storage, "config.ert", [True] * 3) - prior = fs.load_parameters("KW_NAME", range(3)).values.ravel() + prior = fs.load_parameters("KW_NAME", range(3))["values"].values.ravel() expected = np.array([-0.8814228, 1.5847818, 1.009956]) np.testing.assert_almost_equal(prior, expected) @@ -546,9 +550,9 @@ def test_that_sampling_is_fixed_from_name( key_hash = sha256(b"1234" + b"KW_NAME:MY_KEYWORD") seed = np.frombuffer(key_hash.digest(), dtype="uint32") expected = np.random.default_rng(seed).standard_normal(num_realisations) - assert fs.load_parameters("KW_NAME").sel( - names="MY_KEYWORD" - ).values.ravel().tolist() == list(expected) + assert fs.load_parameters("KW_NAME").sel(names="MY_KEYWORD")[ + "values" + ].values.ravel().tolist() == list(expected) @pytest.mark.parametrize( @@ -609,7 +613,7 @@ def test_that_sub_sample_maintains_order(tmpdir, storage, mask, expected): ) assert ( - fs.load_parameters("KW_NAME") + fs.load_parameters("KW_NAME")["values"] .sel(names="MY_KEYWORD") .values.ravel() .tolist() @@ -731,10 +735,12 @@ def test_surface_param_update(tmpdir): prior = storage.get_ensemble_by_name("prior") posterior = storage.get_ensemble_by_name("smoother_update") prior_param = ( - prior.load_parameters("MY_PARAM", range(5)).values.reshape(5, 2 * 3).T + prior.load_parameters("MY_PARAM", range(5))["values"] + .values.reshape(5, 2 * 3) + .T ) posterior_param = ( - posterior.load_parameters("MY_PARAM", range(5)) + posterior.load_parameters("MY_PARAM", range(5))["values"] .values.reshape(5, 2 * 3) .T ) @@ -772,8 +778,8 @@ def test_surface_param_update(tmpdir): assert not (surf.values == surf2.values).any() - assert len(prior.load_parameters("MY_PARAM", 0).x) == 2 - assert len(prior.load_parameters("MY_PARAM", 0).y) == 3 + assert len(prior.load_parameters("MY_PARAM", 0)["values"].x) == 2 + assert len(prior.load_parameters("MY_PARAM", 0)["values"].y) == 3 @pytest.mark.integration_test @@ -908,6 +914,6 @@ def test_gen_kw_optional_template(storage, tmpdir, config_str, expected): fh.writelines("MY_KEYWORD NORMAL 0 1") create_runpath(storage, "config.ert") - assert list(storage.ensembles)[0].load_parameters( - "KW_NAME" - ).values.flatten().tolist() == pytest.approx([expected]) + assert list(storage.ensembles)[0].load_parameters("KW_NAME")[ + "values" + ].values.flatten().tolist() == pytest.approx([expected])