ert 19.0.1__py3-none-any.whl → 20.0.0b1__py3-none-any.whl

ert/storage/local_ensemble.py

@@ -15,15 +15,18 @@ from typing import TYPE_CHECKING
 from uuid import UUID
 
 import numpy as np
-import pandas as pd
 import polars as pl
 import resfo
 import xarray as xr
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
-from ert.config import ParameterCardinality, ParameterConfig, SummaryConfig
-from ert.config.response_config import InvalidResponseFile
+from ert.config import (
+    InvalidResponseFile,
+    ParameterCardinality,
+    ParameterConfig,
+    SummaryConfig,
+)
 from ert.substitutions import substitute_runpath_name
 
 from .load_status import LoadResult
@@ -570,30 +573,51 @@ class LocalEnsemble(BaseMode):
     ) -> pl.DataFrame:
         if keys is None:
             keys = self.experiment.parameter_keys
-        elif set(keys) - set(self.experiment.parameter_keys):
-            missing = set(keys) - set(self.experiment.parameter_keys)
-            raise KeyError(f"Parameters not registered to the experiment: {missing}")
 
         df_lazy = self._load_parameters_lazy(SCALAR_FILENAME)
-        df_lazy = df_lazy.select(["realization", *keys])
+        names = df_lazy.collect_schema().names()
+        matches = [key for key in keys if any(key in item for item in names)]
+        if len(matches) != len(keys):
+            missing = set(keys) - set(matches)
+            raise KeyError(f"Parameters not registered to the experiment: {missing}")
+
+        parameter_keys = [
+            key
+            for e in keys
+            for key in self.experiment.parameter_configuration[e].parameter_keys
+        ]
+        df_lazy_filtered = df_lazy.select(["realization", *parameter_keys])
+
         if realizations is not None:
             if isinstance(realizations, int):
                 realizations = np.array([realizations])
-            df_lazy = df_lazy.filter(pl.col("realization").is_in(realizations))
-        df = df_lazy.collect(engine="streaming")
+            df_lazy_filtered = df_lazy_filtered.filter(
+                pl.col("realization").is_in(realizations)
+            )
+        df = df_lazy_filtered.collect(engine="streaming")
         if df.is_empty():
             raise IndexError(
                 f"No matching realizations {realizations} found for {keys}"
             )
 
         if transformed:
+            tmp_configuration: dict[str, ParameterConfig] = {}
+            for key in keys:
+                for col in df.columns:
+                    if col == "realization":
+                        continue
+                    if col.startswith(key):
+                        tmp_configuration[col] = (
+                            self.experiment.parameter_configuration[key]
+                        )
+
             df = df.with_columns(
                 [
                     pl.col(col)
                     .map_elements(
-                        self.experiment.parameter_configuration[col].transform_data(),
-                        return_dtype=df[col].dtype,
+                        tmp_configuration[col].transform_data(),
                     )
+                    .cast(df[col].dtype)
                     .alias(col)
                     for col in df.columns
                     if col != "realization"
@@ -618,13 +642,11 @@ class LocalEnsemble(BaseMode):
             for p in self.experiment.parameter_configuration.values()
             if group in {p.name, p.group_name}
         ]
-
         if not cfgs:
             raise KeyError(f"{group} is not registered to the experiment.")
 
         # if group refers to a group name, we expect the same cardinality
         cardinality = next(cfg.cardinality for cfg in cfgs)
-
         if cardinality == ParameterCardinality.multiple_configs_per_ensemble_dataset:
             return self.load_scalar_keys(
                 [cfg.name for cfg in cfgs], realizations, transformed
@@ -741,20 +763,21 @@ class LocalEnsemble(BaseMode):
     @staticmethod
     def sample_parameter(
         parameter: ParameterConfig,
-        real_nr: int,
+        active_realizations: list[int],
         random_seed: int,
+        num_realizations: int,
     ) -> pl.DataFrame:
-        parameter_value = parameter.sample_value(
-            str(random_seed),
-            real_nr,
+        parameter_values = parameter.sample_values(
+            str(random_seed), active_realizations, num_realizations=num_realizations
         )
 
-        parameter_dict = {parameter.name: parameter_value[0]}
-        parameter_dict["realization"] = real_nr
-        return pl.DataFrame(
-            parameter_dict,
-            schema={parameter.name: pl.Float64, "realization": pl.Int64},
+        parameters = pl.DataFrame(
+            {parameter.name: parameter_values},
+            schema={parameter.name: pl.Float64},
         )
+        realizations_series = pl.Series("realization", list(active_realizations))
+
+        return parameters.with_columns(realizations_series)
 
     def load_responses(self, key: str, realizations: tuple[int, ...]) -> pl.DataFrame:
         """Load responses for key and realizations into xarray Dataset.
@@ -1054,10 +1077,12 @@ class LocalEnsemble(BaseMode):
                     for col, observed_values in observed_cols.items():
                         if col != "time":
                             responses = responses.filter(
-                                pl.col(col).is_in(observed_values.implode())
+                                pl.col(col).is_in(
+                                    observed_values.implode(), nulls_equal=True
+                                )
                             )
 
-                    pivoted = responses.collect(engine="streaming").pivot(  # noqa: PD010
+                    pivoted = responses.collect(engine="streaming").pivot(
                         on="realization",
                         index=["response_key", *response_cls.primary_key],
                         values="values",
@@ -1097,6 +1122,7 @@ class LocalEnsemble(BaseMode):
                             pivoted,
                             how="left",
                             on=["response_key", *response_cls.primary_key],
+                            nulls_equal=True,
                         )
 
                     # Do not drop primary keys which
@@ -1190,98 +1216,6 @@ class LocalEnsemble(BaseMode):
             self._path / "index.json", self._index.model_dump_json().encode("utf-8")
         )
 
-    @property
-    def all_parameters_and_gen_data(self) -> pl.DataFrame | None:
-        """
-        Only for Everest wrt objectives/constraints,
-        disregards summary data and primary key values
-        """
-        param_dfs = []
-        for param_group in self.experiment.parameter_configuration:
-            params_pd = self.load_parameters(param_group)["values"].to_pandas()
-
-            assert isinstance(params_pd, pd.DataFrame)
-            params_pd = params_pd.reset_index()
-            param_df = pl.from_pandas(params_pd)
-
-            param_columns = [c for c in param_df.columns if c != "realizations"]
-            param_df = param_df.rename(
-                {
-                    **{
-                        c: param_group + "." + c.replace("\0", ".")
-                        for c in param_columns
-                    },
-                    "realizations": "realization",
-                }
-            )
-            param_df = param_df.cast(
-                {
-                    "realization": pl.UInt16,
-                    **{c: pl.Float64 for c in param_df.columns if c != "realization"},
-                }
-            )
-            param_dfs.append(param_df)
-
-        responses = self.load_responses(
-            "gen_data", tuple(self.get_realization_list_with_responses())
-        )
-
-        if responses is None:
-            return pl.concat(param_dfs)
-
-        params_wide = pl.concat(
-            [
-                (
-                    pdf.sort("realization").drop("realization")
-                    if i > 0
-                    else pdf.sort("realization")
-                )
-                for i, pdf in enumerate(param_dfs)
-            ],
-            how="horizontal",
-        )
-
-        responses_wide = responses["realization", "response_key", "values"].pivot(  # noqa: PD010
-            on="response_key", values="values"
-        )
-
-        # If responses are missing for some realizations, this _left_ join will
-        # put null (polars) which maps to nan when doing .to_numpy() into the
-        # response columns for those realizations
-        params_and_responses = params_wide.join(
-            responses_wide, on="realization", how="left"
-        ).with_columns(pl.lit(self.iteration).alias("batch"))
-
-        assert self.everest_realization_info is not None
-
-        model_realization_mapping = {
-            k: v["model_realization"] for k, v in self.everest_realization_info.items()
-        }
-        perturbation_mapping = {
-            k: v["perturbation"] for k, v in self.everest_realization_info.items()
-        }
-
-        params_and_responses = params_and_responses.with_columns(
-            pl.col("realization")
-            .replace(model_realization_mapping)
-            .alias("model_realization"),
-            pl.col("realization")
-            .cast(pl.Int32)
-            .replace(perturbation_mapping)
-            .alias("perturbation"),
-        )
-
-        column_order = [
-            "batch",
-            "model_realization",
-            "perturbation",
-            "realization",
-            *[c for c in responses_wide.columns if c != "realization"],
-            *[c for c in params_wide.columns if c != "realization"],
-        ]
-
-        return params_and_responses[column_order]
-
 
 async def _read_parameters(
     run_path: str,

ert/storage/local_experiment.py

@@ -24,7 +24,10 @@ from ert.config import (
     SurfaceConfig,
 )
 from ert.config import Field as FieldConfig
-from ert.config.parsing.context_values import ContextBoolEncoder
+from ert.config._create_observation_dataframes import (
+    create_observation_dataframes,
+)
+from ert.config._observations import Observation
 
 from .mode import BaseMode, Mode, require_write
 
@@ -54,6 +57,7 @@ class _Index(BaseModel):
     # from a different experiment. For example, a manual update
     # is a separate experiment from the one that created the prior.
     ensembles: list[UUID]
+    experiment: dict[str, Any] = {}
     status: ExperimentStatus | None = Field(default=None)
 
 
@@ -80,9 +84,6 @@ class LocalExperiment(BaseMode):
     arguments. Provides methods to create and access associated ensembles.
     """
 
-    _parameter_file = Path("parameter.json")
-    _responses_file = Path("responses.json")
-    _metadata_file = Path("metadata.json")
     _templates_file = Path("templates.json")
     _index_file = Path("index.json")
 
@@ -118,63 +119,21 @@ class LocalExperiment(BaseMode):
         storage: LocalStorage,
         uuid: UUID,
         path: Path,
-        *,
-        parameters: list[ParameterConfig] | None = None,
-        responses: list[ResponseConfig] | None = None,
-        observations: dict[str, pl.DataFrame] | None = None,
-        simulation_arguments: dict[Any, Any] | None = None,
+        experiment_config: dict[str, Any],
         name: str | None = None,
-        templates: list[tuple[str, str]] | None = None,
     ) -> LocalExperiment:
-        """
-        Create a new LocalExperiment and store its configuration data.
-
-        Parameters
-        ----------
-        storage : LocalStorage
-            Storage instance for experiment creation.
-        uuid : UUID
-            Unique identifier for the new experiment.
-        path : Path
-            File system path for storing experiment data.
-        parameters : list of ParameterConfig, optional
-            List of parameter configurations.
-        responses : list of ResponseConfig, optional
-            List of response configurations.
-        observations : dict of str to encoded observation datasets, optional
-            Observations dictionary.
-        simulation_arguments : SimulationArguments, optional
-            Simulation arguments for the experiment.
-        name : str, optional
-            Experiment name. Defaults to current date if None.
-        templates : list of tuple[str, str], optional
-            Run templates for the experiment. Defaults to None.
-
-        Returns
-        -------
-        local_experiment : LocalExperiment
-            Instance of the newly created experiment.
-        """
         if name is None:
             name = datetime.today().isoformat()
 
         storage._write_transaction(
             path / cls._index_file,
-            _Index(id=uuid, name=name, ensembles=[])
+            _Index(id=uuid, name=name, ensembles=[], experiment=experiment_config)
             .model_dump_json(indent=2, exclude_none=True)
             .encode("utf-8"),
         )
 
-        parameter_data = {}
-        for parameter in parameters or []:
-            parameter.save_experiment_data(path)
-            parameter_data.update({parameter.name: parameter.model_dump(mode="json")})
-        storage._write_transaction(
-            path / cls._parameter_file,
-            json.dumps(parameter_data, indent=2).encode("utf-8"),
-        )
-
-        if templates:
+        templates = experiment_config.get("ert_templates")
+        if templates is not None:
             templates_path = path / "templates"
             templates_path.mkdir(parents=True, exist_ok=True)
             templates_abs: list[tuple[str, str]] = []
@@ -191,27 +150,29 @@ class LocalExperiment(BaseMode):
                 json.dumps(templates_abs).encode("utf-8"),
             )
 
-        response_data = {}
-        for response in responses or []:
-            response_data.update({response.type: response.model_dump(mode="json")})
-        storage._write_transaction(
-            path / cls._responses_file,
-            json.dumps(response_data, default=str, indent=2).encode("utf-8"),
-        )
-
-        if observations:
+        observation_declarations = experiment_config.get("observations")
+        if observation_declarations:
             output_path = path / "observations"
-            output_path.mkdir()
-            for response_type, dataset in observations.items():
-                storage._to_parquet_transaction(
-                    output_path / f"{response_type}", dataset
-                )
+            output_path.mkdir(parents=True, exist_ok=True)
 
-        simulation_data = simulation_arguments or {}
-        storage._write_transaction(
-            path / cls._metadata_file,
-            json.dumps(simulation_data, cls=ContextBoolEncoder).encode("utf-8"),
-        )
+            responses_list = experiment_config.get("response_configuration", [])
+            rft_config_json = next(
+                (r for r in responses_list if r.get("type") == "rft"), None
+            )
+            rft_config = (
+                _responses_adapter.validate_python(rft_config_json)
+                if rft_config_json is not None
+                else None
+            )
+
+            obs_adapter = TypeAdapter(Observation)  # type: ignore
+            obs_objs: list[Observation] = []
+            for od in observation_declarations:
+                obs_objs.append(obs_adapter.validate_python(od))
+
+            datasets = create_observation_dataframes(obs_objs, rft_config)
+            for response_type, df in datasets.items():
+                storage._to_parquet_transaction(output_path / response_type, df)
 
         return cls(storage, path, Mode.WRITE)
 
@@ -285,16 +246,10 @@ class LocalExperiment(BaseMode):
                 return ens
         raise KeyError(f"Ensemble with name '{name}' not found")
 
-    @property
-    def metadata(self) -> dict[str, Any]:
-        path = self.mount_point / self._metadata_file
-        if not path.exists():
-            raise ValueError(f"{self._metadata_file!s} does not exist")
-        return json.loads(path.read_text(encoding="utf-8"))
-
     @property
     def relative_weights(self) -> str:
-        return self.metadata.get("weights", "")
+        assert self.experiment_config is not None
+        return self.experiment_config.get("weights", "")
 
     @property
     def name(self) -> str:
@@ -324,9 +279,8 @@ class LocalExperiment(BaseMode):
 
     @property
     def parameter_info(self) -> dict[str, Any]:
-        return json.loads(
-            (self.mount_point / self._parameter_file).read_text(encoding="utf-8")
-        )
+        parameters_list = self.experiment_config.get("parameter_configuration", [])
+        return {parameter["name"]: parameter for parameter in parameters_list}
 
     @property
     def templates_configuration(self) -> list[tuple[str, str]]:
@@ -348,9 +302,8 @@ class LocalExperiment(BaseMode):
 
     @property
     def response_info(self) -> dict[str, Any]:
-        return json.loads(
-            (self.mount_point / self._responses_file).read_text(encoding="utf-8")
-        )
+        responses_list = self.experiment_config.get("response_configuration", [])
+        return {response["type"]: response for response in responses_list}
 
     def get_surface(self, name: str) -> IrapSurface:
         """
@@ -420,10 +373,50 @@ class LocalExperiment(BaseMode):
 
     @cached_property
     def observations(self) -> dict[str, pl.DataFrame]:
-        observations = sorted(self.mount_point.glob("observations/*"))
+        obs_dir = self.mount_point / "observations"
+
+        if obs_dir.exists():
+            datasets: dict[str, pl.DataFrame] = {}
+            for p in obs_dir.iterdir():
+                if not p.is_file():
+                    continue
+                try:
+                    df = pl.read_parquet(p)
+                except Exception:
+                    continue
+                datasets[p.stem] = df
+            return datasets
+
+        serialized_observations = self.experiment_config.get("observations", None)
+        if not serialized_observations:
+            return {}
+
+        output_path = self.mount_point / "observations"
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        rft_cfg = None
+        try:
+            responses_list = self.experiment_config.get("response_configuration", [])
+            for r in responses_list:
+                if r.get("type") == "rft":
+                    rft_cfg = _responses_adapter.validate_python(r)
+                    break
+        except Exception:
+            rft_cfg = None
+
+        obs_adapter = TypeAdapter(Observation)  # type: ignore
+        obs_objs: list[Observation] = []
+        for od in serialized_observations:
+            obs_objs.append(obs_adapter.validate_python(od))
+
+        datasets = create_observation_dataframes(obs_objs, rft_cfg)
+        for response_type, df in datasets.items():
+            self._storage._to_parquet_transaction(output_path / response_type, df)
+
         return {
-            observation.name: pl.read_parquet(f"{observation}")
-            for observation in observations
+            p.stem: pl.read_parquet(p)
+            for p in (self.mount_point / "observations").iterdir()
+            if p.is_file()
         }
 
     @cached_property
@@ -489,18 +482,22 @@ class LocalExperiment(BaseMode):
             )
 
         config = responses_configuration[response_type]
+
         config.keys = sorted(response_keys)
         config.has_finalized_keys = True
+
+        response_index = next(
+            i
+            for i, c in enumerate(self.experiment_config["response_configuration"])
+            if c["type"] == response_type
+        )
+        self.experiment_config["response_configuration"][response_index] = (
+            config.model_dump(mode="json")
+        )
+
         self._storage._write_transaction(
-            self._path / self._responses_file,
-            json.dumps(
-                {
-                    c.type: c.model_dump(mode="json")
-                    for c in responses_configuration.values()
-                },
-                default=str,
-                indent=2,
-            ).encode("utf-8"),
+            self._path / self._index_file,
+            self._index.model_dump_json(indent=2).encode("utf-8"),
         )
 
         if self.response_key_to_response_type is not None:
@@ -510,17 +507,5 @@ class LocalExperiment(BaseMode):
             del self.response_type_to_response_keys
 
     @property
-    def all_parameters_and_gen_data(self) -> pl.DataFrame | None:
-        if not self.ensembles:
-            return None
-
-        ensemble_dfs = [
-            e.all_parameters_and_gen_data
-            for e in self.ensembles
-            if e.all_parameters_and_gen_data is not None
-        ]
-
-        if not ensemble_dfs:
-            return None
-
-        return pl.concat(ensemble_dfs)
+    def experiment_config(self) -> dict[str, Any]:
+        return self._index.experiment

ert/storage/local_storage.py

@@ -21,7 +21,7 @@ from filelock import FileLock, Timeout
 from pydantic import BaseModel, Field
 
 import ert.storage
-from ert.config import ErtConfig, ParameterConfig, ResponseConfig
+from ert.config import ErtConfig
 from ert.shared import __version__
 
 from .local_ensemble import LocalEnsemble
@@ -31,7 +31,7 @@ from .realization_storage_state import RealizationStorageState
 
 logger = logging.getLogger(__name__)
 
-_LOCAL_STORAGE_VERSION = 23
+_LOCAL_STORAGE_VERSION = 25
 
 
 class _Migrations(BaseModel):
@@ -322,12 +322,8 @@ class LocalStorage(BaseMode):
     @require_write
     def create_experiment(
         self,
-        parameters: list[ParameterConfig] | None = None,
-        responses: list[ResponseConfig] | None = None,
-        observations: dict[str, pl.DataFrame] | None = None,
-        simulation_arguments: dict[Any, Any] | None = None,
+        experiment_config: dict[str, Any] | None = None,
         name: str | None = None,
-        templates: list[tuple[str, str]] | None = None,
     ) -> LocalExperiment:
         """
         Creates a new experiment in the storage.
@@ -352,6 +348,8 @@ class LocalStorage(BaseMode):
         local_experiment : LocalExperiment
             The newly created experiment.
         """
+        if experiment_config is None:
+            experiment_config = {}
 
         exp_id = uuid4()
         path = self._experiment_path(exp_id)
@@ -361,12 +359,8 @@ class LocalStorage(BaseMode):
             self,
             exp_id,
             path,
-            parameters=parameters,
-            responses=responses,
-            observations=observations,
-            simulation_arguments=simulation_arguments,
+            experiment_config=experiment_config,
             name=name,
-            templates=templates,
         )
 
         self._experiments[exp.id] = exp
@@ -519,6 +513,8 @@ class LocalStorage(BaseMode):
             to21,
             to22,
             to23,
+            to24,
+            to25,
         )
 
         try:
@@ -571,6 +567,8 @@ class LocalStorage(BaseMode):
                     20: to21,
                     21: to22,
                     22: to23,
+                    23: to24,
+                    24: to25,
                 }
                 for from_version in range(version, _LOCAL_STORAGE_VERSION):
                     migrations[from_version].migrate(self.path)

ert/storage/migration/to24.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+import polars as pl
+
+info = "Add default None values to RFT observations and responses"
+
+
+def migrate(path: Path) -> None:
+    for rft_obs in path.glob("experiments/*/observations/rft"):
+        rft_obs_df = pl.read_parquet(rft_obs)
+
+        if "zone" not in rft_obs_df.columns:
+            rft_obs_df = rft_obs_df.with_columns(
+                pl.lit(None, dtype=pl.String).alias("zone")
+            )
+
+        rft_obs_df.write_parquet(rft_obs)
+
+    for rft_response in path.glob("ensembles/*/*/rft.parquet"):
+        rft_response_df = pl.read_parquet(rft_response)
+
+        if "zone" not in rft_response_df.columns:
+            rft_response_df = rft_response_df.with_columns(
+                pl.lit(None, dtype=pl.String).alias("zone")
+            )
+        rft_response_df.write_parquet(rft_response)