Skip to content

Commit

Permalink
Add [values]
Browse files Browse the repository at this point in the history
  • Loading branch information
kvashchuka committed Dec 18, 2023
1 parent 9eae710 commit a000af9
Show file tree
Hide file tree
Showing 8 changed files with 40 additions and 41 deletions.
10 changes: 4 additions & 6 deletions src/ert/analysis/_es_update.py
Original file line number Diff line number Diff line change
Expand Up @@ -289,11 +289,11 @@ def _create_temporary_parameter_storage(
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"] # type: ignore
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"] # type: ignore
t_field += time.perf_counter() - t
else:
raise NotImplementedError(f"{type(config_node)} is not supported")
Expand Down Expand Up @@ -748,7 +748,7 @@ 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(
Expand Down Expand Up @@ -777,9 +777,7 @@ 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)
)
prior_dataset = source_fs.load_parameters(parameter_group, int(realization))["values"]
assert isinstance(prior_dataset, xr.Dataset)
target_fs.save_parameters(
parameter_group,
Expand Down
2 changes: 1 addition & 1 deletion src/ert/config/ext_param_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions tests/unit_tests/analysis/test_adaptive_localization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
18 changes: 9 additions & 9 deletions tests/unit_tests/analysis/test_es_update.py
Original file line number Diff line number Diff line change
Expand Up @@ -550,14 +550,14 @@ 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

Expand Down Expand Up @@ -604,7 +604,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)
Expand Down Expand Up @@ -760,8 +760,8 @@ 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"]
# Because of adaptive localization, not all parameters should be updated.
# This would fail if with global updates.
assert np.isclose(prior_da, posterior_da).sum() > 0
Expand Down Expand Up @@ -1006,9 +1006,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"]
)
4 changes: 2 additions & 2 deletions tests/unit_tests/cli/test_integration_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ 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")
Expand All @@ -427,7 +427,7 @@ 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):
Expand Down
30 changes: 15 additions & 15 deletions tests/unit_tests/storage/test_field_parameter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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()


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand All @@ -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


Expand Down Expand Up @@ -547,12 +547,12 @@ 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(
Expand Down Expand Up @@ -674,8 +674,8 @@ 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)
Expand Down
2 changes: 1 addition & 1 deletion tests/unit_tests/storage/test_local_ensemble.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,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)))


Expand Down
11 changes: 6 additions & 5 deletions tests/unit_tests/storage/test_parameter_sample_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,7 @@ def test_that_order_of_input_in_user_input_is_abritrary_for_gen_kw_init_files(
def test_surface_param(
storage,
tmpdir,
config_str,
expect_forward_init,
expect_num_loaded,
error,
Expand Down Expand Up @@ -367,7 +368,7 @@ def test_surface_param(
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
Expand Down Expand Up @@ -612,8 +613,8 @@ def test_that_sub_sample_maintains_order(tmpdir, storage, mask, expected):
)

assert (
fs.load_parameters("KW_NAME")
.sel(names="MY_KEYWORD")["values"]
fs.load_parameters("KW_NAME")["values"]
.sel(names="MY_KEYWORD")
.values.ravel()
.tolist()
== expected
Expand Down Expand Up @@ -777,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
Expand Down

0 comments on commit a000af9

Please sign in to comment.