Make more tests pass

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],
+        }
     )
 
 

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.",

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")

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",
         )