diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index 49d8a04044f..8f740f5c8d0 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -147,6 +147,7 @@ def add(self, realization: int, entries: Set[Tuple[str, str, bool]]): if not isinstance(entries, set): print("wtf") + # print(f"Adding entries: {entries}") real_state.update(entries) def has(self, realization: int, key: str): @@ -161,19 +162,41 @@ def has(self, realization: int, key: str): or any(_has_it and key == _key for _, _key, _has_it in state) ) - def has_true_or_false(self, realization: int, key: str): + def has_entry(self, realization: int, key: str): if realization not in self._states: return False state = self._states[realization] - - return ( + # print( + # f"has_entry({realization},{key})=???, size of states for real" + # f" {realization}: {len(state)}" + # ) + # if (key, key, True) in state: + # print(f"({key}, {key}, True) in state") + # return True + # if (key, key, False) in state: + # print(f"({key}, {key}, False) in state") + # return True + # if any(key == _group for _group, _, _ in state): + # print(f"any({key} == _group for _group, _, _ in state)") + # return True + # if any(key == _key for _, _key, _ in state): + # print(f"any({key} == _key for _, _key, _ in state)") + # return True + result = ( (key, key, True) in state or (key, key, False) in state or any(key == _group for _group, _, _ in state) or any(key == _key for _, _key, _ in state) ) + # print( + # f"has_entry({realization},{key})={result}, size of states for real" + # f" {realization}: {len(state)}" + # ) + # print("return False") + return result + class LocalEnsemble(BaseMode): """ @@ -1155,24 +1178,49 @@ def _refresh_response_state(self, response_key: str, realization: int): # 2. We saved (then everything is OK, it is in memory) # 3. We open a storage and have a saved file (OK) - if self._realization_states.has_true_or_false(realization, response_key): + if self._realization_states.has_entry(realization, response_key): return combined_ds_key = self._find_unified_dataset_for_response(response_key) + # ex: combined_ds_key == gen_data, response_key = WOPR_OP1 + # ex2: response_key = summary, combined_ds_key = summary + is_grouped_ds = combined_ds_key == response_key + + t0 = time.time() if self.has_combined_response_dataset(response_key): - # Refresh from combined for all reals wrt. name ds = xr.open_dataset(self._path / f"{combined_ds_key}.nc") + if is_grouped_ds: + for _real in range(self.ensemble_size): + _reals_with_response = set(ds["realization"].values) + self._realization_states.add( + _real, + { + ( + combined_ds_key, + combined_ds_key, + _real in _reals_with_response, + ) + }, + ) + return + all_names = set(ds["name"].values) - print(f"About to refresh for names: \n {', '.join(all_names)}") - for _key, _ds in ( - ds["values"] - .drop_vars(set(ds["values"].dims) - {"realization", "name"}) - .groupby("name") - ): - assert isinstance(_key, str) - print(f"Refreshing for {_key}") + # print( + # f"About to refresh for ({len(all_names)}) keys in {combined_ds_key}: \n {', '.join(all_names)}" + # ) + for _i, _key in enumerate(ds["name"].values): + _ds = ds.sel(name=_key, drop=True) + # for _key, _ds in ( + # ds["values"] + # .drop_vars(set(ds["values"].dims) - {"realization", "name"}) + # .groupby("name") + # ): + # print( + # f"({_i}/{len(all_names)})Refreshing for {_key}, " + # f"time elapsed: [{round(time.time() - t0, 1)}s]" + # ) # We can assume that all reals not in _real are false reals_with_response = set( _ds.dropna("realization", how="all")["realization"].values @@ -1182,10 +1230,6 @@ def _refresh_response_state(self, response_key: str, realization: int): self._realization_states.add( _real, {(combined_ds_key, _key, _real in reals_with_response)} ) - assert self._realization_states.has_true_or_false(_real, _key) - - if _real in reals_with_response: - assert self._realization_states.has(_real, _key) return @@ -1195,10 +1239,6 @@ def _refresh_response_state(self, response_key: str, realization: int): has_realization_dir = os.path.exists(self._realization_dir(realization)) - # ex: combined_ds_key == gen_data, response_key = WOPR_OP1 - # ex2: response_key = summary, combined_ds_key = summary - is_grouped_ds = combined_ds_key == response_key - if not has_realization_dir: self._realization_states.add( realization,