diff --git a/tests/integration_tests/analysis/test_es_update.py b/tests/integration_tests/analysis/test_es_update.py index faadc9009bf..7aaf27fefd6 100644 --- a/tests/integration_tests/analysis/test_es_update.py +++ b/tests/integration_tests/analysis/test_es_update.py @@ -4,6 +4,7 @@ from textwrap import dedent import numpy as np +import polars import pytest import xarray as xr from scipy.ndimage import gaussian_filter @@ -38,14 +39,16 @@ def uniform_parameter(): @pytest.fixture -def obs(): - return xr.Dataset( +def obs() -> polars.DataFrame: + return polars.DataFrame( { - "observations": (["report_step", "index"], [[1.0, 1.0, 1.0]]), - "std": (["report_step", "index"], [[0.1, 1.0, 10.0]]), - }, - coords={"index": [0, 1, 2], "report_step": [0]}, - attrs={"response": "RESPONSE"}, + "response_key": "RESPONSE", + "observation_key": "OBSERVATION", + "report_step": polars.Series(np.full(3, 0), dtype=polars.UInt16), + "index": polars.Series([0, 1, 2], dtype=polars.UInt16), + "observations": [1.0, 1.0, 1.0], + "std": [0.1, 1.0, 10.0], + } ) diff --git a/tests/unit_tests/storage/test_local_storage.py b/tests/unit_tests/storage/test_local_storage.py index 656563bf62e..5dd3d807eed 100644 --- a/tests/unit_tests/storage/test_local_storage.py +++ b/tests/unit_tests/storage/test_local_storage.py @@ -10,6 +10,7 @@ import hypothesis.strategies as st import numpy as np +import polars import pytest import xarray as xr from hypothesis import assume @@ -88,25 +89,18 @@ def test_that_saving_empty_responses_fails_nicely(tmp_path): ValueError, match="Dataset for response group 'RESPONSE' must contain a 'values' variable", ): - ensemble.save_response( - "RESPONSE", - xr.Dataset(), - 0, - ) + ensemble.save_response("RESPONSE", polars.DataFrame(), 0) # Test for dataset with 'values' but no actual data - empty_data = xr.Dataset( + empty_data = polars.DataFrame( { - "values": ( - ["report_step", "index"], - np.array([], dtype=float).reshape(0, 0), - ) - }, - coords={ - "index": np.array([], dtype=int), - "report_step": np.array([], dtype=int), - }, + "response_key": [], + "report_step": [], + "index": [], + "values": [], + } ) + with pytest.raises( ValueError, match="Responses RESPONSE are empty. Cannot proceed with saving to storage.", diff --git a/tests/unit_tests/test_load_forward_model.py b/tests/unit_tests/test_load_forward_model.py index 706ec11a0fb..ada18db7e77 100644 --- a/tests/unit_tests/test_load_forward_model.py +++ b/tests/unit_tests/test_load_forward_model.py @@ -4,6 +4,7 @@ from textwrap import dedent import numpy as np +import polars import pytest from resdata.summary import Summary @@ -166,13 +167,9 @@ def test_load_forward_model_gen_data(setup_case): facade = LibresFacade(config) facade.load_from_forward_model(prior_ensemble, [True], 0) - assert list( - prior_ensemble.load_responses("RESPONSE", (0,)) - .sel(report_step=0, drop=True) - .to_dataframe() - .dropna() - .values.flatten() - ) == [1.0, 3.0] + df = prior_ensemble.load_responses("gen_data", (0,)) + filter_cond = polars.col("report_step").eq(0), polars.col("values").is_not_nan() + assert df.filter(filter_cond)["values"].to_list() == [1.0, 3.0] def test_single_valued_gen_data_with_active_info_is_loaded(setup_case): @@ -192,9 +189,8 @@ def test_single_valued_gen_data_with_active_info_is_loaded(setup_case): facade = LibresFacade(config) facade.load_from_forward_model(prior_ensemble, [True], 0) - assert list( - prior_ensemble.load_responses("RESPONSE", (0,)).to_dataframe().values.flatten() - ) == [1.0] + df = prior_ensemble.load_responses("RESPONSE", (0,)) + assert df["values"].to_list() == [1.0] def test_that_all_deactivated_values_are_loaded(setup_case): @@ -214,10 +210,8 @@ def test_that_all_deactivated_values_are_loaded(setup_case): facade = LibresFacade(config) facade.load_from_forward_model(prior_ensemble, [True], 0) - response = ( - prior_ensemble.load_responses("RESPONSE", (0,)).to_dataframe().values.flatten() - ) - assert np.isnan(response[0]) + response = prior_ensemble.load_responses("RESPONSE", (0,)) + assert np.isnan(response[0]["values"].to_list()) assert len(response) == 1 @@ -254,12 +248,9 @@ def test_loading_gen_data_without_restart(storage, run_paths, run_args): facade = LibresFacade.from_config_file("config.ert") facade.load_from_forward_model(prior_ensemble, [True], 0) - assert list( - prior_ensemble.load_responses("RESPONSE", (0,)) - .to_dataframe() - .dropna() - .values.flatten() - ) == [1.0, 3.0] + df = prior_ensemble.load_responses("RESPONSE", (0,)) + df_no_nans = df.filter(polars.col("values").is_not_nan()) + assert df_no_nans["values"].to_list() == [1.0, 3.0] @pytest.mark.usefixtures("copy_snake_oil_case_storage") diff --git a/tests/unit_tests/test_summary_response.py b/tests/unit_tests/test_summary_response.py index cdcf01a7be9..8f57fa0d231 100644 --- a/tests/unit_tests/test_summary_response.py +++ b/tests/unit_tests/test_summary_response.py @@ -51,13 +51,12 @@ def test_load_summary_response_restart_not_zero( facade = LibresFacade.from_config_file("config.ert") facade.load_from_forward_model(ensemble, [True], 0) - df = ensemble.load_responses("summary", (0,)).to_dataframe() - df = df.unstack(level="name") - df.columns = [col[1] for col in df.columns.values] - df.index = df.index.rename( - {"time": "Date", "realization": "Realization"} - ).reorder_levels(["Realization", "Date"]) + df = ensemble.load_responses("summary", (0,)) + df = df.pivot(on="response_key", values="values") + df = df[df.columns[:17]] + df = df.rename({"time": "Date", "realization": "Realization"}) + snapshot.assert_match( - df.dropna().iloc[:, :15].to_csv(), + df.to_pandas().to_csv(index=False), "summary_restart", )