Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add run experiment with design matrix to ensemble experiment panel #9186

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 54 additions & 7 deletions src/ert/config/design_matrix.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,10 @@
from ert.config.gen_kw_config import GenKwConfig, TransformFunctionDefinition

from ._option_dict import option_dict
from .parsing import (
ConfigValidationError,
ErrorInfo,
)
from .parsing import ConfigValidationError, ErrorInfo

if TYPE_CHECKING:
from ert.config import (
ParameterConfig,
)
from ert.config import ParameterConfig

DESIGN_MATRIX_GROUP = "DESIGN_MATRIX"

Expand Down Expand Up @@ -74,6 +69,58 @@
default_sheet=default_sheet,
)

def merge_with_existing_parameters(
self, existing_parameters: List[ParameterConfig]
) -> tuple[List[ParameterConfig], ParameterConfig | None]:
"""
This method merges the design matrix parameters with the existing parameters and
returns the new list of existing parameters, wherein we drop GEN_KW group having a full overlap with the design matrix group.
GEN_KW group that was dropped will acquire a new name from the design matrix group.
Additionally, the ParameterConfig which is the design matrix group is returned separately.

Args:
existing_parameters (List[ParameterConfig]): List of existing parameters

Raises:
ConfigValidationError: If there is a partial overlap between the design matrix group and any existing GEN_KW group

Returns:
tuple[List[ParameterConfig], ParameterConfig]: List of existing parameters and the dedicated design matrix group
"""
if self.parameter_configuration is None:
self.read_design_matrix()

if self.parameter_configuration is None or not isinstance(
self.parameter_configuration[DESIGN_MATRIX_GROUP], GenKwConfig
):
return existing_parameters, None

new_param_config: List[ParameterConfig] = []

design_parameter_group = self.parameter_configuration[DESIGN_MATRIX_GROUP]
design_keys = []
if isinstance(design_parameter_group, GenKwConfig):
design_keys = design_parameter_group.getKeyWords()

Check failure on line 103 in src/ert/config/design_matrix.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

"GenKwConfig" has no attribute "getKeyWords"

Check failure on line 103 in src/ert/config/design_matrix.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

"GenKwConfig" has no attribute "getKeyWords"

design_group_added = False
for genkw_group in existing_parameters:
if not isinstance(genkw_group, GenKwConfig):
new_param_config += [genkw_group]
continue
existing_keys = genkw_group.getKeyWords()

Check failure on line 110 in src/ert/config/design_matrix.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

"GenKwConfig" has no attribute "getKeyWords"

Check failure on line 110 in src/ert/config/design_matrix.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

"GenKwConfig" has no attribute "getKeyWords"
if set(existing_keys) == set(design_keys):
design_parameter_group.name = genkw_group.name
design_group_added = True
elif set(design_keys) & set(existing_keys):
raise ConfigValidationError(
"Overlapping parameter names found in design matrix!"
)
else:
new_param_config += [genkw_group]
if not design_group_added:
new_param_config += [design_parameter_group]
return new_param_config, design_parameter_group

def read_design_matrix(
self,
) -> None:
Expand Down
34 changes: 27 additions & 7 deletions src/ert/enkf_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,20 +19,17 @@
)

import orjson
import pandas as pd
import xarray as xr
from numpy.random import SeedSequence

from ert.config.ert_config import forward_model_data_to_json
from ert.config.forward_model_step import ForwardModelStep
from ert.config.model_config import ModelConfig
from ert.substitutions import Substitutions

from .config import (
ExtParamConfig,
Field,
GenKwConfig,
ParameterConfig,
SurfaceConfig,
)
from .config import ExtParamConfig, Field, GenKwConfig, ParameterConfig, SurfaceConfig
from .config.design_matrix import DESIGN_MATRIX_GROUP
from .run_arg import RunArg
from .runpaths import Runpaths

Expand Down Expand Up @@ -162,6 +159,29 @@ def _seed_sequence(seed: Optional[int]) -> int:
return int_seed


def save_design_matrix_to_ensemble(
design_matrix_df: pd.DataFrame,
ensemble: Ensemble,
active_realizations: Iterable[int],
design_group_name: str = DESIGN_MATRIX_GROUP,
) -> None:
assert not design_matrix_df.empty
for realization_nr in active_realizations:
row = design_matrix_df.loc[realization_nr][DESIGN_MATRIX_GROUP]
ds = xr.Dataset(
{
"values": ("names", list(row.values)),
"transformed_values": ("names", list(row.values)),
"names": list(row.keys()),
}
)
ensemble.save_parameters(
design_group_name,
realization_nr,
ds,
)


def sample_prior(
ensemble: Ensemble,
active_realizations: Iterable[int],
Expand Down
9 changes: 8 additions & 1 deletion src/ert/gui/simulation/ensemble_experiment_panel.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from ert.gui.tools.design_matrix.design_matrix_panel import DesignMatrixPanel
from ert.mode_definitions import ENSEMBLE_EXPERIMENT_MODE
from ert.run_models import EnsembleExperiment
from ert.validation import RangeStringArgument
from ert.validation import ActiveRange, RangeStringArgument
from ert.validation.proper_name_argument import ExperimentValidation, ProperNameArgument

from .experiment_config_panel import ExperimentConfigPanel
Expand Down Expand Up @@ -85,6 +85,13 @@ def __init__(

design_matrix = analysis_config.design_matrix
if design_matrix is not None:
if design_matrix.design_matrix_df is None:
design_matrix.read_design_matrix()

if design_matrix.active_realizations:
self._active_realizations_field.setText(
ActiveRange(design_matrix.active_realizations).rangestring
)
show_dm_param_button = QPushButton("Show parameters")
show_dm_param_button.setObjectName("show-dm-parameters")
show_dm_param_button.setMinimumWidth(50)
Expand Down
33 changes: 31 additions & 2 deletions src/ert/run_models/ensemble_experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import numpy as np

from ert.enkf_main import sample_prior
from ert.enkf_main import sample_prior, save_design_matrix_to_ensemble
from ert.ensemble_evaluator import EvaluatorServerConfig
from ert.storage import Ensemble, Experiment, Storage
from ert.trace import tracer
Expand Down Expand Up @@ -63,10 +63,26 @@ def run_experiment(
restart: bool = False,
) -> None:
self.log_at_startup()
# If design matrix is present, we try to merge design matrix parameters
# to the experiment parameters and set new active realizations
parameters_config = self.ert_config.ensemble_config.parameter_configuration
design_matrix = self.ert_config.analysis_config.design_matrix
design_matrix_group = None
if design_matrix is not None:
parameters_config, design_matrix_group = (
design_matrix.merge_with_existing_parameters(parameters_config)
)

assert design_matrix.active_realizations is not None
self.active_realizations = design_matrix.active_realizations
if not restart:
self.experiment = self._storage.create_experiment(
name=self.experiment_name,
parameters=self.ert_config.ensemble_config.parameter_configuration,
parameters=(
[*parameters_config, design_matrix_group]
if design_matrix_group is not None
else parameters_config
),
observations=self.ert_config.observations,
responses=self.ert_config.ensemble_config.response_configuration,
)
Expand All @@ -89,12 +105,25 @@ def run_experiment(
np.array(self.active_realizations, dtype=bool),
ensemble=self.ensemble,
)

sample_prior(
self.ensemble,
np.where(self.active_realizations)[0],
random_seed=self.random_seed,
)

if (
design_matrix_group is not None
and design_matrix is not None
and design_matrix.design_matrix_df is not None
):
save_design_matrix_to_ensemble(
design_matrix.design_matrix_df,
self.ensemble,
np.where(self.active_realizations)[0],
design_matrix_group.name,
)

self._evaluate_and_postprocess(
run_args,
self.ensemble,
Expand Down
93 changes: 93 additions & 0 deletions tests/ert/ui_tests/cli/analysis/test_design_matrix.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
import os
import stat
from textwrap import dedent

import numpy as np
import pandas as pd
import pytest

from ert.config import ErtConfig
from ert.mode_definitions import ENSEMBLE_EXPERIMENT_MODE
from ert.storage import open_storage
from tests.ert.ui_tests.cli.run_cli import run_cli


@pytest.mark.usefixtures("copy_poly_case")
def test_run_poly_example_with_design_matrix():
design_matrix = "poly_design.xlsx"
num_realizations = 10
a_values = list(range(num_realizations))
design_matrix_df = pd.DataFrame(
{
"REAL": list(range(num_realizations)),
"a": a_values,
}
)
default_sheet_df = pd.DataFrame([["b", 1], ["c", 2]])
with pd.ExcelWriter(design_matrix) as xl_write:
design_matrix_df.to_excel(xl_write, index=False, sheet_name="DesignSheet01")
default_sheet_df.to_excel(
xl_write, index=False, sheet_name="DefaultSheet", header=False
)

with open("poly.ert", "w", encoding="utf-8") as fout:
fout.write(
dedent(
"""\
QUEUE_OPTION LOCAL MAX_RUNNING 10
RUNPATH poly_out/realization-<IENS>/iter-<ITER>
NUM_REALIZATIONS 10
MIN_REALIZATIONS 1
GEN_DATA POLY_RES RESULT_FILE:poly.out
DESIGN_MATRIX poly_design.xlsx DESIGN_SHEET:DesignSheet01 DEFAULT_SHEET:DefaultSheet
INSTALL_JOB poly_eval POLY_EVAL
FORWARD_MODEL poly_eval
"""
)
)

with open("poly_eval.py", "w", encoding="utf-8") as f:
f.write(
dedent(
"""\
#!/usr/bin/env python
import numpy as np
import sys
import json

def _load_coeffs(filename):
with open(filename, encoding="utf-8") as f:
return json.load(f)["DESIGN_MATRIX"]

def _evaluate(coeffs, x):
return coeffs["a"] * x**2 + coeffs["b"] * x + coeffs["c"]

if __name__ == "__main__":
coeffs = _load_coeffs("parameters.json")
output = [_evaluate(coeffs, x) for x in range(10)]
with open("poly.out", "w", encoding="utf-8") as f:
f.write("\\n".join(map(str, output)))
"""
)
)
os.chmod(
"poly_eval.py",
os.stat("poly_eval.py").st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH,
)

run_cli(
ENSEMBLE_EXPERIMENT_MODE,
"--disable-monitor",
"poly.ert",
"--experiment-name",
"test-experiment",
)
storage_path = ErtConfig.from_file("poly.ert").ens_path
with open_storage(storage_path) as storage:
experiment = storage.get_experiment_by_name("test-experiment")
params = experiment.get_ensemble_by_name("default").load_parameters(
"DESIGN_MATRIX"
)["values"]
np.testing.assert_array_equal(params[:, 0], a_values)
np.testing.assert_array_equal(params[:, 1], 10 * [1])
np.testing.assert_array_equal(params[:, 2], 10 * [2])
20 changes: 16 additions & 4 deletions tests/ert/unit_tests/gui/simulation/test_run_dialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from queue import SimpleQueue
from unittest.mock import MagicMock, Mock, patch

import pandas as pd
import pytest
from pytestqt.qtbot import QtBot
from qtpy import QtWidgets
Expand Down Expand Up @@ -705,15 +706,26 @@ def test_that_stdout_and_stderr_buttons_react_to_file_content(
def test_that_design_matrix_show_parameters_button_is_visible(
design_matrix_entry, qtbot: QtBot, storage
):
xls_filename = "design_matrix.xls"
with open(f"{xls_filename}", "w", encoding="utf-8"):
pass
xls_filename = "design_matrix.xlsx"
design_matrix_df = pd.DataFrame(
{
"REAL": list(range(3)),
"a": [0, 1, 2],
}
)
default_sheet_df = pd.DataFrame([["b", 1], ["c", 2]])
with pd.ExcelWriter(xls_filename) as xl_write:
design_matrix_df.to_excel(xl_write, index=False, sheet_name="DesignSheet01")
default_sheet_df.to_excel(
xl_write, index=False, sheet_name="DefaultSheet", header=False
)

config_file = "minimal_config.ert"
with open(config_file, "w", encoding="utf-8") as f:
f.write("NUM_REALIZATIONS 1")
if design_matrix_entry:
f.write(
f"\nDESIGN_MATRIX {xls_filename} DESIGN_SHEET:DesignSheet01 DEFAULT_SHEET:DefaultValues"
f"\nDESIGN_MATRIX {xls_filename} DESIGN_SHEET:DesignSheet01 DEFAULT_SHEET:DefaultSheet"
)

args_mock = Mock()
Expand Down
Loading
Loading