Add function to load_subset of parameters from ScalarParameters

Author: xjules

src/ert/analysis/_es_update.py

@@ -5,12 +5,7 @@
 import time
 from collections.abc import Callable, Iterable, Sequence
 from fnmatch import fnmatch
-from typing import (
-    TYPE_CHECKING,
-    Generic,
-    Self,
-    TypeVar,
-)
+from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
 import iterative_ensemble_smoother as ies
 import numpy as np
@@ -19,7 +14,13 @@
 import scipy
 from iterative_ensemble_smoother.experimental import AdaptiveESMDA
 
-from ert.config import ESSettings, GenKwConfig, ObservationGroups, UpdateSettings
+from ert.config import (
+    ESSettings,
+    GenKwConfig,
+    ObservationGroups,
+    ScalarParameters,
+    UpdateSettings,
+)
 
 from . import misfit_preprocessor
 from .event import (
@@ -31,10 +32,7 @@
     AnalysisTimeEvent,
     DataSection,
 )
-from .snapshots import (
-    ObservationAndResponseSnapshot,
-    SmootherSnapshot,
-)
+from .snapshots import ObservationAndResponseSnapshot, SmootherSnapshot
 
 if TYPE_CHECKING:
     import numpy.typing as npt
@@ -123,13 +121,29 @@ def _save_param_ensemble_array_to_disk(
         )
 
 
+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
+    )
+
+
 def _load_param_ensemble_array(
     ensemble: Ensemble,
     param_group: str,
     iens_active_index: npt.NDArray[np.int_],
 ) -> npt.NDArray[np.float64]:
     config_node = ensemble.experiment.parameter_configuration[param_group]
-    return config_node.load_parameters(ensemble, param_group, iens_active_index)
+    if isinstance(config_node, ScalarParameters):
+        raise ValueError("Config node is scalar!")
+    dataset = config_node.load_parameters(ensemble, param_group, iens_active_index)
+    assert isinstance(dataset, np.ndarray), "dataset is not an numpy array"
+    return dataset
 
 
 def _expand_wildcards(
@@ -542,9 +556,14 @@ def correlation_callback(
         cross_correlations_accumulator.append(cross_correlations_of_batch)
 
     for param_group in parameters:
-        param_ensemble_array = _load_param_ensemble_array(
-            source_ensemble, param_group, iens_active_index
-        )
+        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
+            )
         if module.localization:
             config_node = source_ensemble.experiment.parameter_configuration[
                 param_group

src/ert/config/parameter_config.py

@@ -103,7 +103,7 @@ def save_parameters(
     @abstractmethod
     def load_parameters(
         self, ensemble: Ensemble, group: str, realizations: npt.NDArray[np.int_]
-    ) -> npt.NDArray[np.float64]:
+    ) -> npt.NDArray[np.float64] | xr.Dataset:
         """
         Load the parameter from internal storage for the given ensemble.
         Must return array of shape (number of parameters, number of realizations).

src/ert/config/scalar_parameter.py

@@ -278,6 +278,7 @@ class ScalarParameter:
         | TransEmptySettings
     )
     input_source: DataSource
+    update: bool = True
     # dataset_file: PolarsData | None
 
 
@@ -375,13 +376,21 @@ def sample_or_load(
 
         return params_ds
 
+    def load_parameters_scalar(
+        self,
+        ensemble: Ensemble,
+        scalar_params: list[str],
+        iens_active_index: npt.NDArray[np.int_],
+    ) -> npt.NDArray[np.float64]:
+        ds = self.load_parameters(ensemble, "SCALAR_PARAMETERS", iens_active_index)
+        return ds.sel(type="transformed")[scalar_params].to_array().values.T
+
     @staticmethod
     def load_parameters(
         ensemble: Ensemble, group: str, realizations: npt.NDArray[np.int_]
     ) -> xr.Dataset:
         return xr.concat(
             [
-                # ensemble.load_parameters_scalar(realization)
                 ensemble.load_parameters(group, realizations)
                 for realization in realizations
             ],

src/ert/storage/local_experiment.py

@@ -13,7 +13,13 @@
 import xtgeo
 from pydantic import BaseModel
 
-from ert.config import ExtParamConfig, Field, GenKwConfig, SurfaceConfig
+from ert.config import (
+    ExtParamConfig,
+    Field,
+    GenKwConfig,
+    ScalarParameters,
+    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
@@ -298,9 +304,23 @@ 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]:
-        return [p.name for p in self.parameter_configuration.values() if p.update]
+    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
 
     @cached_property
     def observations(self) -> dict[str, pl.DataFrame]: