From 9f2d24ad415b41aa97ae15474da8a0e0461dba3b Mon Sep 17 00:00:00 2001 From: "Yngve S. Kristiansen" Date: Thu, 28 Nov 2024 10:00:28 +0100 Subject: [PATCH] Migrate finalized keys for response configs --- src/ert/storage/local_storage.py | 3 +- src/ert/storage/migration/to9.py | 65 +++++++++++++++++++ .../storage/migration/test_version_2.py | 14 ++-- .../test_that_storage_matches/responses | 2 +- .../storage/test_storage_migration.py | 8 ++- 5 files changed, 85 insertions(+), 7 deletions(-) create mode 100644 src/ert/storage/migration/to9.py diff --git a/src/ert/storage/local_storage.py b/src/ert/storage/local_storage.py index ecaee5f9d6f..59a19424162 100644 --- a/src/ert/storage/local_storage.py +++ b/src/ert/storage/local_storage.py @@ -472,6 +472,7 @@ def _migrate(self, version: int) -> None: to6, to7, to8, + to9, ) try: @@ -516,7 +517,7 @@ def _migrate(self, version: int) -> None: elif version < _LOCAL_STORAGE_VERSION: migrations = list( - enumerate([to2, to3, to4, to5, to6, to7, to8], start=1) + enumerate([to2, to3, to4, to5, to6, to7, to8, to9], start=1) ) for from_version, migration in migrations[version - 1 :]: print(f"* Updating storage to version: {from_version+1}") diff --git a/src/ert/storage/migration/to9.py b/src/ert/storage/migration/to9.py new file mode 100644 index 00000000000..ca7febba4fa --- /dev/null +++ b/src/ert/storage/migration/to9.py @@ -0,0 +1,65 @@ +import json +import os +from pathlib import Path + +import polars + +info = "Migrate finalized response keys into configs" + + +def _migrate_response_configs_wrt_finalized_keys(path: Path) -> None: + for experiment in path.glob("experiments/*"): + ensembles = path.glob("ensembles/*") + + experiment_id = None + with open(experiment / "index.json", encoding="utf-8") as f: + exp_index = json.load(f) + experiment_id = exp_index["id"] + + responses_config = None + with open(experiment / "responses.json", mode="r", encoding="utf-8") as f: + responses_config = json.load(f) + for response_type, config in responses_config.items(): + if not config.get("has_finalized_keys"): + # Read a sample response and write the keys + for ens in ensembles: + with open(ens / "index.json", encoding="utf-8") as f: + ens_file = json.load(f) + if ens_file["experiment_id"] != experiment_id: + continue + + real_dirs = [*ens.glob("realization-*")] + + for real_dir in real_dirs: + if (real_dir / f"{response_type}.parquet").exists(): + df = polars.read_parquet( + real_dir / f"{response_type}.parquet" + ) + response_keys = df["response_key"].unique().to_list() + config["has_finalized_keys"] = True + config["keys"] = sorted(response_keys) + break + + if config["has_finalized_keys"]: + break + + if "has_finalized_keys" not in config: + # At this point in "storage history", + # only gendata and summary response types + # exist, and only summary starts without finalized keys + config["has_finalized_keys"] = ( + config["_ert_kind"] != "SummaryConfig" + ) + + os.remove(experiment / "responses.json") + with open(experiment / "responses.json", mode="w+", encoding="utf-8") as f: + json.dump( + responses_config, + f, + default=str, + indent=2, + ) + + +def migrate(path: Path) -> None: + _migrate_response_configs_wrt_finalized_keys(path) diff --git a/tests/ert/unit_tests/storage/migration/test_version_2.py b/tests/ert/unit_tests/storage/migration/test_version_2.py index 0971f5d0f0d..e025a310777 100644 --- a/tests/ert/unit_tests/storage/migration/test_version_2.py +++ b/tests/ert/unit_tests/storage/migration/test_version_2.py @@ -24,10 +24,16 @@ def test_migrate_responses(setup_case, set_ert_config): response_info = json.loads( (experiment._path / "responses.json").read_text(encoding="utf-8") ) - assert ( - experiment.response_configuration - == ert_config.ensemble_config.response_configs - ) + + response_config_exp = experiment.response_configuration + response_config_ens = ert_config.ensemble_config.response_configs + + response_config_ens["summary"].has_finalized_keys = response_config_exp[ + "summary" + ].has_finalized_keys + response_config_ens["summary"].keys = response_config_exp["summary"].keys + + assert response_config_exp == response_config_ens assert set(response_info) == { "gen_data", diff --git a/tests/ert/unit_tests/storage/snapshots/test_storage_migration/test_that_storage_matches/responses b/tests/ert/unit_tests/storage/snapshots/test_storage_migration/test_that_storage_matches/responses index 598abfea72f..d66a0b0c4d5 100644 --- a/tests/ert/unit_tests/storage/snapshots/test_storage_migration/test_that_storage_matches/responses +++ b/tests/ert/unit_tests/storage/snapshots/test_storage_migration/test_that_storage_matches/responses @@ -1 +1 @@ -{'gen_data': GenDataConfig(name='gen_data', input_files=['gen%d.txt'], keys=['GEN'], has_finalized_keys=True, report_steps_list=[[1]]), 'summary': SummaryConfig(name='summary', input_files=['CASE'], keys=['FOPR', 'RWPR'], has_finalized_keys=False, refcase={})} +{'gen_data': GenDataConfig(name='gen_data', input_files=['gen%d.txt'], keys=['GEN'], has_finalized_keys=True, report_steps_list=[[1]]), 'summary': SummaryConfig(name='summary', input_files=['CASE'], keys=['FOPR'], has_finalized_keys=True, refcase={})} diff --git a/tests/ert/unit_tests/storage/test_storage_migration.py b/tests/ert/unit_tests/storage/test_storage_migration.py index 40a2fbb3229..295ef981531 100644 --- a/tests/ert/unit_tests/storage/test_storage_migration.py +++ b/tests/ert/unit_tests/storage/test_storage_migration.py @@ -125,6 +125,11 @@ def test_that_storage_matches( response_config = experiment.response_configuration response_config["summary"].refcase = {} + assert all( + "has_finalized_keys" in config + for config in experiment.response_info.values() + ) + with open( experiment._path / experiment._responses_file, "w", encoding="utf-8" ) as f: @@ -195,7 +200,8 @@ def test_that_storage_matches( "gen_data", ) - assert not ensemble.experiment._has_finalized_response_keys("summary") + assert ensemble.experiment._has_finalized_response_keys("summary") + assert ensemble.experiment._has_finalized_response_keys("gen_data") ensemble.save_response("summary", ensemble.load_responses("summary", (0,)), 0) assert ensemble.experiment._has_finalized_response_keys("summary") assert ensemble.experiment.response_type_to_response_keys["summary"] == ["FOPR"]