ert 17.1.9__py3-none-any.whl → 18.0.1__py3-none-any.whl

ert/run_models/everest_run_model.py

@@ -17,21 +17,22 @@ from enum import IntEnum, auto
 from functools import cached_property
 from pathlib import Path
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 import numpy as np
 from numpy.typing import NDArray
 from pydantic import PrivateAttr, ValidationError
 from ropt.enums import ExitCode as RoptExitCode
 from ropt.evaluator import EvaluatorContext, EvaluatorResult
-from ropt.plan import BasicOptimizer
 from ropt.results import FunctionResults, Results
 from ropt.transforms import OptModelTransforms
+from ropt.workflow import BasicOptimizer
 from typing_extensions import TypedDict
 
 from ert.config import (
     EverestConstraintsConfig,
     EverestObjectivesConfig,
+    ExtParamConfig,
     GenDataConfig,
     HookRuntime,
     KnownQueueOptionsAdapter,
@@ -57,18 +58,14 @@ from everest.config import (
     InputConstraintConfig,
     ModelConfig,
     OptimizationConfig,
-    OutputConstraintConfig,
 )
-from everest.config.forward_model_config import SummaryResults
+from everest.config.forward_model_config import ForwardModelStepConfig, SummaryResults
 from everest.everest_storage import EverestStorage
 from everest.optimizer.everest2ropt import everest2ropt
 from everest.optimizer.opt_model_transforms import (
     EverestOptModelTransforms,
     get_optimization_domain_transforms,
 )
-from everest.simulator.everest_to_ert import (
-    extract_summary_keys,
-)
 from everest.strings import EVEREST
 
 from ..run_arg import RunArg, create_run_arguments
@@ -76,7 +73,7 @@ from ..storage import ExperimentState, ExperimentStatus
 from ..storage.local_ensemble import EverestRealizationInfo
 from ..substitutions import Substitutions
 from .event import EverestBatchResultEvent, EverestStatusEvent
-from .run_model import RunModel, StatusEvents
+from .run_model import RunModel, RunModelConfig, StatusEvents
 
 if TYPE_CHECKING:
     from ert.storage import Ensemble, Experiment
@@ -225,24 +222,22 @@ def _get_workflows(
     return res_hooks, res_workflows
 
 
-class EverestRunModel(RunModel):
+class EverestRunModelConfig(RunModelConfig):
     optimization_output_dir: str
     simulation_dir: str
 
     parameter_configuration: list[ParameterConfig]
     response_configuration: list[ResponseConfig]
-    controls: list[ControlConfig]
 
     input_constraints: list[InputConstraintConfig]
-
-    output_constraints: list[OutputConstraintConfig]
-
     optimization: OptimizationConfig
     model: ModelConfig
     keep_run_path: bool
     experiment_name: str
     target_ensemble: str
 
+
+class EverestRunModel(RunModel, EverestRunModelConfig):
     _exit_code: EverestExitCode | None = PrivateAttr(default=None)
     _experiment: Experiment | None = PrivateAttr(default=None)
     _eval_server_cfg: EvaluatorServerConfig | None = PrivateAttr(default=None)
@@ -297,15 +292,9 @@ class EverestRunModel(RunModel):
 
         response_configs.append(everest_config.create_ert_objectives_config())
 
-        constraint_names = [c.name for c in everest_config.output_constraints]
-
-        if constraint_names:
-            response_configs.append(
-                EverestConstraintsConfig(
-                    keys=constraint_names,
-                    input_files=constraint_names,
-                )
-            )
+        constraints_config = everest_config.create_ert_output_constraints_config()
+        if constraints_config is not None:
+            response_configs.append(constraints_config)
 
         gen_data_keys = [
             fm.results.file_name
@@ -327,7 +316,7 @@ class EverestRunModel(RunModel):
             eclbase = summary_fm.results.file_name
             response_configs.append(
                 SummaryConfig(
-                    keys=extract_summary_keys(everest_config), input_files=[eclbase]
+                    keys=_extract_summary_keys(everest_config), input_files=[eclbase]
                 )
             )
         else:
@@ -340,7 +329,7 @@ class EverestRunModel(RunModel):
             runpath_format_string=str(
                 Path(everest_config.simulation_dir)
                 / "batch_<ITER>"
-                / "realization_<GEO_ID>"
+                / "realization_<REALIZATION_ID>"
                 / "<SIM_DIR>"
             ),
             eclbase_format_string=eclbase
@@ -382,7 +371,7 @@ class EverestRunModel(RunModel):
         # Only take into account site queue options
         # if and only if they exist and are of same type as user
         # specified queue system
-        applied_site_queue_options = (
+        site_queue_options_to_apply = (
             runtime_plugins.queue_options.model_dump(exclude_unset=True)
             if (
                 runtime_plugins is not None
@@ -393,7 +382,7 @@ class EverestRunModel(RunModel):
             else {}
         )
 
-        queue_options_dict = applied_site_queue_options | queue_options_from_everconfig
+        queue_options_dict = site_queue_options_to_apply | queue_options_from_everconfig
 
         queue_options = KnownQueueOptionsAdapter.validate_python(queue_options_dict)
 
@@ -521,7 +510,6 @@ class EverestRunModel(RunModel):
             objective_names=everest_config.objective_names,
             objective_functions=everest_config.objective_functions,
             input_constraints=everest_config.input_constraints,
-            output_constraints=everest_config.output_constraints,
             optimization=everest_config.optimization,
             model=everest_config.model,
             optimization_output_dir=everest_config.optimization_output_dir,
@@ -546,13 +534,23 @@ class EverestRunModel(RunModel):
             optimization_callback=optimization_callback,
         )
 
+    @property
+    def ext_param_configs(self) -> list[ExtParamConfig]:
+        ext_params = [
+            c for c in self.parameter_configuration if c.type == "everest_parameters"
+        ]
+
+        # There will and must always be one extparam config for an
+        # Everest optimization.
+        return cast(list[ExtParamConfig], ext_params)
+
     @cached_property
     def _transforms(self) -> EverestOptModelTransforms:
         return get_optimization_domain_transforms(
-            self.controls,
+            self.ext_param_configs,
             self.objectives_config,
             self.input_constraints,
-            self.output_constraints,
+            self.output_constraints_config,
             self.model,
             self.optimization.auto_scale,
         )
@@ -646,6 +644,19 @@ class EverestRunModel(RunModel):
                 )
             )
 
+    @property
+    def output_constraints_config(self) -> EverestConstraintsConfig | None:
+        constraints_config = next(
+            (c for c in self.response_configuration if c.type == "everest_constraints"),
+            None,
+        )
+
+        if constraints_config is None:
+            return None
+
+        assert isinstance(constraints_config, EverestConstraintsConfig)
+        return constraints_config
+
     @property
     def objectives_config(self) -> EverestObjectivesConfig:
         obj_config = next(
@@ -682,18 +693,18 @@ class EverestRunModel(RunModel):
         )
 
         formatted_control_names = [
-            name for config in self.controls for name in config.formatted_control_names
+            name for config in self.ext_param_configs for name in config.input_keys
         ]
         self._ever_storage.init(
             formatted_control_names=formatted_control_names,
             objective_functions=self.objectives_config,
-            output_constraints=self.output_constraints,
+            output_constraints=self.output_constraints_config,
             realizations=self.model.realizations,
         )
         optimizer.set_results_callback(self._handle_optimizer_results)
 
         # Run the optimization:
-        optimizer_exit_code = optimizer.run(initial_guesses).exit_code
+        optimizer_exit_code = optimizer.run(initial_guesses)
 
         # Store some final results.
         self._ever_storage.on_optimization_finished()
@@ -751,10 +762,10 @@ class EverestRunModel(RunModel):
 
     def _create_optimizer(self) -> tuple[BasicOptimizer, list[float]]:
         enopt_config, initial_guesses = everest2ropt(
-            self.controls,
+            cast(list[ExtParamConfig], self.parameter_configuration),
             self.objectives_config,
             self.input_constraints,
-            self.output_constraints,
+            self.output_constraints_config,
             self.optimization,
             self.model,
             self.random_seed,
@@ -821,12 +832,7 @@ class EverestRunModel(RunModel):
         for sim_id in range(sim_to_control_vector.shape[0]):
             sim_controls = sim_to_control_vector[sim_id]
             offset = 0
-            for control_config in self.controls:
-                ext_param_config = next(
-                    c
-                    for c in self.parameter_configuration
-                    if c.name == control_config.name
-                )
+            for ext_param_config in self.ext_param_configs:
                 n_param_keys = len(ext_param_config.parameter_keys)
 
                 # Save controls to ensemble
@@ -1004,10 +1010,10 @@ class EverestRunModel(RunModel):
             )
             constraints = (
                 np.zeros(
-                    (num_all_simulations, len(self.output_constraints)),
+                    (num_all_simulations, len(self.output_constraints_config.keys)),
                     dtype=np.float64,
                 )
-                if self.output_constraints
+                if self.output_constraints_config
                 else None
             )
             sim_ids = np.array([-1] * num_all_simulations, dtype=np.int32)
@@ -1070,7 +1076,7 @@ class EverestRunModel(RunModel):
         for sim_id, (model_realization, perturbation) in enumerate(
             zip(sim_to_model_realization, sim_to_perturbation, strict=True)
         ):
-            substitutions[f"<GEO_ID_{sim_id}_{ensemble.iteration}>"] = str(
+            substitutions[f"<REALIZATION_ID_{sim_id}_{ensemble.iteration}>"] = str(
                 int(model_realization)
             )
             if perturbation >= 0:
@@ -1120,7 +1126,11 @@ class EverestRunModel(RunModel):
         objective_names = self.objectives_config.keys
         objectives = np.zeros((ensemble.ensemble_size, len(objective_names)))
 
-        constraint_names = [c.name for c in self.output_constraints]
+        constraint_names = (
+            self.output_constraints_config.keys
+            if self.output_constraints_config is not None
+            else []
+        )
         constraints = np.zeros((ensemble.ensemble_size, len(constraint_names)))
 
         if not any(self.active_realizations):
@@ -1157,3 +1167,104 @@ class EverestRunModel(RunModel):
         return os.path.exists(self.simulation_dir) and any(
             os.listdir(self.simulation_dir)
         )
+
+
+def _extract_summary_keys(ever_config: EverestConfig) -> list[str]:
+    DEFAULT_DATA_SUMMARY_KEYS = ["YEAR", "YEARS", "TCPU", "TCPUDAY", "MONTH", "DAY"]
+
+    DEFAULT_FIELD_SUMMARY_KEYS = [
+        "FOPR",
+        "FOPT",
+        "FOIR",
+        "FOIT",
+        "FWPR",
+        "FWPT",
+        "FWIR",
+        "FWIT",
+        "FGPR",
+        "FGPT",
+        "FGIR",
+        "FGIT",
+        "FVPR",
+        "FVPT",
+        "FVIR",
+        "FVIT",
+        "FWCT",
+        "FGOR",
+        "FOIP",
+        "FOIPL",
+        "FOIPG",
+        "FWIP",
+        "FGIP",
+        "FGIPL",
+        "FGIPG",
+        "FPR",
+        "FAQR",
+        "FAQRG",
+        "FAQT",
+        "FAQTG",
+        "FWGR",
+    ]
+
+    DEFAULT_WELL_SUMMARY_KEYS = [
+        "WOPR",
+        "WOPT",
+        "WOIR",
+        "WOIT",
+        "WWPR",
+        "WWPT",
+        "WWIR",
+        "WWIT",
+        "WGPR",
+        "WGPT",
+        "WGIR",
+        "WGIT",
+        "WVPR",
+        "WVPT",
+        "WVIR",
+        "WVIT",
+        "WWCT",
+        "WGOR",
+        "WWGR",
+        "WBHP",
+        "WTHP",
+        "WPI",
+    ]
+
+    DEFAULT_WELL_TARGET_SUMMARY_KEYS = [
+        well_key + "T"
+        for well_key in DEFAULT_WELL_SUMMARY_KEYS
+        if well_key.endswith("R") and well_key != "WGOR"
+    ]
+
+    summary_fms: list[ForwardModelStepConfig] = [
+        fm
+        for fm in ever_config.forward_model
+        if fm.results is not None and fm.results.type == "summary"
+    ]
+
+    if not summary_fms:
+        return []
+
+    smry_results = summary_fms[0].results
+    assert isinstance(smry_results, SummaryResults)
+
+    requested_keys: list[str] = ["*"] if smry_results.keys == "*" else smry_results.keys
+
+    well_keys = [
+        f"{sum_key}:*"
+        for sum_key in DEFAULT_WELL_SUMMARY_KEYS + DEFAULT_WELL_TARGET_SUMMARY_KEYS
+    ]
+    deprecated_user_specified_keys = (
+        [] if ever_config.export is None else ever_config.export.keywords
+    )
+
+    return list(
+        set(
+            requested_keys
+            + DEFAULT_DATA_SUMMARY_KEYS
+            + DEFAULT_FIELD_SUMMARY_KEYS
+            + well_keys
+            + deprecated_user_specified_keys
+        )
+    )

ert/run_models/initial_ensemble_run_model.py

@@ -1,4 +1,3 @@
-from abc import ABC
 from typing import Annotated, Any, Literal, Self
 
 import numpy as np
@@ -7,26 +6,24 @@ from polars.datatypes import DataTypeClass
 from pydantic import BaseModel, Field, field_validator
 
 from ert.config import (
-    EverestConstraintsConfig,
-    EverestObjectivesConfig,
     ExtParamConfig,
-    GenDataConfig,
     GenKwConfig,
-    SummaryConfig,
+    KnownResponseTypes,
     SurfaceConfig,
 )
 from ert.config import Field as FieldConfig
 from ert.ensemble_evaluator.config import EvaluatorServerConfig
 from ert.run_arg import create_run_arguments
-from ert.run_models.run_model import RunModel
+from ert.run_models.run_model import RunModel, RunModelConfig
 from ert.sample_prior import sample_prior
 from ert.storage.local_ensemble import LocalEnsemble
 
-
 # https://github.com/pola-rs/polars/issues/13152#issuecomment-1864600078
 # PS: Serializing/deserializing schema is scheduled to be added to polars core,
 # ref https://github.com/pola-rs/polars/issues/20426
 # then this workaround can be omitted.
+
+
 def str_to_dtype(dtype_str: str) -> pl.DataType:
     dtype = eval(f"pl.{dtype_str}")
     if isinstance(dtype, DataTypeClass):
@@ -54,7 +51,7 @@ class DictEncodedDataFrame(BaseModel):
         )
 
 
-class InitialEnsembleRunModel(RunModel, ABC):
+class InitialEnsembleRunModelConfig(RunModelConfig):
     experiment_name: str
     design_matrix: DictEncodedDataFrame | None
     parameter_configuration: list[
@@ -65,31 +62,35 @@ class InitialEnsembleRunModel(RunModel, ABC):
     ]
     response_configuration: list[
         Annotated[
-            (
-                GenDataConfig
-                | SummaryConfig
-                | EverestConstraintsConfig
-                | EverestObjectivesConfig
-            ),
+            (KnownResponseTypes),
             Field(discriminator="type"),
         ]
     ]
     ert_templates: list[tuple[str, str]]
-    observations: dict[str, DictEncodedDataFrame] | None
+    observations: dict[str, DictEncodedDataFrame] | None = None
 
     @field_validator("observations", mode="before")
+    @classmethod
     def make_dict_encoded_observations(
-        cls, v: dict[str, pl.DataFrame | DictEncodedDataFrame] | None
+        cls, v: dict[str, pl.DataFrame | DictEncodedDataFrame | dict[str, Any]] | None
     ) -> dict[str, DictEncodedDataFrame] | None:
         if v is None:
             return None
-        return {
-            k: df
-            if isinstance(df, DictEncodedDataFrame)
-            else DictEncodedDataFrame.from_polars(df)
-            for k, df in v.items()
-        }
 
+        encoded = {}
+        for k, df in v.items():
+            match df:
+                case DictEncodedDataFrame():
+                    encoded[k] = df
+                case pl.DataFrame():
+                    encoded[k] = DictEncodedDataFrame.from_polars(df)
+                case dict():
+                    encoded[k] = DictEncodedDataFrame.model_validate(df)
+
+        return encoded
+
+
+class InitialEnsembleRunModel(RunModel, InitialEnsembleRunModelConfig):
     def _sample_and_evaluate_ensemble(
         self,
         evaluator_server_config: EvaluatorServerConfig,

ert/run_models/manual_update.py

@@ -8,7 +8,7 @@ from uuid import UUID
 from pydantic import PrivateAttr
 
 from ert.ensemble_evaluator import EvaluatorServerConfig
-from ert.run_models.update_run_model import UpdateRunModel
+from ert.run_models.update_run_model import UpdateRunModel, UpdateRunModelConfig
 from ert.storage import Ensemble
 
 from ..analysis import smoother_update
@@ -17,10 +17,12 @@ from .run_model import ErtRunError
 logger = logging.getLogger(__name__)
 
 
-class ManualUpdate(UpdateRunModel):
+class ManualUpdateConfig(UpdateRunModelConfig):
     ensemble_id: str
     ert_templates: list[tuple[str, str]]
 
+
+class ManualUpdate(UpdateRunModel, ManualUpdateConfig):
     _prior: Ensemble = PrivateAttr()
 
     def model_post_init(self, ctx: Any) -> None:

ert/run_models/manual_update_enif.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import functools
+import logging
+
+from ert.storage import Ensemble
+
+from ..analysis import enif_update
+from .manual_update import ManualUpdate
+
+logger = logging.getLogger(__name__)
+
+
+class ManualUpdateEnIF(ManualUpdate):
+    @classmethod
+    def name(cls) -> str:
+        return "Manual EnIF update (Experimental)"
+
+    @classmethod
+    def description(cls) -> str:
+        return "Load parameters and responses from existing → EnIF update"
+
+    def update_ensemble_parameters(
+        self, prior: Ensemble, posterior: Ensemble, weight: float
+    ) -> None:
+        enif_update(
+            prior,
+            posterior,
+            parameters=prior.experiment.update_parameters,
+            observations=prior.experiment.observation_keys,
+            random_seed=self.random_seed,
+            progress_callback=functools.partial(
+                self.send_smoother_event,
+                prior.iteration,
+                prior.id,
+            ),
+        )