diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index b32ccca7548..8d2ea69976b 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -225,8 +225,8 @@ def get_realization_mask_with_parameters(self) -> npt.NDArray[np.bool_]: return np.array( [ - self._parameters_exist_for_realization(i) - for i in range(self.ensemble_size) + state == RealizationStorageState.INITIALIZED + for state in self.get_ensemble_state() ] ) @@ -249,73 +249,11 @@ def get_realization_mask_with_responses( return np.array( [ - self._responses_exist_for_realization(i, key) - for i in range(self.ensemble_size) + state == RealizationStorageState.HAS_DATA + for state in self.get_ensemble_state() ] ) - def _parameters_exist_for_realization(self, realization: int) -> bool: - """ - Returns true if all parameters in the experiment have - all been saved in the ensemble. If no parameters, return True - - Parameters - ---------- - realization : int - Realization index. - - Returns - ------- - exists : bool - True if parameters exist for realization. - """ - if not self.experiment.parameter_configuration: - return True - path = self._realization_dir(realization) - return all( - (path / (_escape_filename(parameter) + ".nc")).exists() - for parameter in self.experiment.parameter_configuration - ) - - def _responses_exist_for_realization( - self, realization: int, key: Optional[str] = None - ) -> bool: - """ - Returns true if there are responses in the experiment and they have - all been saved in the ensemble - - Parameters - ---------- - realization : int - Realization index. - key : str, optional - Response key to filter realizations. If None, all responses are considered. - - Returns - ------- - exists : bool - True if responses exist for realization. - """ - - if not self.experiment.response_configuration: - return True - path = self._realization_dir(realization) - - def _has_response(_key: str) -> bool: - if _key in self.experiment.response_key_to_response_type: - _response_type = self.experiment.response_key_to_response_type[_key] - return (path / f"{_response_type}.parquet").exists() - - return (path / f"{_key}.parquet").exists() - - if key: - return _has_response(key) - - return all( - _has_response(response) - for response in self.experiment.response_configuration - ) - def is_initalized(self) -> List[int]: """ Return the realization numbers where all parameters are internalized. In @@ -488,14 +426,76 @@ def get_ensemble_state(self) -> List[RealizationStorageState]: List of realization states. """ + def _parameters_exist_for_realization(realization: int) -> bool: + """ + Returns true if all parameters in the experiment have + all been saved in the ensemble. If no parameters, return True + + Parameters + ---------- + realization : int + Realization index. + + Returns + ------- + exists : bool + True if parameters exist for realization. + """ + if not self.experiment.parameter_configuration: + return True + path = self._realization_dir(realization) + return all( + (path / (_escape_filename(parameter) + ".nc")).exists() + for parameter in self.experiment.parameter_configuration + ) + + def _responses_exist_for_realization( + realization: int, key: Optional[str] = None + ) -> bool: + """ + Returns true if there are responses in the experiment and they have + all been saved in the ensemble + + Parameters + ---------- + realization : int + Realization index. + key : str, optional + Response key to filter realizations. If None, all responses are considered. + + Returns + ------- + exists : bool + True if responses exist for realization. + """ + + if not self.experiment.response_configuration: + return True + path = self._realization_dir(realization) + + def _has_response(_key: str) -> bool: + if _key in self.experiment.response_key_to_response_type: + _response_type = self.experiment.response_key_to_response_type[_key] + return (path / f"{_response_type}.parquet").exists() + + return (path / f"{_key}.parquet").exists() + + if key: + return _has_response(key) + + return all( + _has_response(response) + for response in self.experiment.response_configuration + ) + def _find_state(realization: int) -> RealizationStorageState: if self.has_failure(realization): failure = self.get_failure(realization) assert failure return failure.type - if self._responses_exist_for_realization(realization): + if _responses_exist_for_realization(realization): return RealizationStorageState.HAS_DATA - if self._parameters_exist_for_realization(realization): + if _parameters_exist_for_realization(realization): return RealizationStorageState.INITIALIZED else: return RealizationStorageState.UNDEFINED