Add function to save and copy non updated scalar parameters

Author: xjules

src/ert/analysis/_es_update.py

@@ -109,28 +109,27 @@ def _all_parameters(
 
 
 def _save_param_ensemble_array_to_disk(
-    ensemble: Ensemble,
+    source_ensemble: Ensemble,
+    target_ensemble: Ensemble,
     param_ensemble_array: npt.NDArray[np.float64],
     param_group: str,
     iens_active_index: npt.NDArray[np.int_],
 ) -> None:
-    config_node = ensemble.experiment.parameter_configuration[param_group]
-    for i, realization in enumerate(iens_active_index):
-        config_node.save_parameters(
-            ensemble, param_group, realization, param_ensemble_array[:, i]
-        )
-
-
-def _load_param_ensemble_array_scalar(
-    ensemble: Ensemble,
-    scalar_params: list[str],
-    iens_active_index: npt.NDArray[np.int_],
-) -> npt.NDArray[np.float64]:
-    config_node = ensemble.experiment.parameter_configuration["SCALAR_PARAMETERS"]
-    assert isinstance(config_node, ScalarParameters)
-    return config_node.load_parameters_scalar(
-        ensemble, scalar_params, iens_active_index
-    )
+    config_node = target_ensemble.experiment.parameter_configuration[param_group]
+    if isinstance(config_node, ScalarParameters):
+        for i, realization in enumerate(iens_active_index):
+            config_node.save_updated_parameters_and_copy_remaining(
+                source_ensemble,
+                target_ensemble,
+                param_group,
+                realization,
+                param_ensemble_array[:, i],
+            )
+    else:
+        for i, realization in enumerate(iens_active_index):
+            config_node.save_parameters(
+                target_ensemble, param_group, realization, param_ensemble_array[:, i]
+            )
 
 
 def _load_param_ensemble_array(
@@ -140,7 +139,7 @@ def _load_param_ensemble_array(
 ) -> npt.NDArray[np.float64]:
     config_node = ensemble.experiment.parameter_configuration[param_group]
     if isinstance(config_node, ScalarParameters):
-        raise ValueError("Config node is scalar!")
+        return config_node.load_parameters_to_update(ensemble, iens_active_index)
     dataset = config_node.load_parameters(ensemble, param_group, iens_active_index)
     assert isinstance(dataset, np.ndarray), "dataset is not an numpy array"
     return dataset
@@ -556,14 +555,9 @@ def correlation_callback(
         cross_correlations_accumulator.append(cross_correlations_of_batch)
 
     for param_group in parameters:
-        if isinstance(param_group, list):
-            param_ensemble_array = _load_param_ensemble_array_scalar(
-                source_ensemble, iens_active_index
-            )
-        else:
-            param_ensemble_array = _load_param_ensemble_array(
-                source_ensemble, param_group, iens_active_index
-            )
+        param_ensemble_array = _load_param_ensemble_array(
+            source_ensemble, param_group, iens_active_index
+        )
         if module.localization:
             config_node = source_ensemble.experiment.parameter_configuration[
                 param_group
@@ -629,7 +623,11 @@ def correlation_callback(
         start = time.time()
 
         _save_param_ensemble_array_to_disk(
-            target_ensemble, param_ensemble_array, param_group, iens_active_index
+            source_ensemble,
+            target_ensemble,
+            param_ensemble_array,
+            param_group,
+            iens_active_index,
         )
         logger.info(
             f"Storing data for {param_group} completed in {(time.time() - start) / 60} minutes"

src/ert/config/scalar_parameter.py

@@ -376,14 +376,20 @@ def sample_or_load(
 
         return params_ds
 
-    def load_parameters_scalar(
+    def load_parameters_to_update(
         self,
         ensemble: Ensemble,
-        scalar_params: list[str],
         iens_active_index: npt.NDArray[np.int_],
     ) -> npt.NDArray[np.float64]:
+        params_to_update = [
+            f"{param.group_name}:{param.param_name}"
+            for param in self.scalar_params
+            if param.update
+        ]
+        if not params_to_update:
+            raise ValueError("No parameters to update")
         ds = self.load_parameters(ensemble, "SCALAR_PARAMETERS", iens_active_index)
-        return ds.sel(type="transformed")[scalar_params].to_array().values.T
+        return ds.sel(type="transformed")[params_to_update].to_array().values
 
     @staticmethod
     def load_parameters(
@@ -404,9 +410,31 @@ def save_parameters(
         realization: int,
         data: npt.NDArray[np.float64],
     ) -> None:
-        # This would require to change the API potentially
+        # this function is not used in the current implementation
         pass
 
+    def save_updated_parameters_and_copy_remaining(
+        self,
+        source_ensemble: Ensemble,
+        target_ensemble: Ensemble,
+        group: str,
+        realization: int,
+        data: npt.NDArray[np.float64],
+    ) -> None:
+        params_to_update = [
+            f"{param.group_name}:{param.param_name}"
+            for param in self.scalar_params
+            if param.update
+        ]
+        ds = self.load_parameters(source_ensemble, group, np.array(realization))
+        ds.update(
+            {
+                var: (["type"], np.where(ds["type"] == "raw", data, ds[var]))
+                for i, var in enumerate(params_to_update)
+            }
+        )
+        target_ensemble.save_parameters(group, realization, ds)
+
     def __len__(self) -> int:
         return len(self.scalar_params)
 

src/ert/storage/local_experiment.py

@@ -13,13 +13,7 @@
 import xtgeo
 from pydantic import BaseModel
 
-from ert.config import (
-    ExtParamConfig,
-    Field,
-    GenKwConfig,
-    ScalarParameters,
-    SurfaceConfig,
-)
+from ert.config import ExtParamConfig, Field, GenKwConfig, SurfaceConfig
 from ert.config.parsing.context_values import ContextBoolEncoder
 from ert.config.response_config import ResponseConfig
 from ert.storage.mode import BaseMode, Mode, require_write
@@ -304,23 +298,9 @@ def response_configuration(self) -> dict[str, ResponseConfig]:
 
         return responses
 
-    # @cached_property
-    # def update_parameters(self) -> list[str]:
-    #     return [p.name for p in self.parameter_configuration.values() if p.update]
     @cached_property
-    def update_parameters(self) -> list[str | list[str]]:
-        params_to_update = []
-        for _, config_node in self.parameter_configuration.items():
-            scalar_params = []
-            if isinstance(config_node, ScalarParameters):
-                for param in config_node.scalar_params:
-                    if param.update:
-                        scalar_params.append(f"{param.group_name}:{param.param_name}")
-                if scalar_params:
-                    params_to_update.append(scalar_params)
-            elif config_node.update:
-                params_to_update.append(config_node.name)
-        return params_to_update
+    def update_parameters(self) -> list[str]:
+        return [p.name for p in self.parameter_configuration.values() if p.update]
 
     @cached_property
     def observations(self) -> dict[str, pl.DataFrame]: