|
| 1 | +from datetime import datetime |
| 2 | + |
| 3 | +import polars as pl |
| 4 | + |
| 5 | +from ert.config import BreakthroughConfig, SummaryConfig |
| 6 | +from ert.storage.local_storage import open_storage |
| 7 | + |
| 8 | + |
| 9 | +def test_that_derive_from_storage_frames_are_stackable_regardless_of_breakthrough_time( |
| 10 | + tmp_path, |
| 11 | +): |
| 12 | + with open_storage(tmp_path, mode="w") as storage: |
| 13 | + response_key = "WWCT:OP1" |
| 14 | + time = datetime(2000, 3, 2, 13, 0) # ruff: ignore[call-datetime-without-tzinfo] |
| 15 | + |
| 16 | + breakthrough_config = BreakthroughConfig( |
| 17 | + keys=[f"BREAKTHROUGH:{response_key}"], |
| 18 | + summary_keys=[response_key], |
| 19 | + thresholds=[0.2], |
| 20 | + observed_dates=[time], |
| 21 | + ) |
| 22 | + |
| 23 | + summary_config = SummaryConfig( |
| 24 | + keys=[response_key], |
| 25 | + input_files=["not_relevant"], |
| 26 | + ) |
| 27 | + |
| 28 | + experiment = storage.create_experiment( |
| 29 | + experiment_config={ |
| 30 | + "response_configuration": [ |
| 31 | + summary_config.model_dump(mode="json"), |
| 32 | + breakthrough_config.model_dump(mode="json"), |
| 33 | + ], |
| 34 | + } |
| 35 | + ) |
| 36 | + |
| 37 | + ensemble = storage.create_ensemble( |
| 38 | + experiment, ensemble_size=2, iteration=0, name="prior" |
| 39 | + ) |
| 40 | + |
| 41 | + def create_summary_response_dataframe( |
| 42 | + response_key: str, realization: int, value_modifier |
| 43 | + ) -> pl.DataFrame: |
| 44 | + return pl.DataFrame( |
| 45 | + { |
| 46 | + "realization": [realization] * 5, |
| 47 | + "response_key": [response_key] * 5, |
| 48 | + "time": [datetime(2000, month, 1, 1, 0) for month in range(1, 6)], # ruff: ignore[call-datetime-without-tzinfo] |
| 49 | + "values": [n / value_modifier for n in range(5)], |
| 50 | + } |
| 51 | + ) |
| 52 | + |
| 53 | + value_over_threshold = create_summary_response_dataframe(response_key, 0, 10) |
| 54 | + value_under_threshold = create_summary_response_dataframe(response_key, 1, 100) |
| 55 | + |
| 56 | + ensemble.save_response("summary", value_over_threshold, 0) |
| 57 | + ensemble.save_response("summary", value_under_threshold, 1) |
| 58 | + |
| 59 | + breakthrough_response0 = breakthrough_config.derive_from_storage(0, 0, ensemble) |
| 60 | + breakthrough_response1 = breakthrough_config.derive_from_storage(0, 1, ensemble) |
| 61 | + |
| 62 | + ensemble.save_response("breakthrough", breakthrough_response0, 0) |
| 63 | + ensemble.save_response("breakthrough", breakthrough_response1, 1) |
| 64 | + |
| 65 | + responses = ensemble.load_responses("breakthrough", (0, 1)) |
| 66 | + assert len(responses) == 2 |
0 commit comments