Add support for design matrix in ensemble smoother

Author: JHolba

src/ert/gui/simulation/ensemble_smoother_panel.py

@@ -15,10 +15,14 @@
     TargetEnsembleModel,
     TextModel,
 )
+from ert.gui.tools.design_matrix.design_matrix_panel import DesignMatrixPanel
 from ert.mode_definitions import ENSEMBLE_SMOOTHER_MODE
 from ert.run_models import EnsembleSmoother
-from ert.validation import ProperNameFormatArgument, RangeStringArgument
-from ert.validation.proper_name_argument import ExperimentValidation
+from ert.validation import (
+    ExperimentValidation,
+    ProperNameFormatArgument,
+    RangeStringArgument,
+)
 
 from .experiment_config_panel import ExperimentConfigPanel
 
@@ -44,8 +48,8 @@ def __init__(
     ) -> None:
         super().__init__(EnsembleSmoother)
         self.notifier = notifier
-        layout = QFormLayout()
 
+        layout = QFormLayout()
         self.setObjectName("ensemble_smoother_panel")
 
         self._experiment_name_field = StringBox(
@@ -63,6 +67,7 @@ def __init__(
 
         runpath_label = CopyableLabel(text=run_path)
         layout.addRow("Runpath:", runpath_label)
+
         number_of_realizations_label = QLabel(f"<b>{ensemble_size}</b>")
         layout.addRow(QLabel("Number of realizations:"), number_of_realizations_label)
 
@@ -89,6 +94,15 @@ def __init__(
         self._active_realizations_field.setValidator(RangeStringArgument(ensemble_size))
         layout.addRow("Active realizations", self._active_realizations_field)
 
+        design_matrix = analysis_config.design_matrix
+        if design_matrix is not None:
+            layout.addRow(
+                "Design Matrix",
+                DesignMatrixPanel.get_design_matrix_button(
+                    self._active_realizations_field, design_matrix
+                ),
+            )
+
         self.setLayout(layout)
 
         self._experiment_name_field.getValidationSupport().validationChanged.connect(

src/ert/run_models/ensemble_smoother.py

@@ -8,13 +8,14 @@
 import numpy as np
 
 from ert.config import ErtConfig, ESSettings, HookRuntime, UpdateSettings
-from ert.enkf_main import sample_prior
+from ert.config.parsing.config_errors import ConfigValidationError
+from ert.enkf_main import sample_prior, save_design_matrix_to_ensemble
 from ert.ensemble_evaluator import EvaluatorServerConfig
 from ert.storage import Storage
 from ert.trace import tracer
 
 from ..run_arg import create_run_arguments
-from .base_run_model import StatusEvents, UpdateRunModel
+from .base_run_model import ErtRunError, StatusEvents, UpdateRunModel
 
 if TYPE_CHECKING:
     from ert.config import QueueConfig
@@ -66,6 +67,7 @@ def __init__(
         self.support_restart = False
 
         self._parameter_configuration = config.ensemble_config.parameter_configuration
+        self._design_matrix = config.analysis_config.design_matrix
         self._observations = config.observations
         self._response_configuration = config.ensemble_config.response_configuration
 
@@ -74,14 +76,27 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
+
+        parameters_config = self._parameter_configuration
+        design_matrix = self._design_matrix
+        design_matrix_group = None
+        if design_matrix is not None:
+            try:
+                parameters_config, design_matrix_group = (
+                    design_matrix.merge_with_existing_parameters(parameters_config)
+                )
+            except ConfigValidationError as exc:
+                raise ErtRunError(str(exc)) from exc
+
         self.restart = restart
         self.run_workflows(
             HookRuntime.PRE_EXPERIMENT,
             fixtures={"random_seed": self.random_seed},
         )
         ensemble_format = self.target_ensemble_format
         experiment = self._storage.create_experiment(
-            parameters=self._parameter_configuration,
+            parameters=parameters_config
+            + ([design_matrix_group] if design_matrix_group else []),
             observations=self._observations,
             responses=self._response_configuration,
             name=self.experiment_name,
@@ -106,6 +121,13 @@ def run_experiment(
             random_seed=self.random_seed,
         )
 
+        if design_matrix_group is not None and design_matrix is not None:
+            save_design_matrix_to_ensemble(
+                design_matrix.design_matrix_df,
+                prior,
+                np.where(self.active_realizations)[0],
+                design_matrix_group.name,
+            )
         self._evaluate_and_postprocess(
             prior_args,
             prior,

src/ert/run_models/model_factory.py

@@ -187,9 +187,7 @@ def _setup_ensemble_smoother(
     update_settings: UpdateSettings,
     status_queue: SimpleQueue[StatusEvents],
 ) -> EnsembleSmoother:
-    active_realizations = _realizations(
-        args, config.model_config.num_realizations
-    ).tolist()
+    active_realizations = _get_active_realizations_list(args, config)
     if len(active_realizations) < 2:
         raise ConfigValidationError(
             "Number of active realizations must be at least 2 for an update step"

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

@@ -11,7 +11,11 @@
 
 from ert.cli.main import ErtCliError
 from ert.config import ErtConfig
-from ert.mode_definitions import ENSEMBLE_EXPERIMENT_MODE, ES_MDA_MODE
+from ert.mode_definitions import (
+    ENSEMBLE_EXPERIMENT_MODE,
+    ENSEMBLE_SMOOTHER_MODE,
+    ES_MDA_MODE,
+)
 from ert.storage import open_storage
 from tests.ert.ui_tests.cli.run_cli import run_cli
 
@@ -304,7 +308,14 @@ def _evaluate(coeffs, x):
 
 
 @pytest.mark.usefixtures("copy_poly_case")
-def test_design_matrix_on_esmda():
+@pytest.mark.parametrize(
+    "experiment_mode, ensemble_name, iterations",
+    [
+        (ES_MDA_MODE, "default_", 4),
+        (ENSEMBLE_SMOOTHER_MODE, "iter-", 2),
+    ],
+)
+def test_design_matrix_on_esmda(experiment_mode, ensemble_name, iterations):
     design_path = "design_matrix.xlsx"
     reals = range(10)
     values = [random.uniform(0, 2) for _ in reals]
@@ -375,7 +386,7 @@ def _evaluate(coeffs, x):
         f.write("c UNIFORM 0 5")
 
     run_cli(
-        ES_MDA_MODE,
+        experiment_mode,
         "--disable-monitoring",
         "poly.ert",
         "--experiment-name",
@@ -385,8 +396,8 @@ def _evaluate(coeffs, x):
     coeffs_a_previous = None
     with open_storage(storage_path) as storage:
         experiment = storage.get_experiment_by_name("test-experiment")
-        for i in range(4):
-            ensemble = experiment.get_ensemble_by_name(f"default_{i}")
+        for i in range(iterations):
+            ensemble = experiment.get_ensemble_by_name(f"{ensemble_name}{i}")
 
             # coeffs_a should be different in all realizations
             coeffs_a = ensemble.load_parameters("COEFFS_A")["values"].values.flatten()

tests/ert/unit_tests/cli/test_model_hook_order.py

@@ -134,6 +134,7 @@ def test_hook_call_order_ensemble_smoother(monkeypatch):
     )
     test_class.run_ensemble_evaluator = MagicMock(return_value=[0])
     test_class._storage = storage_mock
+    test_class._design_matrix = None
     test_class.run_experiment(MagicMock())
 
     assert run_wfs_mock.mock_calls == EXPECTED_CALL_ORDER