From 0e48ef241cfbadcd570edbe1cd525598edccdd02 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Wed, 6 May 2026 19:26:13 -0700 Subject: [PATCH 1/9] grad_rho: aggregate per-sub-scenario nonants in proper bundles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pre-fix _get_grad_exprs differentiated the bundle objective wrt s._mpisppy_data.nonant_indices.values() — bundle ref Vars. create_EF picks the first sub-scenario's Var as the ref and links the rest by equality constraints, so this captured only the representative sub-scenario's contribution (refs #673). _eval_grad_exprs likewise set only the bundle ref Var, leaving per-sub-scenario nonants stale. Differentiate wrt every per-sub-scenario nonant Var, sum the resulting partial expressions per bundle position (ndn, k), and on evaluation propagate xhat to all per-sub-scenario nonants tied to that position. Mirrors sputils.nonant_cost_coeffs's bundle branch. Adds test_grad_rho_bundles.py with three solver-free unit tests: linear-obj baseline, bundled equivalence with explicit "not representative-only" guard, and a quadratic-obj test that primes stale per-sub-scenario var values to verify _eval_grad_exprs overrides them. Wired into run_coverage.bash and the unit tests block of test_pr_and_main.yml. sensi_rho and reduced_costs_rho still need analogous fixes per #673. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test_pr_and_main.yml | 3 +- mpisppy/extensions/grad_rho.py | 90 +++++++-- mpisppy/tests/test_grad_rho_bundles.py | 255 +++++++++++++++++++++++++ run_coverage.bash | 3 + 4 files changed, 337 insertions(+), 14 deletions(-) create mode 100644 mpisppy/tests/test_grad_rho_bundles.py diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml index 0f6f3be48..392591d92 100644 --- a/.github/workflows/test_pr_and_main.yml +++ b/.github/workflows/test_pr_and_main.yml @@ -1013,7 +1013,8 @@ jobs: mpisppy/tests/test_solver_spec.py \ mpisppy/tests/test_extensions.py \ mpisppy/tests/test_jensens.py \ - mpisppy/tests/test_proper_bundler.py + mpisppy/tests/test_proper_bundler.py \ + mpisppy/tests/test_grad_rho_bundles.py - name: Upload coverage data if: always() diff --git a/mpisppy/extensions/grad_rho.py b/mpisppy/extensions/grad_rho.py index 01dba2e04..ea3a010b2 100644 --- a/mpisppy/extensions/grad_rho.py +++ b/mpisppy/extensions/grad_rho.py @@ -129,31 +129,95 @@ def _scen_indep_denom(self): return scen_indep_denom def _get_grad_exprs(self): - """ Grabs and caches the gradient expressions for each scenario's objective (without proximal term). """ + """ Grabs and caches the gradient expressions for each scenario's objective (without proximal term). + + Proper bundles need special handling: ``s._mpisppy_data.nonant_indices`` + holds the bundle's ref Vars, but those ref Vars *are* one specific + sub-scenario's nonant Vars (sputils.create_EF picks the first scenario's + Var as the ref); other sub-scenarios' nonants are tied to the ref by + equality constraints. Differentiating the bundle objective w.r.t. the + ref Vars therefore captures only that "representative" sub-scenario's + contribution. Instead, differentiate w.r.t. every per-sub-scenario + nonant Var and sum the resulting partials per bundle position + ``(ndn, k)``. This mirrors the bundle branch in + ``sputils.nonant_cost_coeffs``. + """ self.grad_exprs = dict() + self._bundle_per_scen_vars = dict() for s in self.opt.local_scenarios.values(): - self.grad_exprs[s] = differentiate(sputils.find_active_objective(s), - wrt_list=s._mpisppy_data.nonant_indices.values(), - mode=Modes.reverse_symbolic, - ) + active_obj = sputils.find_active_objective(s) + + if hasattr(s, "_ef_scenario_names"): + per_scen_to_bundle = {} + counters = {} + for (ndn, per_scen_i) in s.ref_vars.keys(): + if (ndn, per_scen_i) in s.ref_surrogate_vars: + continue + counters.setdefault(ndn, 0) + per_scen_to_bundle[(ndn, per_scen_i)] = (ndn, counters[ndn]) + counters[ndn] += 1 + + wrt_vars = [] + wrt_keys = [] + per_scen_vars = {} + for scenario_name in s._ef_scenario_names: + scenario = s.component(scenario_name) + for node in scenario._mpisppy_node_list: + ndn = node.name + for per_scen_i, v in enumerate(node.nonant_vardata_list): + bundle_key = per_scen_to_bundle.get((ndn, per_scen_i)) + if bundle_key is None: + continue + wrt_vars.append(v) + wrt_keys.append(bundle_key) + per_scen_vars.setdefault(bundle_key, []).append(v) + + partials = differentiate(active_obj, + wrt_list=wrt_vars, + mode=Modes.reverse_symbolic, + ) + + grad_dict = {ndn_i: 0 for ndn_i in s._mpisppy_data.nonant_indices} + for bundle_key, partial in zip(wrt_keys, partials): + grad_dict[bundle_key] = grad_dict[bundle_key] + partial + self.grad_exprs[s] = grad_dict + self._bundle_per_scen_vars[s] = per_scen_vars + else: + partials = differentiate(active_obj, + wrt_list=s._mpisppy_data.nonant_indices.values(), + mode=Modes.reverse_symbolic, + ) + self.grad_exprs[s] = {ndn_i: partials[i] + for i, ndn_i in enumerate(s._mpisppy_data.nonant_indices)} - self.grad_exprs[s] = {ndn_i : self.grad_exprs[s][i] for i, ndn_i in enumerate(s._mpisppy_data.nonant_indices)} - - return + return def _eval_grad_exprs(self, s, xhat): """ Evaluates the gradient expressions of the objectives for scenario s at xhat (if available) or the current values. """ - - ci = 0 + grads = {} if self.eval_at_xhat: if True not in np.isnan(self.best_xhat_buf.value_array()): - for ndn_i, var in s._mpisppy_data.nonant_indices.items(): - var.value = xhat[ci] - ci += 1 + ci = 0 + if s in self._bundle_per_scen_vars: + # Bundle: the cached gradient expressions reference + # per-sub-scenario nonant Vars. All per-sub-scenario + # nonants tied to a given bundle position are equal at + # an NA-feasible xhat, so set every one of them to the + # bundle's xhat value at that position. + per_scen_vars = self._bundle_per_scen_vars[s] + for ndn_i, var in s._mpisppy_data.nonant_indices.items(): + var.value = xhat[ci] + for v in per_scen_vars.get(ndn_i, ()): + v.value = xhat[ci] + ci += 1 + else: + for ndn_i, var in s._mpisppy_data.nonant_indices.items(): + var.value = xhat[ci] + ci += 1 for ndn_i, var in s._mpisppy_data.nonant_indices.items(): grads[ndn_i] = pyo.value(self.grad_exprs[s][ndn_i]) diff --git a/mpisppy/tests/test_grad_rho_bundles.py b/mpisppy/tests/test_grad_rho_bundles.py new file mode 100644 index 000000000..8a7c8230a --- /dev/null +++ b/mpisppy/tests/test_grad_rho_bundles.py @@ -0,0 +1,255 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""GradRho gradient extraction must aggregate across all per-sub-scenario +nonants in a proper bundle, not just the bundle's representative ref Var. +See Pyomo/mpi-sppy issue #673. + +These are pure unit tests: they build tiny scenarios and a create_EF bundle +in-process, then drive ``GradRho._get_grad_exprs`` and ``_eval_grad_exprs`` +directly with a stub PH-like object. No solver and no MPI required. +""" + +import unittest + +import numpy as np +import pyomo.environ as pyo + +import mpisppy.utils.sputils as sputils +from mpisppy.extensions.grad_rho import GradRho + + +def _build_linear_scen(name, c_coefs, prob): + """Two-stage scenario with first-stage Var x[0..len-1] and linear + objective sum_i c_coefs[i] * x[i]. No second stage (not needed).""" + m = pyo.ConcreteModel(name=name) + m.x = pyo.Var(range(len(c_coefs)), bounds=(0, None), initialize=0.0) + m.obj = pyo.Objective( + expr=sum(c_coefs[i] * m.x[i] for i in range(len(c_coefs))) + ) + sputils.attach_root_node(m, 0, [m.x[i] for i in range(len(c_coefs))]) + m._mpisppy_probability = prob + return m + + +def _build_quadratic_scen(name, q_coefs, prob): + """Same shape as _build_linear_scen but with a quadratic objective so + the gradient depends on the linearization point.""" + m = pyo.ConcreteModel(name=name) + m.x = pyo.Var(range(len(q_coefs)), bounds=(None, None), initialize=0.0) + m.obj = pyo.Objective( + expr=sum(q_coefs[i] * m.x[i] ** 2 for i in range(len(q_coefs))) + ) + sputils.attach_root_node(m, 0, [m.x[i] for i in range(len(q_coefs))]) + m._mpisppy_probability = prob + return m + + +def _attach_spbase_metadata(scen): + """Populate the ``_mpisppy_data`` Block with the fields SPBase normally + sets (nlens, nonant_indices, varid_to_nonant_index). create_EF already + creates the Block on the bundle; for bare scenarios we create it.""" + if not hasattr(scen, "_mpisppy_data"): + scen._mpisppy_data = pyo.Block( + name="For non-Pyomo mpi-sppy data" + ) + scen._mpisppy_data.nlens = { + node.name: len(node.nonant_vardata_list) + for node in scen._mpisppy_node_list + } + nonant_indices = {} + for node in scen._mpisppy_node_list: + for i, v in enumerate(node.nonant_vardata_list): + nonant_indices[(node.name, i)] = v + scen._mpisppy_data.nonant_indices = nonant_indices + scen._mpisppy_data.varid_to_nonant_index = { + id(v): k for k, v in nonant_indices.items() + } + + +def _build_bundle(scen_creator, scen_names): + """Build a proper-bundle-shaped EF model: create_EF, then attach the + ROOT node from non-surrogate ref_vars (mirroring proper_bundler), then + attach SPBase-style metadata on the bundle's _mpisppy_data Block.""" + bundle = sputils.create_EF(scen_names, scen_creator, EF_name="bundle") + nonantlist = [ + v for idx, v in bundle.ref_vars.items() + if idx[0] == "ROOT" and idx not in bundle.ref_surrogate_vars + ] + surrogates = [ + v for idx, v in bundle.ref_surrogate_vars.items() + if idx[0] == "ROOT" + ] + sputils.attach_root_node(bundle, 0, nonantlist, None, surrogates) + _attach_spbase_metadata(bundle) + return bundle + + +class _FakePH: + """Minimum surface area GradRho._get_grad_exprs touches on self.opt.""" + + def __init__(self, local_scenarios): + self.local_scenarios = local_scenarios + + +class _FakeBuf: + """Stand-in for the BEST_XHAT receive buffer; supplies value_array() + so the NaN-gating in _eval_grad_exprs evaluates as 'best xhat known'.""" + + def __init__(self, arr): + self._arr = np.asarray(arr, dtype=float) + + def value_array(self): + return self._arr + + +def _make_grad_rho(local_scenarios, *, eval_at_xhat=False, xhat_buf=None): + """Construct a GradRho-like object without invoking GradRho.__init__ + (which requires PH options/cfg). _get_grad_exprs and _eval_grad_exprs + only touch a handful of attributes; set them by hand.""" + g = GradRho.__new__(GradRho) + g.opt = _FakePH(local_scenarios) + g.eval_at_xhat = eval_at_xhat + if xhat_buf is not None: + g.best_xhat_buf = xhat_buf + return g + + +class TestGradRhoBundles(unittest.TestCase): + """Equivalence between bundled and non-bundled gradient extraction.""" + + def setUp(self): + # 3 scenarios x 3 first-stage vars, distinct linear cost coeffs. + self.scen_names = ["scen0", "scen1", "scen2"] + self.c = [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0], + ] + self.probs = [1.0 / 3, 1.0 / 3, 1.0 / 3] + + def _build_unbundled(self): + scens = {} + for name, c, p in zip(self.scen_names, self.c, self.probs): + s = _build_linear_scen(name, c, p) + _attach_spbase_metadata(s) + scens[name] = s + return scens + + def _build_bundled(self): + c_by_name = dict(zip(self.scen_names, self.c)) + prob_by_name = dict(zip(self.scen_names, self.probs)) + + def _scen_creator(sname, **kwargs): + return _build_linear_scen( + sname, c_by_name[sname], prob_by_name[sname] + ) + + return {"bundle": _build_bundle(_scen_creator, self.scen_names)} + + def test_unbundled_gradient_is_per_scenario_cost(self): + # Sanity floor: with a linear objective the gradient at every point + # equals the cost coefficient. Establishes the baseline that the + # bundle test compares against. + scens = self._build_unbundled() + g = _make_grad_rho(scens) + g._get_grad_exprs() + for sname, s in scens.items(): + i = self.scen_names.index(sname) + for k in range(len(self.c[i])): + got = pyo.value(g.grad_exprs[s][("ROOT", k)]) + self.assertAlmostEqual( + got, self.c[i][k], places=10, + msg=f"unbundled {sname=} k={k}: {got=} != {self.c[i][k]=}", + ) + + def test_bundle_gradient_is_conditional_prob_weighted_sum(self): + # The bundle objective from create_EF is + # EF_obj = (sum_s prob_s * obj_s) / EF_prob + # so the bundled gradient at position (ROOT, k) must be + # sum_s (prob_s / EF_prob) * c[s][k] + # i.e. the conditional-probability-weighted average of the + # per-sub-scenario cost coefficients. Before the fix for #673, + # differentiating wrt the ref Var only captured one sub-scenario + # (the first one in scen_names), so the bundled gradient would + # equal c[0][k] rather than this weighted sum. + bundled = self._build_bundled() + g = _make_grad_rho(bundled) + g._get_grad_exprs() + bundle = bundled["bundle"] + bundle_prob = bundle._mpisppy_probability + n_vars = len(self.c[0]) + for k in range(n_vars): + expected = sum( + (self.probs[s] / bundle_prob) * self.c[s][k] + for s in range(len(self.c)) + ) + got = pyo.value(g.grad_exprs[bundle][("ROOT", k)]) + self.assertAlmostEqual( + got, expected, places=10, + msg=f"bundle k={k}: {got=} != weighted-sum {expected=}", + ) + # Also confirm we are not accidentally returning the lone first + # sub-scenario's coefficient (the pre-fix behavior). + if self.c[0][k] != expected: + self.assertNotAlmostEqual( + got, self.c[0][k], places=6, + msg=f"bundle k={k} matches representative-only " + f"value {self.c[0][k]}; aggregation regressed", + ) + + def test_bundle_eval_propagates_xhat_to_subscenario_vars(self): + # With a quadratic objective the gradient depends on the + # linearization point, so this test fails if _eval_grad_exprs sets + # only the bundle ref Var and leaves the per-sub-scenario nonant + # Vars (which the cached gradient expressions actually reference) + # at their previous values. + scen_names = ["A", "B"] + q_by_name = {"A": [1.0, 2.0], "B": [3.0, 4.0]} + prob_by_name = {"A": 0.5, "B": 0.5} + + def _scen_creator(sname, **kwargs): + return _build_quadratic_scen( + sname, q_by_name[sname], prob_by_name[sname] + ) + + bundle = _build_bundle(_scen_creator, scen_names) + xhat = np.array([5.0, 7.0]) + g = _make_grad_rho( + {"bundle": bundle}, + eval_at_xhat=True, + xhat_buf=_FakeBuf(xhat), + ) + g._get_grad_exprs() + + # Force per-sub-scenario Vars to a stale value first; if the bundle + # branch of _eval_grad_exprs forgot to override them, the gradient + # would reflect the stale values rather than xhat. + for scenario_name in bundle._ef_scenario_names: + scen = bundle.component(scenario_name) + for v in scen.x.values(): + v.value = -999.0 + + grads = g._eval_grad_exprs(bundle, xhat) + + # bundle_obj = (0.5*sum_i qA[i]*x_Ai^2 + 0.5*sum_i qB[i]*x_Bi^2) / 1.0 + # d/dx_Ai (bundle_obj) = qA[i]*x_Ai + # d/dx_Bi (bundle_obj) = qB[i]*x_Bi + # aggregated grad at (ROOT, i) = qA[i]*x_Ai + qB[i]*x_Bi + # at x_Ai = x_Bi = xhat[i]: = (qA[i] + qB[i]) * xhat[i] + for i in range(2): + expected = (q_by_name["A"][i] + q_by_name["B"][i]) * xhat[i] + self.assertAlmostEqual( + grads[("ROOT", i)], expected, places=10, + msg=f"bundle eval at xhat[i={i}]: " + f"{grads[('ROOT', i)]=} != {expected=}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/run_coverage.bash b/run_coverage.bash index 25265b6cf..e72ad2d70 100755 --- a/run_coverage.bash +++ b/run_coverage.bash @@ -100,6 +100,9 @@ run_phase "test_generic_cylinders (serial)" \ run_phase "test_jensens (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_jensens.py -v +run_phase "test_grad_rho_bundles (serial)" \ + coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_grad_rho_bundles.py -v + run_phase "test_xhat_from_file (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_xhat_from_file.py -v From 135e1f8345199450b458b6501a67ac1fdb9d9608 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 05:29:03 -0700 Subject: [PATCH 2/9] sensi_rho: probe consensus direction for proper bundles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-fix nonant_sensitivies set y_vec = 1 only at the bundle ref Var index (sputils.create_EF picks the first sub-scenario's nonant Var as the ref) and called the KKT inverse on that direction. With the NA equality constraints `x_sub_j == ref` already in K, perturbing only the ref Var is infeasible; the regularized solve returned a representative-scenario-only sensitivity per #673. Add _bundle_consensus_groups: maps each bundle nonant position (ndn, k) to the per-sub-scenario nonant Vars coupled to it by NA equality constraints. In the bundle branch, set y_vec to 1 at every such index — a consensus perturbation, which is the feasible direction the formula is meaningful in. The Rayleigh-quotient form (g'K^{-1}y / y'K^{-1}y) is scale-invariant, so the resulting sensitivity is the rate of objective change per unit consensus shift and stays consistent with the unbundled scalar at that position. Mirrors the bundle branches in sputils.nonant_cost_coeffs and mpisppy.extensions.grad_rho. Tests in test_sensi_rho_bundles.py: - TestBundleConsensusGroups (solver-free, primary correctness check): for an N-sub-scenario bundle the helper returns N vars per position, the ref is one of them, and other sub-scenario Vars are present. Confirmed to fail when the helper is mocked to return only the ref. - TestNonantSensitiviesBundleEndToEnd (skipUnless Ipopt): smoke that the bundle path completes and a one-sub-scenario bundle agrees with the standalone scenario. Wired into run_coverage.bash and the unit tests block of test_pr_and_main.yml. reduced_costs_rho remains TODO under #673. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test_pr_and_main.yml | 3 +- mpisppy/tests/test_sensi_rho_bundles.py | 227 ++++++++++++++++++++++++ mpisppy/utils/nonant_sensitivities.py | 63 ++++++- run_coverage.bash | 3 + 4 files changed, 293 insertions(+), 3 deletions(-) create mode 100644 mpisppy/tests/test_sensi_rho_bundles.py diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml index 392591d92..0aebd0642 100644 --- a/.github/workflows/test_pr_and_main.yml +++ b/.github/workflows/test_pr_and_main.yml @@ -1014,7 +1014,8 @@ jobs: mpisppy/tests/test_extensions.py \ mpisppy/tests/test_jensens.py \ mpisppy/tests/test_proper_bundler.py \ - mpisppy/tests/test_grad_rho_bundles.py + mpisppy/tests/test_grad_rho_bundles.py \ + mpisppy/tests/test_sensi_rho_bundles.py - name: Upload coverage data if: always() diff --git a/mpisppy/tests/test_sensi_rho_bundles.py b/mpisppy/tests/test_sensi_rho_bundles.py new file mode 100644 index 000000000..e58aa8ad8 --- /dev/null +++ b/mpisppy/tests/test_sensi_rho_bundles.py @@ -0,0 +1,227 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""nonant_sensitivies must probe the consensus direction for proper bundles +(every per-sub-scenario nonant Var at each bundle position together) instead +of only the bundle's representative ref Var. See Pyomo/mpi-sppy issue #673. + +The first test class exercises _bundle_consensus_groups directly and runs +without a solver; it is the primary correctness check for the fix and +catches a regression to representative-only behavior. The second class +runs the full nonant_sensitivies pipeline on a real Ipopt solve and is +skipped automatically if Ipopt isn't installed (so it is informational +on environments without Ipopt and a regression check on those that have +it). +""" + +import unittest + +import pyomo.environ as pyo + +import mpisppy.utils.sputils as sputils +from mpisppy.utils.nonant_sensitivities import ( + _bundle_consensus_groups, + nonant_sensitivies, +) + + +def _ipopt_available(): + try: + return bool(pyo.SolverFactory("ipopt").available(exception_flag=False)) + except Exception: + return False + + +def _build_linear_scen(name, c_coefs, prob): + m = pyo.ConcreteModel(name=name) + m.x = pyo.Var(range(len(c_coefs)), bounds=(0, 10), initialize=0.0) + m.obj = pyo.Objective( + expr=sum(c_coefs[i] * m.x[i] for i in range(len(c_coefs))) + ) + sputils.attach_root_node(m, 0, [m.x[i] for i in range(len(c_coefs))]) + m._mpisppy_probability = prob + return m + + +def _attach_metadata(scen): + if not hasattr(scen, "_mpisppy_data"): + scen._mpisppy_data = pyo.Block( + name="For non-Pyomo mpi-sppy data" + ) + scen._mpisppy_data.nlens = { + node.name: len(node.nonant_vardata_list) + for node in scen._mpisppy_node_list + } + nonant_indices = {} + for node in scen._mpisppy_node_list: + for i, v in enumerate(node.nonant_vardata_list): + nonant_indices[(node.name, i)] = v + scen._mpisppy_data.nonant_indices = nonant_indices + + +def _build_bundle(scen_creator, scen_names): + bundle = sputils.create_EF(scen_names, scen_creator, EF_name="bundle") + nonantlist = [ + v for idx, v in bundle.ref_vars.items() + if idx[0] == "ROOT" and idx not in bundle.ref_surrogate_vars + ] + surrogates = [ + v for idx, v in bundle.ref_surrogate_vars.items() + if idx[0] == "ROOT" + ] + sputils.attach_root_node(bundle, 0, nonantlist, None, surrogates) + _attach_metadata(bundle) + return bundle + + +class TestBundleConsensusGroups(unittest.TestCase): + """Solver-free unit tests of the bundle var-grouping helper. + + The helper is what guarantees the KKT sensitivity probe perturbs every + per-sub-scenario nonant at a given bundle position together rather + than only the representative ref Var (issue #673). + """ + + def _build(self, scen_names, c_by_name, prob_by_name): + def _scen_creator(sname, **kwargs): + return _build_linear_scen( + sname, c_by_name[sname], prob_by_name[sname] + ) + return _build_bundle(_scen_creator, scen_names) + + def test_groups_collect_one_var_per_sub_scenario(self): + # 3 sub-scenarios x 2 first-stage vars: groups must have exactly + # 2 keys and exactly 3 vars per key (one from each sub-scenario). + # If the bug came back and the helper returned only the bundle + # ref Var, each group would have 1 var; this test fails loudly. + scen_names = ["scen0", "scen1", "scen2"] + c_by_name = { + "scen0": [1.0, 2.0], + "scen1": [3.0, 4.0], + "scen2": [5.0, 6.0], + } + prob_by_name = dict.fromkeys(scen_names, 1.0 / 3) + + bundle = self._build(scen_names, c_by_name, prob_by_name) + groups = _bundle_consensus_groups(bundle) + + self.assertEqual(set(groups.keys()), {("ROOT", 0), ("ROOT", 1)}) + for k, sub_vars in groups.items(): + self.assertEqual( + len(sub_vars), len(scen_names), + msg=f"position {k}: expected {len(scen_names)} sub-scenario " + f"vars, got {len(sub_vars)}", + ) + owners = {v.parent_component().parent_block() for v in sub_vars} + self.assertEqual( + len(owners), len(scen_names), + msg=f"position {k}: expected {len(scen_names)} distinct " + f"sub-scenario blocks, got {len(owners)}", + ) + + def test_first_var_in_group_is_bundle_ref(self): + # sputils.create_EF picks the first sub-scenario's nonant Var as + # the bundle ref. The helper must include that ref *and* the other + # sub-scenarios' nonants — not just the ref. This test pins both: + # the ref is in the group (sanity), and there is at least one + # other sub-scenario Var in the group that is not the ref + # (the bug-catching half). + scen_names = ["A", "B"] + c_by_name = {"A": [10.0, 20.0], "B": [30.0, 40.0]} + prob_by_name = {"A": 0.5, "B": 0.5} + + bundle = self._build(scen_names, c_by_name, prob_by_name) + groups = _bundle_consensus_groups(bundle) + + for k_idx in range(2): + key = ("ROOT", k_idx) + ref = bundle.ref_vars[key] + self.assertIn( + ref, groups[key], + msg=f"bundle ref Var missing from group {key}", + ) + # At least one entry in the group must be a per-sub-scenario + # Var that is NOT the ref. Without this the helper has + # regressed to representative-only behavior. + non_ref = [v for v in groups[key] if v is not ref] + self.assertTrue( + non_ref, + msg=f"group {key} contains only the ref Var; per-sub-" + f"scenario nonants from other sub-scenarios are missing", + ) + + +@unittest.skipUnless( + _ipopt_available(), + "Ipopt not available; skipping nonant_sensitivies end-to-end test", +) +class TestNonantSensitiviesBundleEndToEnd(unittest.TestCase): + """Smoke tests that go through the actual KKT sensitivity pipeline. + + Skipped automatically when Ipopt is not installed. The unit-test class + above is what guards correctness in solver-less CI; this class runs in + environments that have Ipopt and pins that the bundle path completes + without raising and produces a sensitivity per bundle nonant position. + """ + + def _build(self, scen_names, c_by_name, prob_by_name): + def _scen_creator(sname, **kwargs): + return _build_linear_scen( + sname, c_by_name[sname], prob_by_name[sname] + ) + return _build_bundle(_scen_creator, scen_names) + + def test_bundle_returns_one_sensitivity_per_bundle_position(self): + scen_names = ["A", "B", "C"] + c_by_name = { + "A": [1.0, 2.0], + "B": [3.0, 4.0], + "C": [5.0, 6.0], + } + prob_by_name = dict.fromkeys(scen_names, 1.0 / 3) + bundle = self._build(scen_names, c_by_name, prob_by_name) + sensis = nonant_sensitivies(bundle) + # Keys match the bundle's nonant positions; values are scalars. + self.assertEqual(set(sensis.keys()), {("ROOT", 0), ("ROOT", 1)}) + for k, v in sensis.items(): + # Either NumPy scalar or Python float — just check it is finite. + self.assertTrue( + float(v) == float(v), + msg=f"sensitivity at {k} is NaN ({v=})", + ) + + def test_single_scenario_bundle_matches_unbundled(self): + # A bundle of one sub-scenario must produce the same sensitivities + # as the standalone scenario — proves the bundle branch degenerates + # correctly and isn't introducing extra coupling. + c = [1.0, 2.0] + + def _scen_creator(sname, **kwargs): + return _build_linear_scen(sname, c, 1.0) + + # Standalone scenario. + scen = _scen_creator("standalone") + _attach_metadata(scen) + scen_sensis = nonant_sensitivies(scen) + + # Bundle of one. + bundle = _build_bundle(_scen_creator, ["only"]) + bundle_sensis = nonant_sensitivies(bundle) + + self.assertEqual(set(scen_sensis.keys()), set(bundle_sensis.keys())) + for k in scen_sensis: + self.assertAlmostEqual( + float(scen_sensis[k]), float(bundle_sensis[k]), places=4, + msg=f"single-scen bundle sensitivity at {k} diverges from " + f"standalone: bundle={bundle_sensis[k]} vs " + f"standalone={scen_sensis[k]}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/mpisppy/utils/nonant_sensitivities.py b/mpisppy/utils/nonant_sensitivities.py index 2a06cc6c3..72ec36f86 100644 --- a/mpisppy/utils/nonant_sensitivities.py +++ b/mpisppy/utils/nonant_sensitivities.py @@ -14,6 +14,45 @@ from mpisppy.utils.kkt.interface import InteriorPointInterface + +def _bundle_consensus_groups(bundle): + """Map each bundle nonant position to the per-sub-scenario nonant Vars + coupled to it by NA equality constraints. + + For a proper bundle (a Pyomo EF model with ``_ef_scenario_names``), + ``sputils.create_EF`` picks the first sub-scenario's nonant Var as the + bundle's ref Var and ties the rest by ``x_sub_j == ref`` constraints. + The bundle objective references each sub-scenario's *own* nonant Vars + rather than the ref. Probing the KKT sensitivity at the ref Var alone + therefore captures only the representative sub-scenario; correct + behavior is to perturb the consensus direction — every per-sub-scenario + nonant Var at that bundle position together. See issue #673. + + Returns ``{(ndn, k): [list of per-sub-scenario Var objects]}``. The + construction mirrors the bundle branch of ``sputils.nonant_cost_coeffs`` + and the gradient aggregation in ``mpisppy.extensions.grad_rho``. + """ + per_scen_to_bundle = {} + counters = {} + for (ndn, per_scen_i) in bundle.ref_vars.keys(): + if (ndn, per_scen_i) in bundle.ref_surrogate_vars: + continue + counters.setdefault(ndn, 0) + per_scen_to_bundle[(ndn, per_scen_i)] = (ndn, counters[ndn]) + counters[ndn] += 1 + + groups = {} + for scenario_name in bundle._ef_scenario_names: + scenario = bundle.component(scenario_name) + for node in scenario._mpisppy_node_list: + ndn = node.name + for per_scen_i, v in enumerate(node.nonant_vardata_list): + bundle_key = per_scen_to_bundle.get((ndn, per_scen_i)) + if bundle_key is not None: + groups.setdefault(bundle_key, []).append(v) + return groups + + def nonant_sensitivies(s): """ Compute the sensitivities of noants (w.r.t. the Lagrangian for s) Args: @@ -63,6 +102,12 @@ def nonant_sensitivies(s): grad_vec_kkt_inv = kkt_lu._lu.solve(grad_vec, "T") + consensus_groups = ( + _bundle_consensus_groups(s) + if hasattr(s, "_ef_scenario_names") + else None + ) + nonant_sensis = dict() for ndn_i, v in s._mpisppy_data.nonant_indices.items(): if v.fixed: @@ -70,10 +115,24 @@ def nonant_sensitivies(s): # +infy probably makes more conceptual sense, but 0 seems safer. nonant_sensis[ndn_i] = 0.0 continue - var_idx = kkt_builder._nlp._vardata_to_idx[v] y_vec = np.zeros(kkt.shape[0]) - y_vec[var_idx] = 1.0 + if consensus_groups is None: + y_vec[kkt_builder._nlp._vardata_to_idx[v]] = 1.0 + else: + # Bundle: probe the consensus direction so the perturbation + # matches the NA equality constraints that couple per-sub-scenario + # nonants at this bundle position. Without this, the formula + # would only see the representative ref Var (issue #673). + wrote_any = False + for sub_v in consensus_groups.get(ndn_i, ()): + if sub_v.fixed: + continue + y_vec[kkt_builder._nlp._vardata_to_idx[sub_v]] = 1.0 + wrote_any = True + if not wrote_any: + nonant_sensis[ndn_i] = 0.0 + continue x_denom = y_vec.T @ kkt_lu._lu.solve(y_vec) x = (-1 / x_denom) diff --git a/run_coverage.bash b/run_coverage.bash index e72ad2d70..64fe0d473 100755 --- a/run_coverage.bash +++ b/run_coverage.bash @@ -103,6 +103,9 @@ run_phase "test_jensens (serial)" \ run_phase "test_grad_rho_bundles (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_grad_rho_bundles.py -v +run_phase "test_sensi_rho_bundles (serial)" \ + coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_sensi_rho_bundles.py -v + run_phase "test_xhat_from_file (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_xhat_from_file.py -v From 1f4ed6d0ac8f4a926abeb2703c4aeb7f64581e24 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 07:39:21 -0700 Subject: [PATCH 3/9] reduced_costs_rho: aggregate per-sub-scenario rc in proper bundles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-fix ReducedCostsSpoke.extract_and_store_reduced_costs read only sub.rc[ref_var] when scoring each bundle nonant position. The bundle ref Var is the first sub-scenario's nonant (sputils.create_EF), and the bundle objective references each sub-scenario's *own* nonant Vars (the rest tied by NA equality constraints), so reading only the ref's rc captured one representative sub-scenario's contribution plus the NA constraint dual. The persistent-solver vars_to_load list also restricted itself to the bundle ref Vars, leaving per-sub-scenario rc entries stale. (Issue #673.) Add _consensus_rc_sum: for a bundle, sum scenario.rc over every per-sub-scenario nonant in the consensus group at a given bundle position. Cancels the NA multipliers and yields the rate of bundle- objective change per unit consensus shift — the right input to rc-rho / rc-fixing. Falls back to scenario.rc[ref_var] for the unbundled path so existing behavior is preserved bit-for-bit. Reuse _bundle_consensus_groups from nonant_sensitivities (same helper the sensi_rho fix introduced; it parallels the bundle branch in sputils.nonant_cost_coeffs and grad_rho). Extend vars_to_load in the persistent-solver branch to cover every per-sub-scenario nonant Var. Tests in test_reduced_costs_rho_bundles.py (solver-free): seed the bundle's rc Suffix with distinct values per (sub, position), assert _consensus_rc_sum returns the sum across the consensus group, and add an explicit "not equal to ref-only rc" guard. The unbundled path is also covered. Confirmed to fail under a mocked representative-only helper. Wired into run_coverage.bash and the unit-tests block of test_pr_and_main.yml. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test_pr_and_main.yml | 3 +- mpisppy/cylinders/reduced_costs_spoke.py | 58 +++++++- .../tests/test_reduced_costs_rho_bundles.py | 139 ++++++++++++++++++ run_coverage.bash | 3 + 4 files changed, 199 insertions(+), 4 deletions(-) create mode 100644 mpisppy/tests/test_reduced_costs_rho_bundles.py diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml index 0aebd0642..f27bb907b 100644 --- a/.github/workflows/test_pr_and_main.yml +++ b/.github/workflows/test_pr_and_main.yml @@ -1015,7 +1015,8 @@ jobs: mpisppy/tests/test_jensens.py \ mpisppy/tests/test_proper_bundler.py \ mpisppy/tests/test_grad_rho_bundles.py \ - mpisppy/tests/test_sensi_rho_bundles.py + mpisppy/tests/test_sensi_rho_bundles.py \ + mpisppy/tests/test_reduced_costs_rho_bundles.py - name: Upload coverage data if: always() diff --git a/mpisppy/cylinders/reduced_costs_spoke.py b/mpisppy/cylinders/reduced_costs_spoke.py index d15ee1349..842492386 100644 --- a/mpisppy/cylinders/reduced_costs_spoke.py +++ b/mpisppy/cylinders/reduced_costs_spoke.py @@ -12,8 +12,33 @@ from mpisppy.cylinders.lagrangian_bounder import LagrangianOuterBound from mpisppy.cylinders.spwindow import Field from mpisppy.utils.sputils import is_persistent +from mpisppy.utils.nonant_sensitivities import _bundle_consensus_groups from mpisppy import MPI, global_toc + +def _consensus_rc_sum(scenario, ndn_i, ref_var, consensus_groups): + """Aggregated reduced cost at a nonant position. + + For an unbundled scenario this is just ``scenario.rc[ref_var]`` — the + single nonant Var's reduced cost. + + For a proper bundle the bundle objective references each sub-scenario's + *own* nonant Vars (the bundle ref is the first sub-scenario's Var; the + rest are tied by NA equality constraints). Probing only the ref Var's + reduced cost captures one representative sub-scenario's contribution + plus the NA constraint dual; summing the reduced costs of every + per-sub-scenario nonant in the consensus group cancels the NA + multipliers and gives the rate of bundle-objective change per unit + consensus shift — i.e. the correct "expected reduced cost" input to + rc-rho/rc-fixing. See issue #673. + """ + if consensus_groups is None: + return scenario.rc[ref_var] + return sum( + scenario.rc[v] + for v in consensus_groups.get(ndn_i, (ref_var,)) + ) + class ReducedCostsSpoke(LagrangianOuterBound): send_fields = (*LagrangianOuterBound.send_fields, Field.EXPECTED_REDUCED_COST, Field.SCENARIO_REDUCED_COST, @@ -161,6 +186,11 @@ def extract_and_store_reduced_costs(self): rc = np.zeros(self.nonant_length) for sub in self.opt.local_scenarios.values(): + consensus_groups = ( + _bundle_consensus_groups(sub) + if hasattr(sub, "_ef_scenario_names") + else None + ) if is_persistent(sub._solver_plugin): # Note: what happens with non-persistent solvers? # - if rc is accepted as a model suffix by the solver (e.g. gurobi shell), it is loaded in postsolve @@ -168,7 +198,20 @@ def extract_and_store_reduced_costs(self): # - direct solvers seem to behave the same as persistent solvers # GurobiDirect needs vars_to_load argument # XpressDirect loads for all vars by default - TODO: should notify someone of this inconsistency - vars_to_load = [x for _, x in sub._mpisppy_data.nonant_indices.items()] + if consensus_groups is None: + vars_to_load = [ + x for _, x in sub._mpisppy_data.nonant_indices.items() + ] + else: + # Bundle: load rc for every per-sub-scenario nonant Var, + # not only the ref. The consensus rc sums over the whole + # group; loading just the ref would leave most entries + # at their stale pre-solve values. + vars_to_load = [] + for ndn_i, ref in sub._mpisppy_data.nonant_indices.items(): + vars_to_load.extend( + consensus_groups.get(ndn_i, (ref,)) + ) sub._solver_plugin.load_rc(vars_to_load=vars_to_load) for ci, (ndn_i, xvar) in enumerate(sub._mpisppy_data.nonant_indices.items()): @@ -187,7 +230,9 @@ def extract_and_store_reduced_costs(self): # solver takes care of sign of rc, based on lb, ub and max,min # rc can be of wrong sign if numerically 0 - accepted here, checked in extension if (xvar.lb is not None and xb - xvar.lb <= self.bound_tol) or (xvar.ub is not None and xvar.ub - xb <= self.bound_tol): - rc[ci] += sub._mpisppy_probability * sub.rc[xvar] + rc[ci] += sub._mpisppy_probability * _consensus_rc_sum( + sub, ndn_i, xvar, consensus_groups, + ) # not close to either bound -> rc = nan else: rc[ci] = np.nan @@ -196,12 +241,19 @@ def extract_and_store_reduced_costs(self): assert self._scenario_rc_buffer.size == self.send_buffers[Field.SCENARIO_REDUCED_COST].data_len() ci = 0 # buffer index for sub in self.opt.local_scenarios.values(): + consensus_groups = ( + _bundle_consensus_groups(sub) + if hasattr(sub, "_ef_scenario_names") + else None + ) for ndn_i, xvar in sub._mpisppy_data.nonant_indices.items(): # fixed by modeler if ndn_i in self._modeler_fixed_nonants[sub]: self._scenario_rc_buffer[ci] = np.nan else: - self._scenario_rc_buffer[ci] = sub.rc[xvar] + self._scenario_rc_buffer[ci] = _consensus_rc_sum( + sub, ndn_i, xvar, consensus_groups, + ) ci += 1 self.rc_scenario = self._scenario_rc_buffer # print(f"In ReducedCostsSpoke; {self.rc_scenario=}") diff --git a/mpisppy/tests/test_reduced_costs_rho_bundles.py b/mpisppy/tests/test_reduced_costs_rho_bundles.py new file mode 100644 index 000000000..a1e76c872 --- /dev/null +++ b/mpisppy/tests/test_reduced_costs_rho_bundles.py @@ -0,0 +1,139 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""ReducedCostsSpoke must aggregate per-sub-scenario reduced costs across +the consensus group at each bundle nonant position (issue #673). Pre-fix +behavior used ``sub.rc[ref_var]`` which only saw the bundle's representative +ref Var; in an EF bundle the bundle objective references each sub-scenario's +own nonant Vars (the rest are tied by NA equality constraints), so the rc +of the ref Var alone is not the correct rate-of-change input for rc-fixing. + +Tests are solver-free: they synthesize a bundle, prime the bundle's ``rc`` +Suffix with known values, and call the spoke's aggregation helper directly. +""" + +import unittest + +import pyomo.environ as pyo + +import mpisppy.utils.sputils as sputils +from mpisppy.cylinders.reduced_costs_spoke import _consensus_rc_sum +from mpisppy.utils.nonant_sensitivities import _bundle_consensus_groups + + +def _build_linear_scen(name, c_coefs, prob): + m = pyo.ConcreteModel(name=name) + m.x = pyo.Var(range(len(c_coefs)), bounds=(0, 10), initialize=0.0) + m.obj = pyo.Objective( + expr=sum(c_coefs[i] * m.x[i] for i in range(len(c_coefs))) + ) + sputils.attach_root_node(m, 0, [m.x[i] for i in range(len(c_coefs))]) + m._mpisppy_probability = prob + return m + + +def _attach_metadata(scen): + if not hasattr(scen, "_mpisppy_data"): + scen._mpisppy_data = pyo.Block( + name="For non-Pyomo mpi-sppy data" + ) + scen._mpisppy_data.nlens = { + node.name: len(node.nonant_vardata_list) + for node in scen._mpisppy_node_list + } + nonant_indices = {} + for node in scen._mpisppy_node_list: + for i, v in enumerate(node.nonant_vardata_list): + nonant_indices[(node.name, i)] = v + scen._mpisppy_data.nonant_indices = nonant_indices + + +def _build_bundle(scen_creator, scen_names): + bundle = sputils.create_EF(scen_names, scen_creator, EF_name="bundle") + nonantlist = [ + v for idx, v in bundle.ref_vars.items() + if idx[0] == "ROOT" and idx not in bundle.ref_surrogate_vars + ] + surrogates = [ + v for idx, v in bundle.ref_surrogate_vars.items() + if idx[0] == "ROOT" + ] + sputils.attach_root_node(bundle, 0, nonantlist, None, surrogates) + _attach_metadata(bundle) + return bundle + + +class TestConsensusRcSum(unittest.TestCase): + """Unit tests for the spoke's _consensus_rc_sum helper.""" + + def test_unbundled_returns_single_var_rc(self): + # When consensus_groups is None (non-bundle path), the helper must + # return the ref Var's rc unchanged so existing non-bundle behavior + # is preserved bit-for-bit. + scen = _build_linear_scen("solo", [1.0, 2.0], 1.0) + _attach_metadata(scen) + scen.rc = pyo.Suffix(direction=pyo.Suffix.IMPORT) + scen.rc[scen.x[0]] = 7.5 + scen.rc[scen.x[1]] = 11.25 + for ndn_i, ref in scen._mpisppy_data.nonant_indices.items(): + self.assertEqual( + _consensus_rc_sum(scen, ndn_i, ref, None), + scen.rc[ref], + ) + + def test_bundle_returns_consensus_sum(self): + # 3 sub-scenarios x 2 first-stage vars. Set bundle.rc[per_scen_var] + # to distinct values per (sub, position) and verify the helper + # returns the sum across the consensus group. + scen_names = ["A", "B", "C"] + c_by_name = {"A": [1.0, 2.0], "B": [3.0, 4.0], "C": [5.0, 6.0]} + prob_by_name = dict.fromkeys(scen_names, 1.0 / 3) + + def _scen_creator(sname, **kwargs): + return _build_linear_scen( + sname, c_by_name[sname], prob_by_name[sname] + ) + + bundle = _build_bundle(_scen_creator, scen_names) + bundle.rc = pyo.Suffix(direction=pyo.Suffix.IMPORT) + + # Encode rc[sub.x_k] = 100 * (sub_index + 1) + k so each sub-Var + # carries a distinct, easy-to-reproduce rc value. + rc_values = {} + for sub_idx, sname in enumerate(scen_names): + scen = bundle.component(sname) + for k in range(2): + v = scen.x[k] + bundle.rc[v] = 100.0 * (sub_idx + 1) + k + rc_values[(sname, k)] = bundle.rc[v] + + consensus_groups = _bundle_consensus_groups(bundle) + for ndn_i, ref in bundle._mpisppy_data.nonant_indices.items(): + # Expected: sum across all 3 sub-scenarios at this bundle position. + expected = sum( + rc_values[(sname, ndn_i[1])] for sname in scen_names + ) + got = _consensus_rc_sum(bundle, ndn_i, ref, consensus_groups) + self.assertAlmostEqual( + got, expected, places=10, + msg=f"bundle {ndn_i=}: {got=} != consensus-sum {expected=}", + ) + # Also confirm the helper does NOT silently return only the ref's + # rc — the pre-fix behavior. (This guards against a regression + # to representative-only sourcing.) + ref_only = bundle.rc[ref] + if ref_only != expected: + self.assertNotAlmostEqual( + got, ref_only, places=6, + msg=f"bundle {ndn_i=} matches ref-only rc {ref_only}; " + f"aggregation regressed", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/run_coverage.bash b/run_coverage.bash index 64fe0d473..9897e51a4 100755 --- a/run_coverage.bash +++ b/run_coverage.bash @@ -106,6 +106,9 @@ run_phase "test_grad_rho_bundles (serial)" \ run_phase "test_sensi_rho_bundles (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_sensi_rho_bundles.py -v +run_phase "test_reduced_costs_rho_bundles (serial)" \ + coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_reduced_costs_rho_bundles.py -v + run_phase "test_xhat_from_file (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_xhat_from_file.py -v From 22128281f54c5f9f851a7d6ed114c29191ecfd5a Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 07:49:38 -0700 Subject: [PATCH 4/9] reduced_costs_rho: assert vars_to_load covers the consensus group MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A runtime guard in extract_and_store_reduced_costs that pre-checks every sub-scenario nonant Var in every consensus group has an entry on the bundle's rc Suffix before _consensus_rc_sum reads it. Catches a future regression where vars_to_load is restricted back to the bundle ref Vars (or any other path that skips per-sub-scenario nonants) — the consensus sum would otherwise silently aggregate stale or missing values. No-op for unbundled scenarios so the non-bundle path is unchanged. Tests in test_reduced_costs_rho_bundles.py: - unbundled path is a no-op even with an empty rc Suffix. - passes when every consensus Var has rc. - fails (AssertionError) when only ref Vars have rc — simulating the pre-fix vars_to_load behavior. - fails when one sub-scenario Var is missing — simulating an off-by-one vars_to_load regression. - error message names the missing Var and bundle position so the failure is actionable. Co-Authored-By: Claude Opus 4.7 (1M context) --- mpisppy/cylinders/reduced_costs_spoke.py | 21 +++++ .../tests/test_reduced_costs_rho_bundles.py | 79 ++++++++++++++++++- 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/mpisppy/cylinders/reduced_costs_spoke.py b/mpisppy/cylinders/reduced_costs_spoke.py index 842492386..7f698f811 100644 --- a/mpisppy/cylinders/reduced_costs_spoke.py +++ b/mpisppy/cylinders/reduced_costs_spoke.py @@ -16,6 +16,25 @@ from mpisppy import MPI, global_toc +def _assert_consensus_rc_loaded(scenario, consensus_groups): + """Pre-aggregation guard: every Var in every consensus group must have a + reduced cost on ``scenario.rc`` before _consensus_rc_sum reads it. + + This catches a future regression where ``vars_to_load`` is restricted + back to bundle ref Vars (or any other path) that leaves per-sub-scenario + nonants without an rc entry — the consensus sum would then silently + read stale or zero values. No-op for unbundled scenarios. + """ + if consensus_groups is None: + return + for ndn_i, group in consensus_groups.items(): + for v in group: + assert v in scenario.rc, ( + f"reduced cost not loaded for {v.name}; vars_to_load did " + f"not cover the consensus group at bundle position {ndn_i}" + ) + + def _consensus_rc_sum(scenario, ndn_i, ref_var, consensus_groups): """Aggregated reduced cost at a nonant position. @@ -214,6 +233,8 @@ def extract_and_store_reduced_costs(self): ) sub._solver_plugin.load_rc(vars_to_load=vars_to_load) + _assert_consensus_rc_loaded(sub, consensus_groups) + for ci, (ndn_i, xvar) in enumerate(sub._mpisppy_data.nonant_indices.items()): # fixed by modeler if ndn_i in self._modeler_fixed_nonants[sub]: diff --git a/mpisppy/tests/test_reduced_costs_rho_bundles.py b/mpisppy/tests/test_reduced_costs_rho_bundles.py index a1e76c872..06c579ccf 100644 --- a/mpisppy/tests/test_reduced_costs_rho_bundles.py +++ b/mpisppy/tests/test_reduced_costs_rho_bundles.py @@ -22,7 +22,10 @@ import pyomo.environ as pyo import mpisppy.utils.sputils as sputils -from mpisppy.cylinders.reduced_costs_spoke import _consensus_rc_sum +from mpisppy.cylinders.reduced_costs_spoke import ( + _assert_consensus_rc_loaded, + _consensus_rc_sum, +) from mpisppy.utils.nonant_sensitivities import _bundle_consensus_groups @@ -135,5 +138,79 @@ def _scen_creator(sname, **kwargs): ) +class TestAssertConsensusRcLoaded(unittest.TestCase): + """Structural guard that the rc Suffix carries every Var the consensus + sum is about to read. The runtime assertion in ``ReducedCostsSpoke`` + catches a future regression where ``vars_to_load`` is restricted back + to the bundle ref Vars (or any other path that misses sub-scenario + nonants). + """ + + def setUp(self): + scen_names = ["A", "B", "C"] + c_by_name = {"A": [1.0, 2.0], "B": [3.0, 4.0], "C": [5.0, 6.0]} + prob_by_name = dict.fromkeys(scen_names, 1.0 / 3) + + def _scen_creator(sname, **kwargs): + return _build_linear_scen( + sname, c_by_name[sname], prob_by_name[sname] + ) + + self.scen_names = scen_names + self.bundle = _build_bundle(_scen_creator, scen_names) + self.bundle.rc = pyo.Suffix(direction=pyo.Suffix.IMPORT) + self.consensus_groups = _bundle_consensus_groups(self.bundle) + + def _populate_rc(self, vars_iterable): + for v in vars_iterable: + self.bundle.rc[v] = 0.0 + + def test_unbundled_path_is_noop(self): + # When consensus_groups is None the helper must not raise even with + # an empty rc Suffix — non-bundle code paths are untouched. + scen = _build_linear_scen("solo", [1.0, 2.0], 1.0) + _attach_metadata(scen) + scen.rc = pyo.Suffix(direction=pyo.Suffix.IMPORT) + # Don't populate rc at all — this is the "unbundled, do nothing" + # contract; should return without raising. + _assert_consensus_rc_loaded(scen, None) + + def test_passes_when_every_consensus_var_has_rc(self): + # Populate rc for every sub-scenario nonant Var: assertion must + # succeed silently. + all_consensus_vars = [ + v for group in self.consensus_groups.values() for v in group + ] + self._populate_rc(all_consensus_vars) + _assert_consensus_rc_loaded(self.bundle, self.consensus_groups) + + def test_fails_when_only_ref_vars_have_rc(self): + # Simulate the pre-fix vars_to_load behavior: load only the bundle + # ref Vars. Non-ref sub-scenario nonants have no rc entry, so the + # consensus sum would silently read missing values. The assertion + # must catch this. + ref_vars = list(self.bundle._mpisppy_data.nonant_indices.values()) + self._populate_rc(ref_vars) + with self.assertRaises(AssertionError) as ctx: + _assert_consensus_rc_loaded(self.bundle, self.consensus_groups) + # Error message should name a Var and the bundle position so the + # failure is actionable. + msg = str(ctx.exception) + self.assertIn("vars_to_load did not cover", msg) + self.assertIn("ROOT", msg) + + def test_fails_when_one_sub_scenario_var_missing(self): + # Drop a single sub-scenario's nonant Var to simulate a partial + # vars_to_load regression (e.g. an off-by-one or a filtered list). + all_vars = [ + v for group in self.consensus_groups.values() for v in group + ] + # Skip the first sub-scenario's first-position Var. + first_skip = self.bundle.component(self.scen_names[0]).x[0] + self._populate_rc([v for v in all_vars if v is not first_skip]) + with self.assertRaises(AssertionError): + _assert_consensus_rc_loaded(self.bundle, self.consensus_groups) + + if __name__ == "__main__": unittest.main() From 73cc73341e4d89a34068fa6b1d6eee3ef647eaca Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 07:49:50 -0700 Subject: [PATCH 5/9] deprecate SepRho and ReducedCostsRho per #673 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per the issue thread (bknueven), SepRho is expected to be subsumed into GradRho and reduced-cost rho has not shown to be effective in practice. Emit a DeprecationWarning the first time either rho setter is instantiated so downstream projects have lead time to migrate. The warnings fire as the first statement of each __init__, before any of the option/PH-machinery reads — so they surface even on configurations that would later fail to construct. Tests in test_rho_deprecations.py: instantiate each class with a minimal stub PH that supplies only the option key the warning's guard reads; capture warnings via warnings.catch_warnings; assert a DeprecationWarning was issued with the class name and the issue reference. Wired into run_coverage.bash and the unit tests block of test_pr_and_main.yml alongside the bundle-fix tests. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/test_pr_and_main.yml | 3 +- mpisppy/extensions/reduced_costs_rho.py | 9 +++ mpisppy/extensions/sep_rho.py | 9 +++ mpisppy/tests/test_rho_deprecations.py | 81 +++++++++++++++++++++++++ run_coverage.bash | 3 + 5 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 mpisppy/tests/test_rho_deprecations.py diff --git a/.github/workflows/test_pr_and_main.yml b/.github/workflows/test_pr_and_main.yml index f27bb907b..5110ecd5f 100644 --- a/.github/workflows/test_pr_and_main.yml +++ b/.github/workflows/test_pr_and_main.yml @@ -1016,7 +1016,8 @@ jobs: mpisppy/tests/test_proper_bundler.py \ mpisppy/tests/test_grad_rho_bundles.py \ mpisppy/tests/test_sensi_rho_bundles.py \ - mpisppy/tests/test_reduced_costs_rho_bundles.py + mpisppy/tests/test_reduced_costs_rho_bundles.py \ + mpisppy/tests/test_rho_deprecations.py - name: Upload coverage data if: always() diff --git a/mpisppy/extensions/reduced_costs_rho.py b/mpisppy/extensions/reduced_costs_rho.py index 3ce67f945..368f530f4 100644 --- a/mpisppy/extensions/reduced_costs_rho.py +++ b/mpisppy/extensions/reduced_costs_rho.py @@ -7,6 +7,8 @@ # full copyright and license information. ############################################################################### +import warnings + import numpy as np from mpisppy import global_toc from mpisppy.extensions.sensi_rho import _SensiRhoBase @@ -20,6 +22,13 @@ class ReducedCostsRho(_SensiRhoBase): """ def __init__(self, ph): + warnings.warn( + "ReducedCostsRho is scheduled for deprecation; reduced-cost " + "rho has not been demonstrated to be effective in practice. " + "See https://github.com/Pyomo/mpi-sppy/issues/673.", + DeprecationWarning, + stacklevel=2, + ) cfg = ph.options["reduced_costs_rho_options"]["cfg"] super().__init__(ph, cfg) self.ph = ph diff --git a/mpisppy/extensions/sep_rho.py b/mpisppy/extensions/sep_rho.py index b4685cf4d..bf2b47836 100644 --- a/mpisppy/extensions/sep_rho.py +++ b/mpisppy/extensions/sep_rho.py @@ -7,6 +7,8 @@ # full copyright and license information. ############################################################################### +import warnings + import mpisppy.extensions.dyn_rho_base import numpy as np import mpisppy.MPI as MPI @@ -24,6 +26,13 @@ class SepRho(mpisppy.extensions.dyn_rho_base.Dyn_Rho_extension_base): """ def __init__(self, ph): + warnings.warn( + "SepRho is scheduled for deprecation; its functionality is " + "expected to be subsumed into GradRho. See " + "https://github.com/Pyomo/mpi-sppy/issues/673.", + DeprecationWarning, + stacklevel=2, + ) cfg = ph.options["sep_rho_options"]["cfg"] super().__init__(ph, cfg) self.ph = ph diff --git a/mpisppy/tests/test_rho_deprecations.py b/mpisppy/tests/test_rho_deprecations.py new file mode 100644 index 000000000..602f34506 --- /dev/null +++ b/mpisppy/tests/test_rho_deprecations.py @@ -0,0 +1,81 @@ +############################################################################### +# mpi-sppy: MPI-based Stochastic Programming in PYthon +# +# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for +# Sustainable Energy, LLC, The Regents of the University of California, et al. +# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for +# full copyright and license information. +############################################################################### +"""SepRho and ReducedCostsRho are scheduled for deprecation per +https://github.com/Pyomo/mpi-sppy/issues/673. Both must emit a +DeprecationWarning the first time a user instantiates them so downstream +projects have lead time to migrate (likely to GradRho). + +These tests assert only that the warning is issued; full instantiation +needs PH machinery the rho setters expect, but the deprecation warning +is the first thing each ``__init__`` does, so it fires regardless of +whether the rest of construction succeeds. +""" + +import unittest +import warnings + +from mpisppy.extensions.sep_rho import SepRho +from mpisppy.extensions.reduced_costs_rho import ReducedCostsRho + + +def _capture_first_warning(callable_, *, category): + """Invoke ``callable_`` and return the first warning of ``category`` + captured during the call, or None. Swallows any exception raised after + the warning fires; the deprecation warning is the first statement in + each ``__init__`` so it is independent of downstream construction.""" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", category) + try: + callable_() + except Exception: + pass + matches = [w for w in caught if issubclass(w.category, category)] + return matches[0] if matches else None + + +class _StubPH: + """Just enough surface for the rho setters' first read of options. + + Construction will fail beyond the deprecation warning (Extension's + base class needs more attributes) and that's fine — we only assert on + the warning, which fires before any of those reads.""" + + def __init__(self, key): + self.options = {key: {"cfg": object()}} + + +class TestRhoDeprecations(unittest.TestCase): + + def test_sep_rho_emits_deprecation_warning(self): + warning = _capture_first_warning( + lambda: SepRho(_StubPH("sep_rho_options")), + category=DeprecationWarning, + ) + self.assertIsNotNone( + warning, + msg="SepRho.__init__ did not emit a DeprecationWarning", + ) + self.assertIn("SepRho", str(warning.message)) + self.assertIn("673", str(warning.message)) + + def test_reduced_costs_rho_emits_deprecation_warning(self): + warning = _capture_first_warning( + lambda: ReducedCostsRho(_StubPH("reduced_costs_rho_options")), + category=DeprecationWarning, + ) + self.assertIsNotNone( + warning, + msg="ReducedCostsRho.__init__ did not emit a DeprecationWarning", + ) + self.assertIn("ReducedCostsRho", str(warning.message)) + self.assertIn("673", str(warning.message)) + + +if __name__ == "__main__": + unittest.main() diff --git a/run_coverage.bash b/run_coverage.bash index 9897e51a4..d25170139 100755 --- a/run_coverage.bash +++ b/run_coverage.bash @@ -109,6 +109,9 @@ run_phase "test_sensi_rho_bundles (serial)" \ run_phase "test_reduced_costs_rho_bundles (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_reduced_costs_rho_bundles.py -v +run_phase "test_rho_deprecations (serial)" \ + coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_rho_deprecations.py -v + run_phase "test_xhat_from_file (serial)" \ coverage run --rcfile=.coveragerc -m pytest mpisppy/tests/test_xhat_from_file.py -v From 03dee7a745f40d682cf0f0eca9fcd236d8d187a4 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 08:58:05 -0700 Subject: [PATCH 6/9] nonant_sensitivities: defer scipy/pynumero imports into the function MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI's "unit tests (no solver required)" job runs without scipy installed. The previous top-level imports of pynumero (which transitively imports scipy) and the KKT interface caused scipy-free environments to fail at collection for every test file that imports anything from this module (including just _bundle_consensus_groups), turning the unit tests job red and zeroing out codecov/patch on PR #679. Move the scipy-dependent imports inside nonant_sensitivies(). The function still requires scipy when called (it's the KKT solve path) but importing the module — and importing scipy-free helpers like _bundle_consensus_groups from other modules — no longer does. Verified by simulating scipy-less __import__ and confirming mpisppy.utils.nonant_sensitivities, mpisppy.cylinders.reduced_costs_spoke, mpisppy.extensions.sensi_rho, and mpisppy.extensions.reduced_costs_rho all import successfully. Co-Authored-By: Claude Opus 4.7 (1M context) --- mpisppy/utils/nonant_sensitivities.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/mpisppy/utils/nonant_sensitivities.py b/mpisppy/utils/nonant_sensitivities.py index 72ec36f86..4a085b6b7 100644 --- a/mpisppy/utils/nonant_sensitivities.py +++ b/mpisppy/utils/nonant_sensitivities.py @@ -10,9 +10,12 @@ import numpy as np import pyomo.environ as pyo -from pyomo.contrib.pynumero.linalg.scipy_interface import ScipyLU -from mpisppy.utils.kkt.interface import InteriorPointInterface +# pynumero / KKT helpers pull in scipy at import time; defer them into +# nonant_sensitivies() so importing this module — and re-using its +# scipy-free helpers like _bundle_consensus_groups from other modules — +# does not require scipy. The "unit tests (no solver required)" CI job +# runs in an environment without scipy installed. def _bundle_consensus_groups(bundle): @@ -61,6 +64,13 @@ def nonant_sensitivies(s): nonant_sensis (dict): [ndn_i]: sensitivity for the Var """ + # Lazy import: scipy is required to actually compute sensitivities + # but not to import this module. Keeping these inside the function + # lets scipy-free environments import _bundle_consensus_groups and + # the rest of the module without failure. + from pyomo.contrib.pynumero.linalg.scipy_interface import ScipyLU + from mpisppy.utils.kkt.interface import InteriorPointInterface + # first, solve the subproblems with Ipopt, # and gather sensitivity information ipopt = pyo.SolverFactory("ipopt") From afbc3c43037f08d9004f2d5450cce42877cd4e9a Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 09:35:23 -0700 Subject: [PATCH 7/9] docs(_consensus_rc_sum): explain why no explicit prob weighting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The consensus sum doesn't multiply by per-sub-scenario probability and that is intentional: sputils.create_EF builds the bundle objective as sum_s cond_prob_s * obj_s, so each per-sub-scenario nonant's reduced cost already carries cond_prob_s. The caller multiplies the consensus sum by EF_prob, recovering sum_s prob_s * unbundled_rc_s,k — exactly the unbundled formula. Adding explicit prob weighting would double-count. Spell that derivation out in the docstring so future readers don't re-litigate it. Co-Authored-By: Claude Opus 4.7 (1M context) --- mpisppy/cylinders/reduced_costs_spoke.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/mpisppy/cylinders/reduced_costs_spoke.py b/mpisppy/cylinders/reduced_costs_spoke.py index 7f698f811..ac7b15fe2 100644 --- a/mpisppy/cylinders/reduced_costs_spoke.py +++ b/mpisppy/cylinders/reduced_costs_spoke.py @@ -48,8 +48,27 @@ def _consensus_rc_sum(scenario, ndn_i, ref_var, consensus_groups): plus the NA constraint dual; summing the reduced costs of every per-sub-scenario nonant in the consensus group cancels the NA multipliers and gives the rate of bundle-objective change per unit - consensus shift — i.e. the correct "expected reduced cost" input to + consensus shift — the correct "expected reduced cost" input to rc-rho/rc-fixing. See issue #673. + + Note on probability weighting (no explicit weighting needed): + + ``sputils.create_EF`` builds the bundle objective as + ``(sum_s prob_s * obj_s) / EF_prob = sum_s cond_prob_s * obj_s``, + where ``cond_prob_s = prob_s / EF_prob``. Pyomo's reduced cost for + each per-sub-scenario nonant is therefore + ``bundle.rc[x_s,k] = cond_prob_s * unbundled_rc_s,k`` (the per- + sub-scenario constraint duals scale with ``cond_prob_s`` too because + the objective scaling carries through to the dual LP), modulo the NA + multipliers that cancel in the consensus sum. So the consensus sum + already equals + ``sum_s cond_prob_s * unbundled_rc_s,k = (1/EF_prob) * sum_s prob_s * unbundled_rc_s,k`` + — the ``cond_prob_s`` is baked in by ``create_EF``. + + The caller multiplies the result by ``sub._mpisppy_probability`` + (= ``EF_prob`` for a bundle), recovering ``sum_s prob_s * unbundled_rc_s,k`` + — exactly the unbundled formula. Adding an explicit ``prob_s`` + weighting inside this sum would double-count. """ if consensus_groups is None: return scenario.rc[ref_var] From c4297554b1cf000ed429fd7eba905de68c1bb251 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Thu, 7 May 2026 09:51:15 -0700 Subject: [PATCH 8/9] docs(_eval_grad_exprs): sentinel comment on the current-values branch Note that the implicit "fall through to evaluate at current Var values" path is correct for bundles only because post-solve NA equality leaves every per-sub-scenario nonant at a given bundle position equal to the consensus value. The cached gradient expression is a sum of partials over those Vars; if a future caller invokes _eval_grad_exprs while the bundle is in a non-NA-feasible intermediate state, this branch would silently return a non-consensus gradient. Spell that out so the invariant is obvious to a future reader. Co-Authored-By: Claude Opus 4.7 (1M context) --- mpisppy/extensions/grad_rho.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mpisppy/extensions/grad_rho.py b/mpisppy/extensions/grad_rho.py index ea3a010b2..d17471daf 100644 --- a/mpisppy/extensions/grad_rho.py +++ b/mpisppy/extensions/grad_rho.py @@ -219,6 +219,18 @@ def _eval_grad_exprs(self, s, xhat): var.value = xhat[ci] ci += 1 + # Current-values branch (the default path: cfg.eval_at_xhat is False, + # or eval_at_xhat is True but no best xhat has been received yet). + # Pyomo evaluates the cached gradient expressions against whatever + # values are currently on the Vars — the most recent post-solve + # state. For a bundle this is correct ONLY because the NA equality + # constraints leave every per-sub-scenario nonant at bundle + # position k equal to the bundle's consensus value at that + # position; the cached expression is a sum of partials over those + # Vars and therefore evaluates the consensus-direction gradient. + # If a future caller invokes _eval_grad_exprs while the bundle is + # in a non-NA-feasible intermediate state, this branch would + # silently return a non-consensus gradient. for ndn_i, var in s._mpisppy_data.nonant_indices.items(): grads[ndn_i] = pyo.value(self.grad_exprs[s][ndn_i]) From dab1da8cf89466903d5acad7d701cf60df7db435 Mon Sep 17 00:00:00 2001 From: Dave Woodruff Date: Fri, 8 May 2026 13:45:38 -0700 Subject: [PATCH 9/9] address copilot review on PR #679 - _assert_consensus_rc_loaded: raise RuntimeError instead of plain `assert` so the guard survives `python -O`. - extract_and_store_reduced_costs: cache _bundle_consensus_groups(sub) per local scenario; the second pass populating _scenario_rc_buffer now reads from the cache instead of recomputing it. - grad_rho._get_grad_exprs: drop the inline bundle position-mapping duplicate and call the shared _bundle_consensus_groups helper from mpisppy.utils.nonant_sensitivities. Same behavior; one source of truth for the (ndn, k) -> per-sub-scenario Vars mapping. - test_reduced_costs_rho_bundles: assertRaises(AssertionError) -> assertRaises(RuntimeError) for the two guard tests. Co-Authored-By: Claude Opus 4.7 (1M context) --- mpisppy/cylinders/reduced_costs_spoke.py | 30 ++++++++++++------- mpisppy/extensions/grad_rho.py | 30 +++++++------------ .../tests/test_reduced_costs_rho_bundles.py | 6 ++-- 3 files changed, 33 insertions(+), 33 deletions(-) diff --git a/mpisppy/cylinders/reduced_costs_spoke.py b/mpisppy/cylinders/reduced_costs_spoke.py index ac7b15fe2..ff5936aa4 100644 --- a/mpisppy/cylinders/reduced_costs_spoke.py +++ b/mpisppy/cylinders/reduced_costs_spoke.py @@ -29,10 +29,14 @@ def _assert_consensus_rc_loaded(scenario, consensus_groups): return for ndn_i, group in consensus_groups.items(): for v in group: - assert v in scenario.rc, ( - f"reduced cost not loaded for {v.name}; vars_to_load did " - f"not cover the consensus group at bundle position {ndn_i}" - ) + # Plain `assert` would be stripped under `python -O`, disabling + # the guard exactly where silent stale-rc reads would do real + # damage. Raise unconditionally instead. + if v not in scenario.rc: + raise RuntimeError( + f"reduced cost not loaded for {v.name}; vars_to_load did " + f"not cover the consensus group at bundle position {ndn_i}" + ) def _consensus_rc_sum(scenario, ndn_i, ref_var, consensus_groups): @@ -223,12 +227,20 @@ def extract_and_store_reduced_costs(self): # would probably need additional info about where scenarios disagree rc = np.zeros(self.nonant_length) - for sub in self.opt.local_scenarios.values(): - consensus_groups = ( + # _bundle_consensus_groups walks every sub-scenario/node/var in the + # bundle; cache once per `sub` and reuse below in the second pass that + # populates _scenario_rc_buffer. + consensus_groups_by_sub = { + sub: ( _bundle_consensus_groups(sub) if hasattr(sub, "_ef_scenario_names") else None ) + for sub in self.opt.local_scenarios.values() + } + + for sub in self.opt.local_scenarios.values(): + consensus_groups = consensus_groups_by_sub[sub] if is_persistent(sub._solver_plugin): # Note: what happens with non-persistent solvers? # - if rc is accepted as a model suffix by the solver (e.g. gurobi shell), it is loaded in postsolve @@ -281,11 +293,7 @@ def extract_and_store_reduced_costs(self): assert self._scenario_rc_buffer.size == self.send_buffers[Field.SCENARIO_REDUCED_COST].data_len() ci = 0 # buffer index for sub in self.opt.local_scenarios.values(): - consensus_groups = ( - _bundle_consensus_groups(sub) - if hasattr(sub, "_ef_scenario_names") - else None - ) + consensus_groups = consensus_groups_by_sub[sub] for ndn_i, xvar in sub._mpisppy_data.nonant_indices.items(): # fixed by modeler if ndn_i in self._modeler_fixed_nonants[sub]: diff --git a/mpisppy/extensions/grad_rho.py b/mpisppy/extensions/grad_rho.py index d17471daf..664af0b20 100644 --- a/mpisppy/extensions/grad_rho.py +++ b/mpisppy/extensions/grad_rho.py @@ -15,6 +15,7 @@ import mpisppy.MPI as MPI from mpisppy import global_toc import mpisppy.utils.sputils as sputils +from mpisppy.utils.nonant_sensitivities import _bundle_consensus_groups from mpisppy.cylinders.spwindow import Field class GradRho(mpisppy.extensions.dyn_rho_base.Dyn_Rho_extension_base): @@ -150,29 +151,18 @@ def _get_grad_exprs(self): active_obj = sputils.find_active_objective(s) if hasattr(s, "_ef_scenario_names"): - per_scen_to_bundle = {} - counters = {} - for (ndn, per_scen_i) in s.ref_vars.keys(): - if (ndn, per_scen_i) in s.ref_surrogate_vars: - continue - counters.setdefault(ndn, 0) - per_scen_to_bundle[(ndn, per_scen_i)] = (ndn, counters[ndn]) - counters[ndn] += 1 + # Shared helper builds the bundle position -> per-sub-scenario + # Vars mapping (see mpisppy/utils/nonant_sensitivities.py); we + # then flatten it into parallel wrt_vars/wrt_keys so we can + # differentiate once and re-aggregate the partials by key. + per_scen_vars = _bundle_consensus_groups(s) wrt_vars = [] wrt_keys = [] - per_scen_vars = {} - for scenario_name in s._ef_scenario_names: - scenario = s.component(scenario_name) - for node in scenario._mpisppy_node_list: - ndn = node.name - for per_scen_i, v in enumerate(node.nonant_vardata_list): - bundle_key = per_scen_to_bundle.get((ndn, per_scen_i)) - if bundle_key is None: - continue - wrt_vars.append(v) - wrt_keys.append(bundle_key) - per_scen_vars.setdefault(bundle_key, []).append(v) + for bundle_key, vars_at_pos in per_scen_vars.items(): + for v in vars_at_pos: + wrt_vars.append(v) + wrt_keys.append(bundle_key) partials = differentiate(active_obj, wrt_list=wrt_vars, diff --git a/mpisppy/tests/test_reduced_costs_rho_bundles.py b/mpisppy/tests/test_reduced_costs_rho_bundles.py index 06c579ccf..c03aeb297 100644 --- a/mpisppy/tests/test_reduced_costs_rho_bundles.py +++ b/mpisppy/tests/test_reduced_costs_rho_bundles.py @@ -191,7 +191,9 @@ def test_fails_when_only_ref_vars_have_rc(self): # must catch this. ref_vars = list(self.bundle._mpisppy_data.nonant_indices.values()) self._populate_rc(ref_vars) - with self.assertRaises(AssertionError) as ctx: + # RuntimeError, not AssertionError: the guard must survive `python -O` + # (which strips `assert` statements). + with self.assertRaises(RuntimeError) as ctx: _assert_consensus_rc_loaded(self.bundle, self.consensus_groups) # Error message should name a Var and the bundle position so the # failure is actionable. @@ -208,7 +210,7 @@ def test_fails_when_one_sub_scenario_var_missing(self): # Skip the first sub-scenario's first-position Var. first_skip = self.bundle.component(self.scen_names[0]).x[0] self._populate_rc([v for v in all_vars if v is not first_skip]) - with self.assertRaises(AssertionError): + with self.assertRaises(RuntimeError): _assert_consensus_rc_loaded(self.bundle, self.consensus_groups)