Remove LocalEnsemble.load_all_gen_kw_data()

This commit removes the function, and replaces it with
`LocalEnsemble.load_scalars()` as it gradually moves from pandas towards
polars.

Author: jonathan-eq

src/ert/plugins/hook_implementations/workflows/csv_export.py

@@ -90,7 +90,10 @@ def run(
                     f"The ensemble '{ensemble.name}' does not have any data!"
                 )
 
-            ensemble_data = ensemble.load_all_gen_kw_data()
+            ensemble_data = ensemble.load_scalars().to_pandas().set_index("realization")
+            ensemble_data.columns.name = None
+            ensemble_data.index.name = "Realization"
+            ensemble_data = ensemble_data.sort_index(axis=1)
 
             if design_matrix_path is not None:
                 design_matrix_data = loadDesignMatrix(design_matrix_path)

src/ert/storage/local_ensemble.py

@@ -783,50 +783,6 @@ def _load_responses_lazy(
 
         return pl.concat(loaded) if loaded else pl.DataFrame().lazy()
 
-    def load_all_gen_kw_data(
-        self,
-        group: str | None = None,
-        realization_index: int | None = None,
-    ) -> pd.DataFrame:
-        """Loads scalar parameters (GEN_KWs) into a pandas DataFrame
-        with columns <PARAMETER_GROUP>:<PARAMETER_NAME> and
-        "Realization" as index.
-
-        Parameters
-        ----------
-        group : str, optional
-            Name of parameter group to load.
-        relization_index : int, optional
-            The realization to load.
-
-        Returns
-        -------
-        data : DataFrame
-            A pandas DataFrame containing the GEN_KW data.
-
-        Notes
-        -----
-        Any provided keys that are not gen_kw will be ignored.
-        """
-        if realization_index is not None:
-            realizations = np.array([realization_index])
-        else:
-            ens_mask = (
-                self.get_realization_mask_with_responses()
-                + self.get_realization_mask_with_parameters()
-            )
-            realizations = np.flatnonzero(ens_mask)
-
-        df = self.load_scalars(group, realizations)
-
-        if df.is_empty():
-            return pd.DataFrame()
-
-        dataframe = df.to_pandas().set_index("realization")
-        dataframe.columns.name = None
-        dataframe.index.name = "Realization"
-        return dataframe.sort_index(axis=1)
-
     @require_write
     def save_parameters(
         self,

tests/ert/ui_tests/cli/analysis/test_es_update.py

@@ -53,9 +53,9 @@ def test_that_posterior_has_lower_variance_than_prior():
     with open_storage("storage") as storage:
         experiment = storage.get_experiment_by_name("es-test")
         prior_ensemble = experiment.get_ensemble_by_name("iter-0")
-        df_default = prior_ensemble.load_all_gen_kw_data()
+        df_default = prior_ensemble.load_scalars()
         posterior_ensemble = experiment.get_ensemble_by_name("iter-1")
-        df_target = posterior_ensemble.load_all_gen_kw_data()
+        df_target = posterior_ensemble.load_scalars()
 
         # The std for the ensemble should decrease
         assert float(
@@ -68,8 +68,8 @@ def test_that_posterior_has_lower_variance_than_prior():
     # generalized variance for the parameters.
     assert (
         0
-        < np.linalg.det(df_target.cov().to_numpy())
-        < np.linalg.det(df_default.cov().to_numpy())
+        < np.linalg.det(df_target.to_pandas().cov().to_numpy())
+        < np.linalg.det(df_default.to_pandas().cov().to_numpy())
     )
 
 

tests/ert/ui_tests/cli/test_cli.py

@@ -540,7 +540,11 @@ def test_that_es_mda_on_poly_case_matches_snapshot(snapshot):
         experiment = storage.get_experiment_by_name("es-mda")
         for iter_nr in range(4):
             ensemble = experiment.get_ensemble_by_name(f"iter-{iter_nr}")
-            data.append(ensemble.load_all_gen_kw_data())
+            ensemble_data = ensemble.load_scalars().to_pandas().set_index("realization")
+            ensemble_data.columns.name = None
+            ensemble_data.index.name = "Realization"
+            ensemble_data = ensemble_data.sort_index(axis=1)
+            data.append(ensemble_data)
     result = pd.concat(
         data,
         keys=[f"iter-{iter_}" for iter_ in range(len(data))],
@@ -576,7 +580,11 @@ def test_that_enif_on_poly_case_matches_snapshot(snapshot):
         experiment = storage.get_experiment_by_name("enif")
         for iter_nr in range(2):
             ensemble = experiment.get_ensemble_by_name(f"iter-{iter_nr}")
-            data.append(ensemble.load_all_gen_kw_data())
+            ensemble_data = ensemble.load_scalars().to_pandas().set_index("realization")
+            ensemble_data.columns.name = None
+            ensemble_data.index.name = "Realization"
+            ensemble_data = ensemble_data.sort_index(axis=1)
+            data.append(ensemble_data)
     result = pd.concat(
         data,
         keys=[f"iter-{i}" for i in range(len(data))],

tests/ert/ui_tests/cli/test_update.py

@@ -232,12 +232,10 @@ def test_update_lowers_generalized_variance_or_deactivates_observations(
         if success:
             with open_storage("storage") as storage:
                 experiment = storage.get_experiment_by_name("experiment")
-                prior = experiment.get_ensemble_by_name("iter-0").load_all_gen_kw_data()
-                posterior = experiment.get_ensemble_by_name(
-                    "iter-1"
-                ).load_all_gen_kw_data()
+                prior = experiment.get_ensemble_by_name("iter-0").load_scalars()
+                posterior = experiment.get_ensemble_by_name("iter-1").load_scalars()
 
             assert (
-                np.linalg.det(posterior.cov().to_numpy())
-                <= np.linalg.det(prior.cov().to_numpy()) + 0.001
+                np.linalg.det(posterior.to_pandas().cov().to_numpy())
+                <= np.linalg.det(prior.to_pandas().cov().to_numpy()) + 0.001
             )

tests/ert/ui_tests/gui/test_csv_export.py

@@ -68,7 +68,7 @@ def verify_exported_content(file_name, gui, ensemble_select):
     for name in ensemble_names:
         experiment = gui.notifier.storage.get_experiment_by_name("es_mda")
         ensemble = experiment.get_ensemble_by_name(name)
-        gen_kw_data = ensemble.load_all_gen_kw_data()
+        gen_kw_data = ensemble.load_scalars().to_pandas()
 
         facade = LibresFacade.from_config_file("poly.ert")
         misfit_data = facade.load_all_misfit_data(ensemble)

tests/ert/ui_tests/gui/test_full_manual_update_workflow.py

@@ -88,9 +88,9 @@ def test_manual_analysis_workflow(ensemble_experiment_has_run, qtbot):
         10,
     )
 
-    df_prior = ensemble_prior.load_all_gen_kw_data()
+    df_prior = ensemble_prior.load_scalars().to_pandas()
     ensemble_posterior = experiment.get_ensemble_by_name("iter-0_1")
-    df_posterior = ensemble_posterior.load_all_gen_kw_data()
+    df_posterior = ensemble_posterior.load_scalars().to_pandas()
 
     # Making sure measured data works with failed realizations
     MeasuredData(experiment.get_ensemble_by_name("iter-0"), ["POLY_OBS"])

tests/ert/unit_tests/storage/test_local_storage.py

@@ -1021,40 +1021,69 @@ def test_load_gen_kw_not_sorted(storage, tmpdir, snapshot):
         )
 
         sample_prior(ensemble, range(ensemble_size), random_seed=1234)
-
-        data = ensemble.load_all_gen_kw_data()
+        data = ensemble.load_scalars().to_pandas().set_index("realization")
+        data.columns.name = None
+        data.index.name = "Realization"
+        data = data.sort_index(axis=1)
         snapshot.assert_match(data.round(12).to_csv(), "gen_kw_unsorted")
 
 
 def test_gen_kw_collector(snake_oil_default_storage, snapshot):
-    data = snake_oil_default_storage.load_all_gen_kw_data()
+    data = snake_oil_default_storage.load_scalars().to_pandas().set_index("realization")
+    data.columns.name = None
+    data.index.name = "Realization"
+    data = data.sort_index(axis=1)
     snapshot.assert_match(data.round(6).to_csv(), "gen_kw_collector.csv")
 
     with pytest.raises(KeyError):
         # realization 60:
         _ = data.loc[60]
 
-    data = snake_oil_default_storage.load_all_gen_kw_data(
-        "SNAKE_OIL_PARAM",
-    )[["SNAKE_OIL_PARAM:OP1_PERSISTENCE", "SNAKE_OIL_PARAM:OP1_OFFSET"]]
+    data = (
+        snake_oil_default_storage.load_scalars(
+            "SNAKE_OIL_PARAM",
+        )
+        .to_pandas()
+        .set_index("realization")
+    )
+    data.columns.name = None
+    data.index.name = "Realization"
+    data = data.sort_index(axis=1)
+    data = data[["SNAKE_OIL_PARAM:OP1_PERSISTENCE", "SNAKE_OIL_PARAM:OP1_OFFSET"]]
     snapshot.assert_match(data.round(6).to_csv(), "gen_kw_collector_2.csv")
 
     with pytest.raises(KeyError):
         _ = data["SNAKE_OIL_PARAM:OP1_DIVERGENCE_SCALE"]
 
     realization_index = 3
-    data = snake_oil_default_storage.load_all_gen_kw_data(
-        "SNAKE_OIL_PARAM",
-        realization_index=realization_index,
-    )["SNAKE_OIL_PARAM:OP1_PERSISTENCE"]
+    data = (
+        snake_oil_default_storage.load_scalars(
+            "SNAKE_OIL_PARAM",
+            realization_index=realization_index,
+        )
+        .to_pandas()
+        .set_index("realization")
+    )
+    data.columns.name = None
+    data.index.name = "Realization"
+    data = data.sort_index(axis=1)
+    data = data["SNAKE_OIL_PARAM:OP1_PERSISTENCE"]
     snapshot.assert_match(data.round(6).to_csv(), "gen_kw_collector_3.csv")
 
     non_existing_realization_index = 150
     with pytest.raises((IndexError, KeyError)):
-        _ = snake_oil_default_storage.load_all_gen_kw_data(
-            "SNAKE_OIL_PARAM",
-            realization_index=non_existing_realization_index,
-        )["SNAKE_OIL_PARAM:OP1_PERSISTENCE"]
+        data = (
+            snake_oil_default_storage.load_scalars(
+                "SNAKE_OIL_PARAM",
+                realization_index=non_existing_realization_index,
+            )
+            .to_pandas()
+            .set_index("realization")
+        )
+        data.columns.name = None
+        data.index.name = "Realization"
+        data = data.sort_index(axis=1)
+        data = data["SNAKE_OIL_PARAM:OP1_PERSISTENCE"]
 
 
 def test_keyword_type_checks(snake_oil_default_storage):
@@ -1083,12 +1112,12 @@ def test_data_fetching_missing_key(snake_oil_case):
         empty_case = experiment.create_ensemble(name="new_case", ensemble_size=25)
 
         data = [
-            empty_case.load_all_gen_kw_data("nokey", None),
+            empty_case.load_scalars("nokey", None),
         ]
 
         for dataframe in data:
-            assert isinstance(dataframe, DataFrame)
-            assert dataframe.empty
+            assert isinstance(dataframe, pl.DataFrame)
+            assert dataframe.is_empty()
 
 
 @dataclass