Reimplement load_all_gen_kwn_data

xjules · xjules · commit e4124c5847e0 · 2025-03-10T13:29:05.000+01:00
diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py
@@ -768,14 +768,33 @@ def load_all_gen_kw_data(
             for config in self.experiment.parameter_configuration.values()
             if isinstance(config, ScalarParameters)
         ]
-        if group:
-            scalar_groups = [config for config in scalar_groups if config.name == group]
         for scalars in scalar_groups:
             with contextlib.suppress(KeyError):
-                da = ScalarParameters.load_parameters(
-                    self, scalars.name, realizations
-                ).filter(pl.col("realization").is_in(realizations))
-                df = da.to_pandas()
+                df = self.load_parameters_scalar(scalars.name, realizations)
+                if df.height != len(realizations):
+                    missing_realizations = set(realizations) - set(
+                        df["realization"].to_list()
+                    )
+                    raise IndexError(
+                        f"Missing realizations in {missing_realizations} in {scalars.name}"
+                    )
+                if group:
+                    df = df.select(
+                        ["realization"]
+                        + [col for col in df.columns if col.startswith(f"{group}:")]
+                    )
+
+                df = df.select(
+                    ["realization"]
+                    + [col for col in df.columns if col.endswith(".transformed")]
+                ).rename(
+                    {
+                        col: col.replace(".transformed", "")
+                        for col in df.columns
+                        if col != "realization"
+                    }
+                )
+                df = df.to_pandas().reset_index().set_index("realization")
                 dataframes.append(df)
         if not dataframes:
             return pd.DataFrame()
@@ -808,7 +827,7 @@ def load_parameters_scalar(
             if key not in df_lazy.columns:
                 raise KeyError(f"No such key {key} in scalar parameters!")
             return df_lazy.select(["realization", key]).collect()
-        else:
+        elif realizations is not None and key is None:
             df_lazy = pl.scan_parquet(self._path / f"{_escape_filename(group)}.parquet")
             return df_lazy.filter(pl.col("realization").is_in(realizations)).collect()
 
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_gen_kw_collector/0/gen_kw_collector.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_gen_kw_collector/0/gen_kw_collector.csv
@@ -1,6 +1,6 @@
-Realization,SNAKE_OIL_PARAM:BPR_138_PERSISTENCE,SNAKE_OIL_PARAM:BPR_555_PERSISTENCE,SNAKE_OIL_PARAM:OP1_DIVERGENCE_SCALE,SNAKE_OIL_PARAM:OP1_OCTAVES,SNAKE_OIL_PARAM:OP1_OFFSET,SNAKE_OIL_PARAM:OP1_PERSISTENCE,SNAKE_OIL_PARAM:OP2_DIVERGENCE_SCALE,SNAKE_OIL_PARAM:OP2_OCTAVES,SNAKE_OIL_PARAM:OP2_OFFSET,SNAKE_OIL_PARAM:OP2_PERSISTENCE
-0,0.455637,0.272046,0.838939,3.033689,0.055062,0.291683,0.626204,7.833912,-0.139835,0.3927
-1,0.576027,0.403038,0.863516,3.478627,-0.091152,0.066219,0.564774,10.105767,0.015113,0.111248
-2,0.670747,0.122072,0.893288,4.760788,-0.014215,0.155447,0.824429,8.868052,-0.166693,0.565635
-3,0.555205,0.284146,0.255134,4.424522,0.087298,0.070156,0.938802,6.315156,0.194929,0.343156
-4,0.519822,0.476484,1.033044,4.571844,0.024175,0.01015,1.345163,7.55014,-0.144591,0.152632
+Realization,SNAKE_OIL_PARAM:BPR_138_PERSISTENCE,SNAKE_OIL_PARAM:BPR_555_PERSISTENCE,SNAKE_OIL_PARAM:OP1_DIVERGENCE_SCALE,SNAKE_OIL_PARAM:OP1_OCTAVES,SNAKE_OIL_PARAM:OP1_OFFSET,SNAKE_OIL_PARAM:OP1_PERSISTENCE,SNAKE_OIL_PARAM:OP2_DIVERGENCE_SCALE,SNAKE_OIL_PARAM:OP2_OCTAVES,SNAKE_OIL_PARAM:OP2_OFFSET,SNAKE_OIL_PARAM:OP2_PERSISTENCE,index
+0,0.455637,0.272046,0.838939,3.033689,0.055062,0.291683,0.626204,7.833912,-0.139835,0.3927,0
+1,0.576027,0.403038,0.863516,3.478627,-0.091152,0.066219,0.564774,10.105767,0.015113,0.111248,1
+2,0.670747,0.122072,0.893288,4.760788,-0.014215,0.155447,0.824429,8.868052,-0.166693,0.565635,2
+3,0.555205,0.284146,0.255134,4.424522,0.087298,0.070156,0.938802,6.315156,0.194929,0.343156,3
+4,0.519822,0.476484,1.033044,4.571844,0.024175,0.01015,1.345163,7.55014,-0.144591,0.152632,4
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_misfit_collector/0/misfit_collector.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_misfit_collector/0/misfit_collector.csv
@@ -1,6 +1,6 @@
 Realization,MISFIT:FOPR,MISFIT:WOPR_OP1_108,MISFIT:WOPR_OP1_144,MISFIT:WOPR_OP1_190,MISFIT:WOPR_OP1_36,MISFIT:WOPR_OP1_72,MISFIT:WOPR_OP1_9,MISFIT:WPR_DIFF_1,MISFIT:TOTAL
-0,1572.4551,4.6631575,1.2280039,24.150873,0.16579537,16.603199,0.5786172,17.52338,1637.3682
-1,564.73254,4.3687825,32.65306,2.25,7.513238,7.502955,4.0,3.9172122,626.9378
-2,760.21344,0.67585987,0.04954808,0.8789783,0.53148407,10.315068,0.56918764,21.326956,794.5605
-3,762.2885,0.0573729,2.0035706,89.3921,1.0729967,0.2363393,1.9635296,4.454344,861.46875
-4,978.6855,0.60999763,11.165132,2.3617368,0.51387596,41.033993,0.04177265,27.461798,1061.8737
+0,5640.708,74.44868,350.77222,5739.49,111.28959,185.8407,218.16934,1258.7874,13579.506
+1,1665.7919,16.0,32.65306,2.25,100.0,100.0,4.0,449.67307,2370.368
+2,4174.4463,14.789415,32.65306,2.25,7.928078,25.793795,4.0,223.25809,4485.119
+3,57495.29,655.511,3331.6685,48630.008,464.07623,1183.6646,909.6749,305.27356,112975.164
+4,1263.8804,7.0019083,34.925747,954.34784,17.13819,21.70592,33.607265,639.1491,2971.7563
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_1.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_1.csv
@@ -1,5 +1,5 @@
 Realization,Date,"BPR:1,3,8","BPR:5,5,5",FGIP,FGIPH,FGOR,FGORH,FGPR,FGPRH,FGPT,FGPTH,FOIP,FOIPH,FOPR,FOPRH,FOPT,FOPTH,FWCT,FWCTH,FWIP,FWIPH,FWPR,FWPRH,FWPT,FWPTH,WGOR:OP1,WGOR:OP2,WGORH:OP1,WGORH:OP2,WGPR:OP1,WGPR:OP2,WGPRH:OP1,WGPRH:OP2,WOPR:OP1,WOPR:OP2,WOPRH:OP1,WOPRH:OP2,WWCT:OP1,WWCT:OP2,WWCTH:OP1,WWCTH:OP2,WWPR:OP1,WWPR:OP2,WWPRH:OP1,WWPRH:OP2
-0,2010-01-10,0.9996,0.9996,2499.4473,2499.9956,1.0,1.0,0.0557,0.0012,0.5528,0.0044,1999.4462,1999.994,0.056,0.0017,0.5538,0.0059,0.1776,0.0002,2249.4492,2249.9998,0.0551,0.0,0.5507,1e-04,1.0,1.0,1.0,1.0,0.0557,0.0,0.0006,0.0006,0.056,0.0,0.0008,0.0008,0.3552,0.0,1e-04,0.0002,0.0551,0.0,0.0,0.0
-1,2010-01-10,0.9996,0.9996,2499.8467,2499.9956,1.0,1.0,0.0157,0.0012,0.1533,0.0044,1999.8458,1999.994,0.016,0.0017,0.1542,0.0059,0.0657,0.0002,2249.8489,2249.9998,0.0151,0.0,0.1512,1e-04,1.0,1.0,1.0,1.0,0.0,0.0157,0.0006,0.0006,0.0,0.016,0.0008,0.0008,0.0,0.1314,1e-04,0.0002,0.0,0.0151,0.0,0.0
-2,2010-01-10,0.9996,0.9996,2500.0,2499.9956,1.0,1.0,0.0,0.0012,0.0,0.0044,2000.0,1999.994,0.0,0.0017,0.0,0.0059,0.0,0.0002,2250.0,2249.9998,0.0,0.0,0.0,1e-04,1.0,1.0,1.0,1.0,0.0,0.0,0.0006,0.0006,0.0,0.0,0.0008,0.0008,0.0,0.0,1e-04,0.0002,0.0,0.0,0.0,0.0
-3,2010-01-10,0.9996,0.9996,2497.1733,2499.9956,0.9994,1.0,0.2835,0.0012,2.8267,0.0044,1997.1716,1999.994,0.284,0.0017,2.8285,0.0059,0.4825,0.0002,2247.1775,2249.9998,0.2823,0.0,2.8224,1e-04,1.0,0.9987,1.0,1.0,0.0879,0.1956,0.0006,0.0006,0.0882,0.1958,0.0008,0.0008,0.4661,0.4989,1e-04,0.0002,0.0873,0.195,0.0,0.0
+0,2010-01-10,0.9996,0.9996,2492.4333,2499.9956,0.9998,1.0,0.7571,0.0012,7.5667,0.0044,1992.4324,1999.994,0.7573,0.0017,7.5675,0.0059,0.2499,0.0002,2242.4355,2249.9998,0.7565,0.0,7.5645,1e-04,0.9997,1.0,1.0,1.0,0.7571,0.0,0.0006,0.0006,0.7573,0.0,0.0008,0.0008,0.4997,0.0,1e-04,0.0002,0.7565,0.0,0.0,0.0
+1,2010-01-10,0.9996,0.9996,2499.0493,2499.9956,1.0,1.0,0.0955,0.0012,0.9507,0.0044,1999.0483,1999.994,0.0957,0.0017,0.9515,0.0059,0.2434,0.0002,2249.0515,2249.9998,0.0949,0.0,0.9485,1e-04,1.0,1.0,1.0,1.0,0.0,0.0955,0.0006,0.0006,0.0,0.0957,0.0008,0.0008,0.0,0.4868,1e-04,0.0002,0.0,0.0949,0.0,0.0
+2,2010-01-10,0.9996,0.9995,2500.0,2499.9956,1.0,1.0,0.0,0.0012,0.0,0.0044,2000.0,1999.994,0.0,0.0017,0.0,0.0059,0.0,0.0002,2250.0,2249.9998,0.0,0.0,0.0,1e-04,1.0,1.0,1.0,1.0,0.0,0.0,0.0006,0.0006,0.0,0.0,0.0008,0.0008,0.0,0.0,1e-04,0.0002,0.0,0.0,0.0,0.0
+3,2010-01-10,0.9996,0.9996,2462.3762,2499.9956,0.9999,1.0,3.7632,0.0012,37.6238,0.0044,1962.3744,1999.994,3.7637,0.0017,37.6256,0.0059,0.4999,0.0002,2212.3806,2249.9998,3.762,0.0,37.6195,1e-04,0.9998,0.9999,1.0,1.0,1.5266,2.2366,0.0006,0.0006,1.5268,2.2369,0.0008,0.0008,0.4999,0.4999,1e-04,0.0002,1.526,2.236,0.0,0.0
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_2.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_2.csv
@@ -1,5 +1,5 @@
 Realization,Date,WWCT:OP1,WWCT:OP2
-0,2010-01-10,0.35516918,0.0
-1,2010-01-10,0.0,0.1313817
+0,2010-01-10,0.49971545,0.0
+1,2010-01-10,0.0,0.48680916
 2,2010-01-10,0.0,0.0
-3,2010-01-10,0.46613,0.49889395
+3,2010-01-10,0.4998589,0.49990368
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_3.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_collector/0/summary_collector_3.csv
@@ -1,5 +1,5 @@
 Realization,Date,WWCT:OP1,WWCT:OP2
-4,2010-01-10,0.19479288,0.0
-4,2010-01-20,0.1951685,0.0
-4,2010-01-30,0.19581304,0.0
-4,2010-02-09,0.19672726,0.0
+4,2010-01-10,0.4993012,0.0
+4,2010-01-20,0.49690545,0.0
+4,2010-01-30,0.49286872,0.0
+4,2010-02-09,0.48729724,0.0
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_data_verify_indices_and_values/0/summary_head.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_data_verify_indices_and_values/0/summary_head.csv
@@ -1,8 +1,8 @@
 ,FOPR,FOPR,FOPR,FOPR,FOPR
 Realization,0,1,2,3,4
 Date,,,,,
-2010-01-10,0.055961154,0.0159828,0.0,0.2839923,0.02509654
-2010-01-20,0.05905991,0.01898478,0.0,0.2900899,0.028275477
-2010-01-30,0.064337626,0.024110613,0.0,0.30051348,0.03370005
-2010-02-09,0.07174286,0.031328063,0.0009306805,0.31520578,0.041331254
-2010-02-19,0.081208415,0.04059234,0.009585856,0.33403203,0.0511149
+2010-01-10,0.757321,0.09572082,0.0,3.763688,0.308652
+2010-01-20,0.7603367,0.098736525,0.0,3.7697194,0.31166774
+2010-01-30,0.7654863,0.10388611,0.0,3.7800186,0.3168173
+2010-02-09,0.77273804,0.11113784,0.0,3.794522,0.32406905
+2010-02-19,0.7820472,0.120446995,0.0,3.8131404,0.3333782
diff --git a/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_data_verify_indices_and_values/0/summary_tail.csv b/tests/ert/unit_tests/snapshots/test_libres_facade/test_summary_data_verify_indices_and_values/0/summary_tail.csv
@@ -1,8 +1,8 @@
 ,FOPR,FOPR,FOPR,FOPR,FOPR
 Realization,0,1,2,3,4
 Date,,,,,
-2015-05-14,0.05978195,0.028228229,0.0,0.2988855,0.03184779
-2015-05-24,0.06022726,0.02732692,0.0,0.2993037,0.032080103
-2015-06-03,0.060678575,0.026679888,0.0,0.29959962,0.032258306
-2015-06-13,0.06101494,0.026284758,0.0,0.2997273,0.03237228
-2015-06-23,0.061153173,0.026137633,0.0,0.29976672,0.032415867
+2015-05-14,0.76747996,0.10587975,0.0,0.0,0.31881094
+2015-05-24,0.76703846,0.10543829,0.0,0.0,0.31836948
+2015-06-03,0.76671875,0.10511856,0.0,0.0,0.31804976
+2015-06-13,0.0,0.104921415,0.0,0.0,0.31785262
+2015-06-23,0.0,0.10484742,0.0,0.0,0.31777862
