Merge PRs #651 + #652: stage2_ef_solver_name and stoch_admm changes

[DLWoodruff]

Combined merge of #651 and #652 with #651 taking precedence on naming.

From #651 (unchanged)

• Renames stage2EFsolvern → stage2_ef_solver_name throughout.
• Registers the option via Config.xhatshuffle_args() so --stage2-ef-solver-name is available from generic_cylinders.py.
• Relaxes the multistage xhatshuffle hard assert to a warning in mpisppy/generic/parsing.py.
• Adds docs and a hydro test_xhatshuffle_stage2ef MPI test.

From #652 (most of it, renamed to match #651)

• admm_subproblem_names_creator and stoch_scenario_names_creator now take the cfg object instead of specific scalar args (enables use cases like CEP that split by date range).
• add_stage2_ef_solver_name_arg() helper registered from both multistage() and admm_args(), with dedup guards so calling from multiple CLIs doesn't error.
• add_branching_factors() also guarded against double-registration.
• find_component fallback in Stoch_AdmmWrapper.assign_variable_probs to quote-wrap bracketed index strings.
• get_scenario_unscaled() helper on Stoch_AdmmWrapper (covered by new test).
• Indentation fix in mpisppy/spopt.py.

Dropped from #652

• Rank-slicing disablement in Stoch_AdmmWrapper (would have made every rank hold all scenarios — suspected regression).
• Verbose debug prints in xhatxbar_bounder.py.

Conflict resolution

• 652's add_stage2EF() helper renamed to add_stage2_ef_solver_name_arg() and key changed to stage2_ef_solver_name to match 651's convention.
• Registration in multistage(), xhatshuffle_args(), and admm_args() all go through the guarded helper; no duplicate-key errors.

Test plan

• ruff check . clean
• python -m unittest mpisppy.tests.test_admm_bundler — 18/18 pass
• python -m unittest mpisppy.tests.test_config — pass
• python -m unittest mpisppy.tests.test_stoch_admmWrapper.TestStochAdmmWrapper.test_get_scenario_unscaled — pass
• mpiexec -np 2 python -m mpi4py -m unittest mpisppy.tests.test_with_cylinders.Test_hydro_with_cylinders.test_xhatshuffle_stage2ef — pass
• Full CI

Note: a few pre-existing failures on upstream/main (test_scenario_lpwriter_extension, test_admmWrapper.test_values, test_with_cylinders.test_xhatshuffle_coverage) are unaffected by this merge.

Addendum

stoch-admm EF partial-consensus fix + user-surrogate preservation

[codecov]

Codecov Report

❌ Patch coverage is 83.63636% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.73%. Comparing base (b726347) to head (a90eeb9).

Files with missing lines	Patch %	Lines
mpisppy/generic/admm.py	33.33%	4 Missing ⚠️
mpisppy/utils/sputils.py	25.00%	3 Missing ⚠️
mpisppy/generic_cylinders.py	0.00%	1 Missing ⚠️
mpisppy/opt/fwph.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #653      +/-   ##
==========================================
+ Coverage   70.71%   70.73%   +0.02%     
==========================================
  Files         153      153              
  Lines       19005    19041      +36     
==========================================
+ Hits        13439    13469      +30     
- Misses       5566     5572       +6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

This PR merges prior work to standardize the stage2_ef_solver_name option across the codebase/CLI/docs and to improve stochastic-ADMM interoperability with EF/PH (including partial-consensus behavior and config-driven scenario/subproblem naming).

Changes:

• Renames and wires up stage2_ef_solver_name end-to-end (Config registration, generic CLI/spokes, examples, and docs) and relaxes multistage xhatshuffle’s hard requirement to a warning.
• Updates stochastic-ADMM plumbing: creators now accept cfg, Stoch_AdmmWrapper better handles partial consensus with EF (surrogate-tagged dummies + EF ref handling), and exposes get_scenario_unscaled.
• Adds/updates tests and documentation to cover the new option wiring and ADMM/EF behaviors.

Reviewed changes

Copilot reviewed 29 out of 30 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
mpisppy/utils/stoch_admmWrapper.py	Tags dummy consensus vars as surrogates and preserves user nonant metadata during node rewrites; adds component-name fallback and get_scenario_unscaled.
mpisppy/utils/sputils.py	Adjusts EF creation to “upgrade” surrogate refs to real vars for partial-consensus ADMM cases.
mpisppy/utils/config.py	Adds guarded registration helpers for branching_factors and stage2_ef_solver_name; registers stage2 solver name from multistage()/xhatshuffle_args().
mpisppy/utils/cfg_vanilla.py	Forwards turn_off_names_check into EF options dict (and auto-disables when bundled).
mpisppy/tests/test_with_cylinders.py	Adds hydro multistage test for xhatshuffle with stage2_ef_solver_name.
mpisppy/tests/test_stoch_admmWrapper.py	Adds regression tests for user surrogate/EF-supplemental preservation and EF partial-consensus correctness.
mpisppy/tests/test_admm_bundler.py	Updates creator signatures to take cfg; adds tests for multistage warning and EF option forwarding.
mpisppy/tests/examples/stoch_distr/stoch_distr_admm_cylinders.py	Updates stoch_distr creator calls to pass cfg.
mpisppy/tests/examples/stoch_distr/stoch_distr.py	Updates scenario/subproblem name creators to accept cfg.
mpisppy/spopt.py	Fixes indentation for extension hook calls.
mpisppy/opt/fwph.py	Updates option name usage to stage2_ef_solver_name.
mpisppy/generic_cylinders.py	Auto-disables nonant name checking for ADMM wrappers (dummy/renamed components).
mpisppy/generic/spokes.py	Passes stage2_ef_solver_name + branching_factors into xhatshuffle spoke options for multistage.
mpisppy/generic/parsing.py	Replaces multistage xhatshuffle assert with a warning when stage2_ef_solver_name is missing.
mpisppy/generic/admm.py	Uses cfg-driven name creators; registers branching factors + stage2 EF solver arg for ADMM CLI path.
mpisppy/extensions/xhatbase.py	Renames stage2 EF solver kwarg and updates usage in xhat evaluation path.
mpisppy/cylinders/xhatshufflelooper_bounder.py	Renames option key plumbing for stage2 EF solve path.
examples/stoch_distr/stoch_distr_admm_cylinders.py	Updates creator calls to pass cfg.
examples/stoch_distr/stoch_distr.py	Updates scenario/subproblem name creators to accept cfg.
examples/run_more.py	Updates comment text to new option name.
examples/run_all.py	Updates CLI flag to --stage2-ef-solver-name.
examples/hydro/hydro_cylinders.py	Updates config key/usage to stage2_ef_solver_name.
examples/hydro/hydro.py	Updates config key registration to stage2_ef_solver_name.
examples/hydro/demo.bash	Updates CLI flag to --stage2-ef-solver-name.
examples/hydro/cfg.bash	Updates CLI flag to --stage2-ef-solver-name.
examples/generic_tester.py	Updates CLI flag to --stage2-ef-solver-name.
examples/generic_cylinders.bash	Updates CLI flag to --stage2-ef-solver-name.
doc/src/spokes.rst	Documents the new flag and option key name for stage2 EF behavior.
doc/src/generic_cylinders.rst	Adds multistage documentation for --stage2-ef-solver-name.
doc/src/generic_admm.rst	Updates docs to reflect cfg-driven creator signatures and related usage notes.

Comments suppressed due to low confidence (1)

mpisppy/utils/stoch_admmWrapper.py:227

• varlist[stage-1].append(v) runs even when v is None (e.g., when a consensus var is listed but not found on the scenario model). This can trigger an AttributeError in sputils.build_vardatalist() during ScenarioNode construction before the later error_list1/error_list2 RuntimeError is raised. Consider raising immediately when v is None, or skip appending/creating nodes until after validation succeeds.

                        dummy_varlist[stage-1].append(v)
                    else:
                        error_list2.append((sname,vstr))
                varlist[stage-1].append(v)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

When rebuilding each existing ScenarioNode, user_surrogates = list(old_node.surrogate_vardatas) converts a set-like ComponentSet to a list. This can produce non-deterministic ordering, which can change the positional layout of surrogate nonants across scenarios. Prefer preserving the original user-provided ordering by passing old_node.surrogate_nonant_list (and similarly old_node.nonant_ef_suppl_list) into the new ScenarioNode instead of reconstructing from the vardata sets/lists.

[DLWoodruff]

Good catch to flag this, but node.surrogate_vardatas is a Pyomo ComponentSet, which is backed by a dict keyed by an id-based hasher. Iteration yields self._data.values(), so order is deterministic (Python dict insertion order) — list(old_node.surrogate_vardatas) would have preserved the order it was built in.

That said, the edit is still worth taking: passing the raw user surrogate_nonant_list / nonant_ef_suppl_list avoids the vardata round-trip and matches what the user originally supplied, so the change landed as suggested.

[Copilot]

The quote-wrapping fallback assumes vstr contains [...] and uses find('[') / find(']') unconditionally. If vstr has no brackets, find returns -1 and this will generate a truncated/invalid component name. Guard this logic with an explicit bracket check (or use parsing that can handle missing/complex indices) before attempting the wrapped lookup.

Suggested change

	# try quote wrap around variable name
	vvstr = vstr[:vstr.find('[')]
	tvstr = vstr[vstr.find('[')+1:vstr.find(']')]
	v = s.find_component(vvstr+'["'+tvstr+'"]')
	# try quote wrap around variable name when vstr has a
	# simple bracketed index (e.g., name[index])
	lbracket = vstr.find('[')
	rbracket = vstr.find(']')
	if lbracket != -1 and rbracket != -1 and lbracket < rbracket:
	vvstr = vstr[:lbracket]
	tvstr = vstr[lbracket+1:rbracket]
	v = s.find_component(vvstr+'["'+tvstr+'"]')

[Copilot]

The new warnings.warn(...) will point at parsing.py as the source. Setting stacklevel=2 (and optionally a warning category) will make the warning more actionable by pointing users to their call site / CLI invocation that triggered multistage xhatshuffle without stage2_ef_solver_name.

Suggested change

	warnings.warn("stage2_ef_solver_name is recommended for multistage xhatshuffle")
	warnings.warn(
	"stage2_ef_solver_name is recommended for multistage xhatshuffle",
	UserWarning,
	stacklevel=2,
	)

[Copilot]

cfg.multistage() now registers stage2_ef_solver_name (via Config.add_stage2_ef_solver_name_arg()), so this manual add_to_config(name="stage2_ef_solver_name", ...) is redundant and will be silently ignored as a duplicate. Consider removing it (or switching to the shared helper) to avoid confusion about which option definition/description is actually active.

Suggested change

	cfg.add_to_config(name ="stage2_ef_solver_name",
	description="Solver to use for xhatlooper stage2ef option (default None)",
	domain = str,
	default=None)

[Copilot]

This file calls cfg.multistage() (which now registers stage2_ef_solver_name) and then manually calls add_to_config(name="stage2_ef_solver_name", ...) again. Since Config.add_to_config ignores duplicates, this second declaration is a no-op and can be misleading—consider removing it or using the shared helper consistently.