diff --git a/src/ert/analysis/_es_update.py b/src/ert/analysis/_es_update.py index 0bc9bf08b1c..00a08421f8e 100644 --- a/src/ert/analysis/_es_update.py +++ b/src/ert/analysis/_es_update.py @@ -506,8 +506,9 @@ def _copy_unupdated_parameters( # Copy the non-updated parameter groups from source to target for each active realization for parameter_group in not_updated_parameter_groups: - ds = source_ensemble.load_parameters(parameter_group, iens_active_index) - target_ensemble.save_parameters(parameter_group, iens_active_index, ds) + for realization in iens_active_index: + ds = source_ensemble.load_parameters(parameter_group, realization) + target_ensemble.save_parameters(parameter_group, realization, ds) def analysis_ES( diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index b287a45585b..5672b55b80e 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -548,10 +548,10 @@ def _load_single_dataset( def _load_dataset( self, group: str, - realizations: Union[int, npt.NDArray[np.int_], None], + realizations: Union[int, np.int64, npt.NDArray[np.int_], None], ) -> xr.Dataset: - if isinstance(realizations, int): - return self._load_single_dataset(group, realizations).isel( + if isinstance(realizations, (int, np.int64)): + return self._load_single_dataset(group, int(realizations)).isel( realizations=0, drop=True ) @@ -792,7 +792,7 @@ def load_all_gen_kw_data( def save_parameters( self, group: str, - realization: Union[int, npt.NDArray[np.int_]], + realization: int, dataset: xr.Dataset, ) -> None: """ @@ -820,30 +820,13 @@ def save_parameters( if group not in self.experiment.parameter_configuration: raise ValueError(f"{group} is not registered to the experiment.") - # Convert to numpy array if it's an integer - realizations: npt.NDArray[np.int_] = ( - np.array([realization]) - if isinstance(realization, (int, np.integer)) - else np.asarray(realization) - ) - - if realizations.size > 1 and "realizations" not in dataset.dims: - raise ValueError( - "Dataset must have 'realizations' dimension when saving multiple realizations" - ) - + path = self._realization_dir(realization) / f"{_escape_filename(group)}.nc" + path.parent.mkdir(exist_ok=True) if "realizations" in dataset.dims: - dataset = dataset.sel(realizations=realizations) - for real, data_to_save in dataset.groupby("realizations"): - path = self._realization_dir(real) / f"{_escape_filename(group)}.nc" - path.parent.mkdir(exist_ok=True) - self._storage._to_netcdf_transaction(path, data_to_save) + data_to_save = dataset.sel(realizations=[realization]) else: - for real in realizations: - path = self._realization_dir(real) / f"{_escape_filename(group)}.nc" - path.parent.mkdir(exist_ok=True) - data_to_save = dataset.expand_dims(realizations=[real]) - self._storage._to_netcdf_transaction(path, data_to_save) + data_to_save = dataset.expand_dims(realizations=[realization]) + self._storage._to_netcdf_transaction(path, data_to_save) @require_write def save_response( diff --git a/tests/ert/ui_tests/cli/test_field_parameter.py b/tests/ert/ui_tests/cli/test_field_parameter.py index b2262e3eb21..6c1c7cb193a 100644 --- a/tests/ert/ui_tests/cli/test_field_parameter.py +++ b/tests/ert/ui_tests/cli/test_field_parameter.py @@ -266,7 +266,8 @@ def test_field_param_update_using_heat_equation_zero_var_params_and_adaptive_loc name="prior-zero-var", ) cond["values"][:, :, :5, 0] = 1.0 - new_prior.save_parameters("COND", range(prior.ensemble_size), cond) + for real in range(prior.ensemble_size): + new_prior.save_parameters("COND", real, cond) # Copy responses from existing prior to new prior. # Note that we ideally should generate new responses by running the