From 0259cdfba206fe868e363cc7501716ac83a66a9d Mon Sep 17 00:00:00 2001 From: xjules Date: Wed, 22 Nov 2023 09:49:01 +0100 Subject: [PATCH] Rename RealizationState to RealizationStorageState --- src/ert/analysis/_es_update.py | 28 ++++++++----- src/ert/callbacks.py | 6 +-- src/ert/data/_measured_data.py | 5 ++- src/ert/enkf_main.py | 22 +++------- src/ert/gui/ertwidgets/models/ertmodel.py | 10 ++--- src/ert/job_queue/job_queue_node.py | 4 +- src/ert/libres_facade.py | 10 ++--- src/ert/realization_state.py | 2 +- src/ert/run_models/ensemble_experiment.py | 8 ++-- src/ert/run_models/ensemble_smoother.py | 12 ++---- .../run_models/iterated_ensemble_smoother.py | 12 ++---- .../run_models/multiple_data_assimilation.py | 12 ++---- .../workflows/export_misfit_data.py | 4 +- .../scripts/gen_data_rft_export.py | 11 +++-- src/ert/simulator/simulation_context.py | 10 ++--- src/ert/storage/local_ensemble.py | 42 ++++++++++--------- src/ert/storage/local_storage.py | 12 +++--- src/ert/storage/migration/block_fs.py | 12 +++--- tests/performance_tests/test_memory_usage.py | 9 ++-- tests/unit_tests/analysis/test_es_update.py | 8 ++-- tests/unit_tests/gui/tools/test_case_tool.py | 6 +-- .../simulator/test_simulation_context.py | 6 +-- .../status/test_tracking_integration.py | 12 +++--- tests/unit_tests/test_load_forward_model.py | 6 +-- 24 files changed, 128 insertions(+), 141 deletions(-) diff --git a/src/ert/analysis/_es_update.py b/src/ert/analysis/_es_update.py index 496e3812329..3904f53756b 100644 --- a/src/ert/analysis/_es_update.py +++ b/src/ert/analysis/_es_update.py @@ -27,7 +27,7 @@ ) from ert.config import Field, GenKwConfig, SurfaceConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ..config.analysis_module import ESSettings, IESSettings from . import misfit_preprocessor @@ -510,10 +510,13 @@ def analysis_ES( AnalysisStatusEvent(msg="Loading observations and responses..") ) try: - S, ( - observation_values, - observation_errors, - update_snapshot, + ( + S, + ( + observation_values, + observation_errors, + update_snapshot, + ), ) = _load_observations_and_responses( source_fs, alpha, @@ -687,10 +690,13 @@ def analysis_IES( AnalysisStatusEvent(msg="Loading observations and responses..") ) try: - S, ( - observation_values, - observation_errors, - update_snapshot, + ( + S, + ( + observation_values, + observation_errors, + update_snapshot, + ), ) = _load_observations_and_responses( source_fs, alpha, @@ -829,7 +835,7 @@ def smoother_update( analysis_config = UpdateSettings() if analysis_config is None else analysis_config es_settings = ESSettings() if es_settings is None else es_settings ens_mask = prior_storage.get_realization_mask_from_state( - [RealizationState.HAS_DATA] + [RealizationStorageState.HAS_DATA] ) _assert_has_enough_realizations(ens_mask, analysis_config) @@ -887,7 +893,7 @@ def iterative_smoother_update( alpha = analysis_config.alpha std_cutoff = analysis_config.std_cutoff ens_mask = prior_storage.get_realization_mask_from_state( - [RealizationState.HAS_DATA] + [RealizationStorageState.HAS_DATA] ) _assert_has_enough_realizations(ens_mask, analysis_config) diff --git a/src/ert/callbacks.py b/src/ert/callbacks.py index 7ac28d020a4..2af59d32e30 100644 --- a/src/ert/callbacks.py +++ b/src/ert/callbacks.py @@ -9,7 +9,7 @@ from ert.run_arg import RunArg from .load_status import LoadResult, LoadStatus -from .realization_state import RealizationState +from .realization_state import RealizationStorageState logger = logging.getLogger(__name__) @@ -105,9 +105,9 @@ def forward_model_ok( final_result = response_result run_arg.ensemble_storage.state_map[run_arg.iens] = ( - RealizationState.HAS_DATA + RealizationStorageState.HAS_DATA if final_result.status == LoadStatus.LOAD_SUCCESSFUL - else RealizationState.LOAD_FAILURE + else RealizationStorageState.LOAD_FAILURE ) return final_result diff --git a/src/ert/data/_measured_data.py b/src/ert/data/_measured_data.py index 1f3c67435a3..6fcb258eb0a 100644 --- a/src/ert/data/_measured_data.py +++ b/src/ert/data/_measured_data.py @@ -13,7 +13,7 @@ import numpy as np import pandas as pd -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState if TYPE_CHECKING: import numpy.typing as npt @@ -107,7 +107,8 @@ def _get_data( group = obs.attrs["response"] try: response = ensemble.load_responses( - group, tuple(ensemble.realization_list(RealizationState.HAS_DATA)) + group, + tuple(ensemble.realization_list(RealizationStorageState.HAS_DATA)), ) _msg = f"No response loaded for observation key: {key}" if not response: diff --git a/src/ert/enkf_main.py b/src/ert/enkf_main.py index 7ff1146b590..d67c48363ef 100644 --- a/src/ert/enkf_main.py +++ b/src/ert/enkf_main.py @@ -7,25 +7,15 @@ from copy import copy from datetime import datetime from pathlib import Path -from typing import ( - TYPE_CHECKING, - Dict, - Iterable, - List, - Mapping, - Optional, - Union, -) +from typing import TYPE_CHECKING, Dict, Iterable, List, Mapping, Optional, Union import numpy as np from numpy.random import SeedSequence from .analysis.configuration import UpdateConfiguration, UpdateStep -from .config import ( - ParameterConfig, -) +from .config import ParameterConfig from .job_queue import WorkflowRunner -from .realization_state import RealizationState +from .realization_state import RealizationStorageState from .run_context import RunContext from .runpaths import Runpaths from .substitution_list import SubstitutionList @@ -215,10 +205,10 @@ def sample_prior( ensemble.update_realization_state( realization_nr, [ - RealizationState.UNDEFINED, - RealizationState.LOAD_FAILURE, + RealizationStorageState.UNDEFINED, + RealizationStorageState.LOAD_FAILURE, ], - RealizationState.INITIALIZED, + RealizationStorageState.INITIALIZED, ) ensemble.sync() diff --git a/src/ert/gui/ertwidgets/models/ertmodel.py b/src/ert/gui/ertwidgets/models/ertmodel.py index c79bc70e7bf..24c3e38d041 100644 --- a/src/ert/gui/ertwidgets/models/ertmodel.py +++ b/src/ert/gui/ertwidgets/models/ertmodel.py @@ -1,4 +1,4 @@ -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.storage import StorageReader @@ -17,9 +17,9 @@ def get_runnable_realizations_mask(storage: StorageReader, casename: str): return [] runnable_states = [ - RealizationState.UNDEFINED, - RealizationState.INITIALIZED, - RealizationState.LOAD_FAILURE, - RealizationState.HAS_DATA, + RealizationStorageState.UNDEFINED, + RealizationStorageState.INITIALIZED, + RealizationStorageState.LOAD_FAILURE, + RealizationStorageState.HAS_DATA, ] return ensemble.get_realization_mask_from_state(runnable_states) diff --git a/src/ert/job_queue/job_queue_node.py b/src/ert/job_queue/job_queue_node.py index c73323dd1ac..70f5bf001cc 100644 --- a/src/ert/job_queue/job_queue_node.py +++ b/src/ert/job_queue/job_queue_node.py @@ -13,7 +13,7 @@ from ert._clib.queue import _get_submit_attempt, _kill, _refresh_status, _submit from ert.callbacks import forward_model_ok from ert.load_status import LoadStatus -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from . import ResPrototype from .job_status import JobStatus @@ -165,7 +165,7 @@ def run_timeout_callback(self) -> None: def run_exit_callback(self) -> None: self.run_arg.ensemble_storage.state_map[ self.run_arg.iens - ] = RealizationState.LOAD_FAILURE + ] = RealizationStorageState.LOAD_FAILURE def is_running(self, given_status: Optional[JobStatus] = None) -> bool: status = given_status or self.queue_status diff --git a/src/ert/libres_facade.py b/src/ert/libres_facade.py index a6236dbdeb5..d1587c46756 100644 --- a/src/ert/libres_facade.py +++ b/src/ert/libres_facade.py @@ -23,7 +23,7 @@ ) from ert.data import MeasuredData from ert.data._measured_data import ObservationError, ResponseError -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.shared.version import __version__ from ert.storage import EnsembleReader @@ -186,7 +186,7 @@ def get_ensemble_size(self) -> int: return self.config.model_config.num_realizations def get_active_realizations(self, ensemble: EnsembleReader) -> List[int]: - return ensemble.realization_list(RealizationState.HAS_DATA) + return ensemble.realization_list(RealizationStorageState.HAS_DATA) def get_queue_config(self) -> "QueueConfig": return self.config.queue_config @@ -264,7 +264,7 @@ def load_gen_data( report_step: int, realization_index: Optional[int] = None, ) -> DataFrame: - realizations = ensemble.realization_list(RealizationState.HAS_DATA) + realizations = ensemble.realization_list(RealizationStorageState.HAS_DATA) if realization_index is not None: if realization_index not in realizations: raise IndexError(f"No such realization {realization_index}") @@ -338,8 +338,8 @@ def load_all_gen_kw_data( """ ens_mask = ensemble.get_realization_mask_from_state( [ - RealizationState.INITIALIZED, - RealizationState.HAS_DATA, + RealizationStorageState.INITIALIZED, + RealizationStorageState.HAS_DATA, ] ) realizations = ( diff --git a/src/ert/realization_state.py b/src/ert/realization_state.py index 9e3238adafb..068af91fe84 100644 --- a/src/ert/realization_state.py +++ b/src/ert/realization_state.py @@ -1,7 +1,7 @@ from enum import Enum -class RealizationState(Enum): +class RealizationStorageState(Enum): UNDEFINED = 1 INITIALIZED = 2 HAS_DATA = 4 diff --git a/src/ert/run_models/ensemble_experiment.py b/src/ert/run_models/ensemble_experiment.py index 6e54f1a859c..03a6f951c50 100644 --- a/src/ert/run_models/ensemble_experiment.py +++ b/src/ert/run_models/ensemble_experiment.py @@ -8,7 +8,7 @@ from ert.enkf_main import sample_prior from ert.ensemble_evaluator import EvaluatorServerConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.run_context import RunContext from ert.storage import EnsembleAccessor, StorageAccessor @@ -84,10 +84,10 @@ def runSimulations__( state_map = prior_context.sim_fs.state_map for realization_nr in prior_context.active_realizations: if state_map[realization_nr] in [ - RealizationState.UNDEFINED, - RealizationState.LOAD_FAILURE, + RealizationStorageState.UNDEFINED, + RealizationStorageState.LOAD_FAILURE, ]: - state_map[realization_nr] = RealizationState.INITIALIZED + state_map[realization_nr] = RealizationStorageState.INITIALIZED iteration = prior_context.iteration phase_count = iteration + 1 self.setPhaseCount(phase_count) diff --git a/src/ert/run_models/ensemble_smoother.py b/src/ert/run_models/ensemble_smoother.py index 548ddef4772..a5b737c7d74 100644 --- a/src/ert/run_models/ensemble_smoother.py +++ b/src/ert/run_models/ensemble_smoother.py @@ -11,7 +11,7 @@ from ert.config import ErtConfig, HookRuntime from ert.enkf_main import sample_prior from ert.ensemble_evaluator import EvaluatorServerConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.run_context import RunContext from ert.run_models.run_arguments import ESRunArguments from ert.storage import StorageAccessor @@ -19,11 +19,7 @@ from ..analysis._es_update import UpdateSettings from ..config.analysis_module import ESSettings from .base_run_model import BaseRunModel, ErtRunError -from .event import ( - RunModelStatusEvent, - RunModelUpdateBeginEvent, - RunModelUpdateEndEvent, -) +from .event import RunModelStatusEvent, RunModelUpdateBeginEvent, RunModelUpdateEndEvent if TYPE_CHECKING: from ert.config import QueueConfig @@ -108,8 +104,8 @@ def run_experiment( ) states = [ - RealizationState.HAS_DATA, - RealizationState.INITIALIZED, + RealizationStorageState.HAS_DATA, + RealizationStorageState.INITIALIZED, ] self.send_event( diff --git a/src/ert/run_models/iterated_ensemble_smoother.py b/src/ert/run_models/iterated_ensemble_smoother.py index 573e6175a91..f63fb4711e8 100644 --- a/src/ert/run_models/iterated_ensemble_smoother.py +++ b/src/ert/run_models/iterated_ensemble_smoother.py @@ -12,7 +12,7 @@ from ert.config import ErtConfig, HookRuntime from ert.enkf_main import sample_prior from ert.ensemble_evaluator import EvaluatorServerConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.run_context import RunContext from ert.run_models.run_arguments import SIESRunArguments from ert.storage import EnsembleAccessor, StorageAccessor @@ -20,11 +20,7 @@ from ..analysis._es_update import UpdateSettings from ..config.analysis_module import IESSettings from .base_run_model import BaseRunModel, ErtRunError -from .event import ( - RunModelStatusEvent, - RunModelUpdateBeginEvent, - RunModelUpdateEndEvent, -) +from .event import RunModelStatusEvent, RunModelUpdateBeginEvent, RunModelUpdateEndEvent if TYPE_CHECKING: from ert.config import QueueConfig @@ -144,8 +140,8 @@ def run_experiment( ) for current_iter in range(1, iteration_count + 1): states = [ - RealizationState.HAS_DATA, - RealizationState.INITIALIZED, + RealizationStorageState.HAS_DATA, + RealizationStorageState.INITIALIZED, ] self.send_event(RunModelUpdateBeginEvent(iteration=current_iter - 1)) self.send_event( diff --git a/src/ert/run_models/multiple_data_assimilation.py b/src/ert/run_models/multiple_data_assimilation.py index 7f142894cd2..0782f4ffd8e 100644 --- a/src/ert/run_models/multiple_data_assimilation.py +++ b/src/ert/run_models/multiple_data_assimilation.py @@ -11,7 +11,7 @@ from ert.config import ErtConfig, HookRuntime from ert.enkf_main import sample_prior from ert.ensemble_evaluator import EvaluatorServerConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.run_context import RunContext from ert.run_models.run_arguments import ESMDARunArguments from ert.storage import EnsembleAccessor, StorageAccessor @@ -19,11 +19,7 @@ from ..analysis._es_update import UpdateSettings from ..config.analysis_module import ESSettings from .base_run_model import BaseRunModel, ErtRunError -from .event import ( - RunModelStatusEvent, - RunModelUpdateBeginEvent, - RunModelUpdateEndEvent, -) +from .event import RunModelStatusEvent, RunModelUpdateBeginEvent, RunModelUpdateEndEvent if TYPE_CHECKING: from ert.config import QueueConfig @@ -146,8 +142,8 @@ def run_experiment( ) self.ert.runWorkflows(HookRuntime.PRE_UPDATE, self._storage, prior) states = [ - RealizationState.HAS_DATA, - RealizationState.INITIALIZED, + RealizationStorageState.HAS_DATA, + RealizationStorageState.INITIALIZED, ] self.send_event( RunModelStatusEvent( diff --git a/src/ert/shared/hook_implementations/workflows/export_misfit_data.py b/src/ert/shared/hook_implementations/workflows/export_misfit_data.py index 9ef12af9ecd..93a12fdf46c 100644 --- a/src/ert/shared/hook_implementations/workflows/export_misfit_data.py +++ b/src/ert/shared/hook_implementations/workflows/export_misfit_data.py @@ -2,7 +2,7 @@ from ert import ErtScript from ert.exceptions import StorageError -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState class ExportMisfitDataJob(ErtScript): @@ -24,7 +24,7 @@ def run(self, target_file: Optional[str] = None) -> None: if self.ensemble is None: raise StorageError("No responses loaded") - realizations = self.ensemble.realization_list(RealizationState.HAS_DATA) + realizations = self.ensemble.realization_list(RealizationStorageState.HAS_DATA) if not realizations: raise StorageError("No responses loaded") diff --git a/src/ert/shared/share/ert/workflows/jobs/internal-gui/scripts/gen_data_rft_export.py b/src/ert/shared/share/ert/workflows/jobs/internal-gui/scripts/gen_data_rft_export.py index a766c5b59fd..650d627cbfc 100644 --- a/src/ert/shared/share/ert/workflows/jobs/internal-gui/scripts/gen_data_rft_export.py +++ b/src/ert/shared/share/ert/workflows/jobs/internal-gui/scripts/gen_data_rft_export.py @@ -7,15 +7,12 @@ from PyQt5.QtWidgets import QCheckBox from ert import LibresFacade -from ert.config import ( - CancelPluginException, - ErtPlugin, -) +from ert.config import CancelPluginException, ErtPlugin from ert.gui.ertwidgets.customdialog import CustomDialog from ert.gui.ertwidgets.listeditbox import ListEditBox from ert.gui.ertwidgets.models.path_model import PathModel from ert.gui.ertwidgets.pathchooser import PathChooser -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState def load_args(filename, column_names=None): @@ -176,7 +173,9 @@ def run( ) rft_data = facade.load_gen_data(ensemble, data_key, report_step) - realizations = ensemble.realization_list(RealizationState.HAS_DATA) + realizations = ensemble.realization_list( + RealizationStorageState.HAS_DATA + ) # Trajectory trajectory_file = os.path.join(trajectory_path, f"{well}.txt") diff --git a/src/ert/simulator/simulation_context.py b/src/ert/simulator/simulation_context.py index 488936e6b73..7660248a488 100644 --- a/src/ert/simulator/simulation_context.py +++ b/src/ert/simulator/simulation_context.py @@ -11,7 +11,7 @@ from ert.config import HookRuntime from ert.enkf_main import create_run_path from ert.job_queue import JobQueue, JobStatus -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.run_context import RunContext from ert.runpaths import Runpaths @@ -38,10 +38,10 @@ def _run_forward_model( run_context.sim_fs.update_realization_state( realization_nr, [ - RealizationState.UNDEFINED, - RealizationState.LOAD_FAILURE, + RealizationStorageState.UNDEFINED, + RealizationStorageState.LOAD_FAILURE, ], - RealizationState.INITIALIZED, + RealizationStorageState.INITIALIZED, ) run_context.sim_fs.sync() @@ -112,7 +112,7 @@ def __init__( for realization_nr in self._run_context.active_realizations: self._run_context.sim_fs.state_map[ realization_nr - ] = RealizationState.INITIALIZED + ] = RealizationStorageState.INITIALIZED create_run_path(self._run_context, global_substitutions, self._ert.ert_config) self._ert.runWorkflows( HookRuntime.PRE_SIMULATION, None, self._run_context.sim_fs diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index 7a18bc7ee28..0941ec92402 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -15,7 +15,7 @@ from ert.callbacks import forward_model_ok from ert.load_status import LoadResult, LoadStatus -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState if TYPE_CHECKING: import numpy.typing as npt @@ -37,14 +37,14 @@ def _load_realization( ) -> Tuple[LoadResult, int]: sim_fs.update_realization_state( realisation, - [RealizationState.UNDEFINED], - RealizationState.INITIALIZED, + [RealizationStorageState.UNDEFINED], + RealizationStorageState.INITIALIZED, ) result = forward_model_ok(run_args[realisation]) sim_fs.state_map[realisation] = ( - RealizationState.HAS_DATA + RealizationStorageState.HAS_DATA if result.status == LoadStatus.LOAD_SUCCESSFUL - else RealizationState.LOAD_FAILURE + else RealizationStorageState.LOAD_FAILURE ) return result, realisation @@ -101,7 +101,7 @@ def iteration(self) -> int: return self._index.iteration @property - def state_map(self) -> List[RealizationState]: + def state_map(self) -> List[RealizationStorageState]: return self._state_map @property @@ -115,30 +115,32 @@ def sync(self) -> None: pass def get_realization_mask_from_state( - self, states: List[RealizationState] + self, states: List[RealizationStorageState] ) -> npt.NDArray[np.bool_]: return np.array([s in states for s in self._state_map], dtype=bool) - def _load_state_map(self) -> List[RealizationState]: + def _load_state_map(self) -> List[RealizationStorageState]: state_map_file = self._experiment_path / "state_map.json" if state_map_file.exists(): with open(state_map_file, "r", encoding="utf-8") as f: data = json.load(f) - return [RealizationState(v) for v in data["state_map"]] + return [RealizationStorageState(v) for v in data["state_map"]] else: - return [RealizationState.UNDEFINED for _ in range(self.ensemble_size)] + return [ + RealizationStorageState.UNDEFINED for _ in range(self.ensemble_size) + ] @property def is_initalized(self) -> bool: - return RealizationState.INITIALIZED in self.state_map or self.has_data + return RealizationStorageState.INITIALIZED in self.state_map or self.has_data @property def has_data(self) -> bool: - return RealizationState.HAS_DATA in self.state_map + return RealizationStorageState.HAS_DATA in self.state_map def realizations_initialized(self, realizations: List[int]) -> bool: initialized_realizations = set( - self.realization_list(RealizationState.INITIALIZED) + self.realization_list(RealizationStorageState.INITIALIZED) ) return all(real in initialized_realizations for real in realizations) @@ -154,7 +156,7 @@ def get_summary_keyset(self) -> List[str]: keys = sorted(response["name"].values) return keys - def realization_list(self, state: RealizationState) -> List[int]: + def realization_list(self, state: RealizationStorageState) -> List[int]: """ Will return list of realizations with state == the specified state. """ @@ -271,8 +273,8 @@ def _save_state_map(self) -> None: def update_realization_state( self, realization: int, - old_states: List[RealizationState], - new_state: RealizationState, + old_states: List[RealizationStorageState], + new_state: RealizationStorageState, ) -> None: if self._state_map[realization] in old_states: self._state_map[realization] = new_state @@ -304,7 +306,7 @@ def load_from_run_path( if status == LoadStatus.LOAD_SUCCESSFUL: loaded += 1 - self.state_map[iens] = RealizationState.HAS_DATA + self.state_map[iens] = RealizationStorageState.HAS_DATA else: logger.error(f"Realization: {iens}, load failure: {message}") @@ -340,10 +342,10 @@ def save_parameters( self.update_realization_state( realization, [ - RealizationState.UNDEFINED, - RealizationState.LOAD_FAILURE, + RealizationStorageState.UNDEFINED, + RealizationStorageState.LOAD_FAILURE, ], - RealizationState.INITIALIZED, + RealizationStorageState.INITIALIZED, ) def save_response(self, group: str, data: xr.Dataset, realization: int) -> None: diff --git a/src/ert/storage/local_storage.py b/src/ert/storage/local_storage.py index d1e719ae6fb..44f261afd00 100644 --- a/src/ert/storage/local_storage.py +++ b/src/ert/storage/local_storage.py @@ -28,7 +28,7 @@ from pydantic import BaseModel, Field from ert.config import ErtConfig -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.shared import __version__ from ert.storage.local_ensemble import LocalEnsembleAccessor, LocalEnsembleReader from ert.storage.local_experiment import LocalExperimentAccessor, LocalExperimentReader @@ -309,11 +309,13 @@ def create_ensemble( state_map = prior_ensemble.state_map for realization_nr, state in enumerate(state_map[: len(ens.state_map)]): if state in [ - RealizationState.LOAD_FAILURE, - RealizationState.PARENT_FAILURE, - RealizationState.UNDEFINED, + RealizationStorageState.LOAD_FAILURE, + RealizationStorageState.PARENT_FAILURE, + RealizationStorageState.UNDEFINED, ]: - ens.state_map[realization_nr] = RealizationState.PARENT_FAILURE + ens.state_map[ + realization_nr + ] = RealizationStorageState.PARENT_FAILURE self._ensembles[ens.id] = ens return ens diff --git a/src/ert/storage/migration/block_fs.py b/src/ert/storage/migration/block_fs.py index bf105eced81..158e4c2aaf6 100644 --- a/src/ert/storage/migration/block_fs.py +++ b/src/ert/storage/migration/block_fs.py @@ -24,7 +24,7 @@ SurfaceConfig, field_transform, ) -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.storage import EnsembleAccessor, StorageAccessor from ert.storage.local_storage import LocalStorageAccessor, local_storage_get_ert_config from ert.storage.migration._block_fs_native import DataFile, Kind @@ -162,7 +162,7 @@ def _load_timestamps(path: Path) -> npt.NDArray[np.datetime64]: return np.frombuffer(f.read(size * sizeof_time_t), dtype="datetime64[s]") -def _load_states(path: Path) -> List[RealizationState]: +def _load_states(path: Path) -> List[RealizationStorageState]: if not path.exists(): return [] @@ -172,13 +172,13 @@ def _load_states(path: Path) -> List[RealizationState]: size = struct.unpack("I", f.read(4))[0] f.read(sizeof_int) # int default_value; (unused) return [ - RealizationState(x) + RealizationStorageState(x) for x in np.frombuffer(f.read(size * sizeof_int), dtype=np.int32) ] def _copy_state_map( - ensemble: EnsembleAccessor, states: Sequence[RealizationState] + ensemble: EnsembleAccessor, states: Sequence[RealizationStorageState] ) -> None: for index, state in enumerate(states): ensemble.state_map[index] = state @@ -351,7 +351,9 @@ def _migrate_gen_data( ) ) ensemble.save_response( - name, xr.combine_by_coords(datasets), iens # type: ignore + name, + xr.combine_by_coords(datasets), # type: ignore + iens, ) diff --git a/tests/performance_tests/test_memory_usage.py b/tests/performance_tests/test_memory_usage.py index 5f94ee9b884..00abeb7f762 100644 --- a/tests/performance_tests/test_memory_usage.py +++ b/tests/performance_tests/test_memory_usage.py @@ -10,12 +10,9 @@ from flaky import flaky from ert.analysis import UpdateConfiguration, smoother_update -from ert.config import ( - ErtConfig, - SummaryConfig, -) +from ert.config import ErtConfig, SummaryConfig from ert.enkf_main import sample_prior -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.storage import open_storage from tests.performance_tests.performance_utils import make_poly_example @@ -100,7 +97,7 @@ def fill_storage_with_data(poly_template: Path, ert_config: ErtConfig) -> None: ), real, ) - source.state_map[real] = RealizationState.HAS_DATA + source.state_map[real] = RealizationStorageState.HAS_DATA sample_prior(source, realizations, ens_config.parameters) diff --git a/tests/unit_tests/analysis/test_es_update.py b/tests/unit_tests/analysis/test_es_update.py index e3bbf4e0710..26fab265186 100644 --- a/tests/unit_tests/analysis/test_es_update.py +++ b/tests/unit_tests/analysis/test_es_update.py @@ -26,7 +26,7 @@ from ert.cli.main import run_cli from ert.config import AnalysisConfig, ErtConfig, GenDataConfig, GenKwConfig from ert.config.analysis_module import ESSettings, IESSettings -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.storage import open_storage @@ -433,7 +433,7 @@ def test_snapshot_alpha( ) rng = np.random.default_rng(1234) for iens in range(prior.ensemble_size): - prior.state_map[iens] = RealizationState.HAS_DATA + prior.state_map[iens] = RealizationStorageState.HAS_DATA data = rng.uniform(0, 1) prior.save_parameters( "PARAMETER", @@ -632,7 +632,7 @@ def test_gen_data_obs_data_mismatch(storage, uniform_parameter, update_config): ) rng = np.random.default_rng(1234) for iens in range(prior.ensemble_size): - prior.state_map[iens] = RealizationState.HAS_DATA + prior.state_map[iens] = RealizationStorageState.HAS_DATA data = rng.uniform(0, 1) prior.save_parameters( "PARAMETER", @@ -688,7 +688,7 @@ def test_gen_data_missing(storage, update_config, uniform_parameter, obs): ) rng = np.random.default_rng(1234) for iens in range(prior.ensemble_size): - prior.state_map[iens] = RealizationState.HAS_DATA + prior.state_map[iens] = RealizationStorageState.HAS_DATA data = rng.uniform(0, 1) prior.save_parameters( "PARAMETER", diff --git a/tests/unit_tests/gui/tools/test_case_tool.py b/tests/unit_tests/gui/tools/test_case_tool.py index a2fc12b56fd..07fc99d75e3 100644 --- a/tests/unit_tests/gui/tools/test_case_tool.py +++ b/tests/unit_tests/gui/tools/test_case_tool.py @@ -7,7 +7,7 @@ from ert.gui.tools.manage_cases.case_init_configuration import ( CaseInitializationConfigurationPanel, ) -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState @pytest.mark.usefixtures("copy_poly_case") @@ -24,7 +24,7 @@ def test_case_tool_init_prior(qtbot, storage): notifier.set_current_case(ensemble) assert ( ensemble.state_map - == [RealizationState.UNDEFINED] * config.model_config.num_realizations + == [RealizationStorageState.UNDEFINED] * config.model_config.num_realizations ) tool = CaseInitializationConfigurationPanel( config, notifier, config.model_config.num_realizations @@ -35,7 +35,7 @@ def test_case_tool_init_prior(qtbot, storage): ) assert ( ensemble.state_map - == [RealizationState.INITIALIZED] * config.model_config.num_realizations + == [RealizationStorageState.INITIALIZED] * config.model_config.num_realizations ) diff --git a/tests/unit_tests/simulator/test_simulation_context.py b/tests/unit_tests/simulator/test_simulation_context.py index c5f65a91f9f..0b1effadfb4 100644 --- a/tests/unit_tests/simulator/test_simulation_context.py +++ b/tests/unit_tests/simulator/test_simulation_context.py @@ -1,5 +1,5 @@ from ert.enkf_main import EnKFMain -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.simulator import SimulationContext from tests.utils import wait_until @@ -60,10 +60,10 @@ def test_simulation_context(setup_case, storage): assert not even_ctx.didRealizationFail(iens) assert even_ctx.isRealizationFinished(iens) - assert even_half.state_map[iens] == RealizationState.HAS_DATA + assert even_half.state_map[iens] == RealizationStorageState.HAS_DATA else: assert odd_ctx.didRealizationSucceed(iens) assert not odd_ctx.didRealizationFail(iens) assert odd_ctx.isRealizationFinished(iens) - assert odd_half.state_map[iens] == RealizationState.HAS_DATA + assert odd_half.state_map[iens] == RealizationStorageState.HAS_DATA diff --git a/tests/unit_tests/status/test_tracking_integration.py b/tests/unit_tests/status/test_tracking_integration.py index 2bb596ef3ef..b744f3f6644 100644 --- a/tests/unit_tests/status/test_tracking_integration.py +++ b/tests/unit_tests/status/test_tracking_integration.py @@ -30,7 +30,7 @@ JOB_STATE_START, REALIZATION_STATE_FINISHED, ) -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.shared.feature_toggling import FeatureToggling @@ -74,7 +74,7 @@ def check_expression(original, path_expression, expected, msg_start): "The run is cancelled due to reaching MAX_RUNTIME", ), ], - [RealizationState.LOAD_FAILURE] * 2, + [RealizationStorageState.LOAD_FAILURE] * 2, id="ee_poly_experiment_cancelled_by_max_runtime", ), pytest.param( @@ -90,7 +90,7 @@ def check_expression(original, path_expression, expected, msg_start): 1, 1.0, [(".*", "reals.*.jobs.*.status", JOB_STATE_FINISHED)], - [RealizationState.HAS_DATA] * 2, + [RealizationStorageState.HAS_DATA] * 2, id="ee_poly_experiment", ), pytest.param( @@ -108,7 +108,7 @@ def check_expression(original, path_expression, expected, msg_start): 2, 1.0, [(".*", "reals.*.jobs.*.status", JOB_STATE_FINISHED)], - [RealizationState.HAS_DATA] * 2, + [RealizationStorageState.HAS_DATA] * 2, id="ee_poly_smoother", ), pytest.param( @@ -132,8 +132,8 @@ def check_expression(original, path_expression, expected, msg_start): (".*", "reals.'1'.jobs.*.status", JOB_STATE_FINISHED), ], [ - RealizationState.LOAD_FAILURE, - RealizationState.HAS_DATA, + RealizationStorageState.LOAD_FAILURE, + RealizationStorageState.HAS_DATA, ], id="ee_failing_poly_smoother", ), diff --git a/tests/unit_tests/test_load_forward_model.py b/tests/unit_tests/test_load_forward_model.py index 5fa5d779c5d..c39a1d3e9b8 100644 --- a/tests/unit_tests/test_load_forward_model.py +++ b/tests/unit_tests/test_load_forward_model.py @@ -12,7 +12,7 @@ from ert.config import ErtConfig from ert.enkf_main import create_run_path, ensemble_context from ert.libres_facade import LibresFacade -from ert.realization_state import RealizationState +from ert.realization_state import RealizationStorageState from ert.storage import open_storage @@ -87,7 +87,7 @@ def test_load_inconsistent_time_map_summary(caplog): storage = open_storage(facade.enspath, mode="w") ensemble = storage.get_ensemble_by_name("default_0") assert ( - ensemble.state_map[realisation_number] == RealizationState.HAS_DATA + ensemble.state_map[realisation_number] == RealizationStorageState.HAS_DATA ) # Check prior state # Create a result that is incompatible with the refcase @@ -128,7 +128,7 @@ def test_load_forward_model(snake_oil_default_storage): loaded = facade.load_from_forward_model(default, realizations, 0) assert loaded == 1 assert ( - default.state_map[realisation_number] == RealizationState.HAS_DATA + default.state_map[realisation_number] == RealizationStorageState.HAS_DATA ) # Check that status is as expected