From 494347af38132129cb822b2c2a4a5f530ff6e9c3 Mon Sep 17 00:00:00 2001 From: Julius Parulek Date: Fri, 18 Oct 2024 12:28:18 +0200 Subject: [PATCH] Test that save parameters internalize DataFrame parameters in the storage --- tests/ert/unit_tests/test_libres_facade.py | 55 +++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/tests/ert/unit_tests/test_libres_facade.py b/tests/ert/unit_tests/test_libres_facade.py index be24bc0d8ae..55155765221 100644 --- a/tests/ert/unit_tests/test_libres_facade.py +++ b/tests/ert/unit_tests/test_libres_facade.py @@ -2,12 +2,15 @@ from datetime import datetime, timedelta from textwrap import dedent +import numpy as np import pytest +from pandas import ExcelWriter from pandas.core.frame import DataFrame from resdata.summary import Summary from ert.config import ErtConfig -from ert.enkf_main import sample_prior +from ert.config.design_matrix import DESIGN_MATRIX_GROUP, DesignMatrix +from ert.enkf_main import sample_prior, save_design_matrix_to_ensemble from ert.libres_facade import LibresFacade from ert.storage import open_storage @@ -241,3 +244,53 @@ def test_load_gen_kw_not_sorted(storage, tmpdir, snapshot): data = ensemble.load_all_gen_kw_data() snapshot.assert_match(data.round(12).to_csv(), "gen_kw_unsorted") + + +@pytest.mark.parametrize( + "reals, expect_error", + [ + pytest.param( + list(range(10)), + False, + id="correct_active_realizations", + ), + pytest.param([10, 11], True, id="incorrect_active_realizations"), + ], +) +def test_save_parameters_to_storage_from_design_dataframe( + tmp_path, reals, expect_error +): + design_path = tmp_path / "design_matrix.xlsx" + ensemble_size = 10 + a_values = np.random.default_rng().uniform(-5, 5, 10) + b_values = np.random.default_rng().uniform(-5, 5, 10) + c_values = np.random.default_rng().uniform(-5, 5, 10) + design_matrix_df = DataFrame({"a": a_values, "b": b_values, "c": c_values}) + with ExcelWriter(design_path) as xl_write: + design_matrix_df.to_excel(xl_write, index=False, sheet_name="DesignSheet01") + DataFrame().to_excel( + xl_write, index=False, sheet_name="DefaultValues", header=False + ) + design_matrix = DesignMatrix(design_path, "DesignSheet01", "DefaultValues") + design_matrix.read_design_matrix() + with open_storage(tmp_path / "storage", mode="w") as storage: + experiment_id = storage.create_experiment( + parameters=[design_matrix.parameter_configuration[DESIGN_MATRIX_GROUP]] + ) + ensemble = storage.create_ensemble( + experiment_id, name="default", ensemble_size=ensemble_size + ) + if expect_error: + with pytest.raises(KeyError): + save_design_matrix_to_ensemble( + design_matrix.design_matrix_df, ensemble, reals + ) + else: + save_design_matrix_to_ensemble( + design_matrix.design_matrix_df, ensemble, reals + ) + params = ensemble.load_parameters(DESIGN_MATRIX_GROUP)["values"] + all(params.names.values == ["a", "b", "c"]) + np.testing.assert_array_almost_equal(params[:, 0], a_values) + np.testing.assert_array_almost_equal(params[:, 1], b_values) + np.testing.assert_array_almost_equal(params[:, 2], c_values)