Adapt to new parameter setup

Author: Yngve S. Kristiansen

src/semeio/workflows/ahm_analysis/ahmanalysis.py

@@ -21,7 +21,7 @@
     SmootherSnapshot,
     smoother_update,
 )
-from ert.config import ESSettings, Field, GenKwConfig, ObservationSettings
+from ert.config import ESSettings, Field, ObservationSettings
 from ert.storage import Ensemble, Storage, open_storage
 from scipy.stats import ks_2samp
 
@@ -200,39 +200,46 @@ def _replace(s: str) -> str:
         os.makedirs(ahmanalysis_reports_dir, exist_ok=True)
 
         # create dataframe with observations vectors (1 by 1 obs and also all_obs)
-        combinations = make_obs_groups(key_map)
+        obs_group_to_obs_key = make_obs_groups(key_map)
 
         field_parameters = [
             p.name
             for p in prior_experiment.parameter_configuration.values()
             if isinstance(p, Field)
         ]
-        gen_kws = [
+        gen_kw_names = [
             p.name
             for p in prior_experiment.parameter_configuration.values()
-            if isinstance(p, GenKwConfig)
+            if p.type == "gen_kw"
         ]
         if field_parameters:
             logger.warning(
                 f"AHM_ANALYSIS will only evaluate scalar parameters, skipping: {field_parameters}"
             )
 
-        scalar_parameters = sorted(gen_kws)
         # identify the set of actual parameters that was updated for now just go
         # through scalar parameters but in future if easier access to field parameter
         # updates should also include field parameters
-        dkeysf = get_updated_parameters(prior_data, scalar_parameters)
+        # dkeysf = get_updated_parameters(prior_data, gen_kw_groups)
+        updated_genkws = [
+            p.name
+            for p in prior_experiment.parameter_configuration.values()
+            if p.update and p.type == "gen_kw"
+        ]
         # setup dataframe for calculated data
         kolmogorov_smirnov_data, active_obs, misfitval = (
-            pd.DataFrame(sorted(dkeysf), columns=["Parameters"]),
+            pd.DataFrame(sorted(updated_genkws), columns=["Parameters"]),
             pd.DataFrame(),
             pd.DataFrame(index=["misfit"]),
         )
+
+        active_obs_per_obs_group = {}
+        misfits_per_obs_group = {}
         # loop over keys and calculate the KS matrix,
         # conditioning one parameter at the time.
-        updated_combinations = deepcopy(combinations)
-        for group_name, obs_group in combinations.items():
-            print("Processing:", group_name)
+        updated_combinations = deepcopy(obs_group_to_obs_key)
+        for obs_group_name, obs_keys in obs_group_to_obs_key.items():
+            print("Processing:", obs_group_name)
 
             #  Use localization to evaluate change of parameters for each observation
             # The order of the context managers is important, as we want to create a new
@@ -255,46 +262,49 @@ def _replace(s: str) -> str:
                     update_log = _run_ministep(
                         prior_storage=prior_ensemble,
                         target_storage=target_ensemble,
-                        obs_group=obs_group,
-                        data_parameters=field_parameters + scalar_parameters,
+                        obs_group=obs_keys,
+                        data_parameters=field_parameters + gen_kw_names,
                         observation_settings=observation_settings,
                         es_settings=es_settings,
                         random_seed=random_seed,
                     )
                     # Get the active vs total observation info
                     df_update_log = make_update_log_df(update_log)
                 except ErtAnalysisError:
-                    logger.error(f"Analysis failed for: {obs_group}")
-                    del updated_combinations[group_name]
+                    logger.error(f"Analysis failed for: {obs_keys}")
+                    del updated_combinations[obs_group_name]
                     continue
                 # Get the updated scalar parameter distributions
                 target_ensemble.load_all_gen_kw_data().to_csv(
-                    ahmanalysis_reports_dir / f"{group_name}.csv"
+                    ahmanalysis_reports_dir / f"{obs_group_name}.csv"
                 )
 
-                active_obs.at["ratio", group_name] = (
+                active_obs.at["ratio", obs_group_name] = (
                     str(count_active_observations(df_update_log))
                     + " active/"
                     + str(len(df_update_log.index))
                 )
+                active_obs_per_obs_group[obs_group_name] = count_active_observations(
+                    df_update_log
+                )
+
                 # Get misfit values
-                misfitval[group_name] = [
+                misfitval[obs_group_name] = [
                     calc_observationsgroup_misfit(
-                        group_name,
+                        obs_group_name,
                         df_update_log,
                         LibresFacade.load_all_misfit_data(prior_ensemble),
                     )
                 ]
+
+                target_data = target_ensemble.load_all_gen_kw_data()
                 # Calculate Ks matrix for scalar parameters
-                kolmogorov_smirnov_data[group_name] = kolmogorov_smirnov_data[
-                    "Parameters"
-                ].map(
-                    calc_kolmogorov_smirnov(
-                        dkeysf,
-                        prior_data,
-                        target_ensemble.load_all_gen_kw_data(),
-                    )
-                )
+
+                kolmogorov_smirnov_data[obs_group_name] = [
+                    ks_2samp(prior_data[dkey], target_data[dkey])[0]
+                    for dkey in sorted(updated_genkws)
+                ]
+
         kolmogorov_smirnov_data.set_index("Parameters", inplace=True)
 
         # save/export the Ks matrix, active_obs, misfitval and prior data

tests/workflows/ahm_analysis/test_ahm_analysis.py

@@ -241,21 +241,6 @@ def test_make_obs_groups(input_map, expected_keys):
 @pytest.mark.parametrize(
     "prior_data, expected_result",
     [
-        [
-            {
-                "SNAKE_OIL_PARAM:OP1_PERSISTENCE": [0, 1, 2],
-                "SNAKE_OIL_PARAM:OP1_OCTAVES": [0, 1, 2],
-                "SNAKE_OIL_PARAM:OP1_DIVERGENCE_SCALE": [0, 1, 2],
-                "SNAKE_OIL_PARAM:OP1_OFFSET": [0, 0, 0],
-                "SNAKE_OIL_PRES:BPR_138_PERSISTENCE": [0, 1, 2],
-            },
-            [
-                "SNAKE_OIL_PARAM:OP1_PERSISTENCE",
-                "SNAKE_OIL_PARAM:OP1_OCTAVES",
-                "SNAKE_OIL_PARAM:OP1_DIVERGENCE_SCALE",
-                "SNAKE_OIL_PRES:BPR_138_PERSISTENCE",
-            ],
-        ],
         [
             {
                 "SNAKE_OIL_PARAM:OP1_PERSISTENCE": [0, 1, 2],