From 873565489bfc886955aaac98af3c24489e5a0a56 Mon Sep 17 00:00:00 2001 From: Mathias Methlie Nilsen Date: Thu, 8 Oct 2026 15:23:01 +0200 Subject: [PATCH] Evaluate a single control vector on every model in robust optimization With num_models > 1, GaussianEnsemble.function built one simulator input for a single control vector but asked for num_models members, so _simulator_input raised IndexError. Repeat the control vector once per model and pair each copy with its own model index; the pairing that _aux_input sets for the gradient ensemble ([0, 0, 1, 1, ...]) would also have sent the copies to the wrong models. Add tests for both the single-point evaluation and the gradient pairing. Co-Authored-By: Claude Opus 5.5 --- src/popt/ensembles/ensemble_base.py | 11 +++- tests/optimization/test_robust_evaluation.py | 66 ++++++++++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) create mode 100644 tests/optimization/test_robust_evaluation.py diff --git a/src/popt/ensembles/ensemble_base.py b/src/popt/ensembles/ensemble_base.py index 6ef3b935..53ce1b38 100644 --- a/src/popt/ensembles/ensemble_base.py +++ b/src/popt/ensembles/ensemble_base.py @@ -106,8 +106,15 @@ def function(self, x, *args, **kwargs): if ensemble_input: self.ne = x.shape[1] else: - x = x[:, np.newaxis] - self.ne = self.num_models # In case of robust optimization + # A single control vector is evaluated on every geological model in + # robust optimization: one copy per model, each paired with its own + # model index. Without the copies the simulator got one input for + # num_models members, and `_aux_input`'s ensemble pairing + # ([0, 0, 1, 1, ...]) would have sent them to the wrong models. + x = np.repeat(x[:, np.newaxis], self.num_models, axis=1) + self.ne = self.num_models + if self.num_models > 1: + self.aux_input = list(range(self.num_models)) if isinstance(self.sim, noSimulation): func_values = self.obj_func(x, **kwargs) diff --git a/tests/optimization/test_robust_evaluation.py b/tests/optimization/test_robust_evaluation.py new file mode 100644 index 00000000..7c4bfd55 --- /dev/null +++ b/tests/optimization/test_robust_evaluation.py @@ -0,0 +1,66 @@ +"""Robust optimization: every control vector is evaluated on every geological model. + +With ``num_models > 1`` a single control vector must run once per model, each copy +with its own model index, and the gradient ensemble must pair its perturbations +with the models in the order the gradient subtracts the single-point values. +""" + +import numpy as np +import pandas as pd + +from popt import GaussianEnsemble + + +class _ModelRecorder: + """A simulator whose output is the sum of the controls plus 100 per model index.""" + + def __init__(self): + self.input_dict = {'datatype': ['v'], 'parallel': 1} + self.datatype = ['v'] + self.true_order = None + self.redund_sim = None + self.models = [] + + def setup_fwd_run(self, **kwargs): + pass + + def run_fwd_sim(self, state, member_index): + self.models.append(int(state['aux_input'])) + return pd.DataFrame({'v': [float(np.sum(state['x'])) + 100 * state['aux_input']]}, index=[0]) + + +def _objective(data, **kwargs): + return np.array([np.asarray(v).ravel() for v in data['v']]).ravel() + + +def _ensemble(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + sim = _ModelRecorder() + keys = { + 'ne': 10, 'num_models': 5, 'seed': 1, 'save_prior': False, 'disable_tqdm': True, + 'controls': {'x': {'mean': [1.0, 2.0], 'var': 0.01, 'limits': [0, 10]}}, + } + return GaussianEnsemble(keys, sim, _objective), sim + + +def test_a_single_control_vector_runs_on_every_model(tmp_path, monkeypatch): + ensemble, sim = _ensemble(tmp_path, monkeypatch) + + values = ensemble.function(ensemble.get_state()) + + assert sim.models == [0, 1, 2, 3, 4] + np.testing.assert_allclose(values, [3.0, 103.0, 203.0, 303.0, 403.0]) + + +def test_the_gradient_pairs_each_perturbation_with_its_model(tmp_path, monkeypatch): + ensemble, sim = _ensemble(tmp_path, monkeypatch) + x = ensemble.get_state() + ensemble.function(x) + sim.models.clear() + + gradient = ensemble.gradient(x, ensemble.get_cov()) + + assert sim.models == [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] + # The model offsets cancel against the single-point values, leaving a small + # natural gradient of sum(x) rather than one of order 100. + assert np.all(np.abs(gradient) < 1.0)