Pyomo · DLWoodruff · Aug 15, 2025 · Aug 15, 2025 · Aug 15, 2025 · Aug 15, 2025
diff --git a/doc/src/generic_admm.rst b/doc/src/generic_admm.rst
@@ -6,7 +6,7 @@ ADMM with ``generic_cylinders``
 The ``--admm`` and ``--stoch-admm`` flags allow ADMM-based decomposition
 to be used with any compatible model module through
 :ref:`generic_cylinders <generic_cylinders>`,
-eliminating the need for a bespoke driver script per problem.
+eliminating the need for a custom driver script per problem.
 
 There are two modes:
 
@@ -57,7 +57,7 @@ Here:
 - ``-np 3`` is one MPI rank per cylinder (1 hub + 2 spokes).
 
 The output will show PH iterations with bounds converging, just as with
-a bespoke ADMM driver.
+a custom ADMM driver.
 
 
 Running stochastic ADMM
@@ -77,8 +77,8 @@ with stochastic scenarios (random production losses). From that directory:
 Here:
 
 - ``--stoch-admm`` enables stochastic ADMM.
-- ``--num-admm-subproblems 3`` specifies three ADMM subproblems (regions).
-- ``--num-stoch-scens 3`` specifies three stochastic scenarios per region.
+- ``--num-admm-subproblems 3`` specifies three ADMM subproblems (regions). These are loaded into the ``admm_subproblem_names_creator`` via the config object.
+- ``--num-stoch-scens 3`` specifies three stochastic scenarios per region. These are loaded into the ``stoch_scenario_names_creator`` via the config object.
 - The total number of "scenarios" seen by mpi-sppy is
   ``num_admm_subproblems * num_stoch_scens = 9``.
 
@@ -136,14 +136,14 @@ Additional functions for ``--stoch-admm``
    :returns: dict mapping subproblem names to lists of
        ``(variable_name, stage)`` tuples
 
-.. py:function:: admm_subproblem_names_creator(num_admm_subproblems)
+.. py:function:: admm_subproblem_names_creator(cfg)
 
-   :param int num_admm_subproblems: number of ADMM subproblems
+   :param cfg: config object
    :returns: list of ADMM subproblem name strings
 
-.. py:function:: stoch_scenario_names_creator(num_stoch_scens)
+.. py:function:: stoch_scenario_names_creator(cfg)
 
-   :param int num_stoch_scens: number of stochastic scenarios
+   :param cfg: config object
    :returns: list of stochastic scenario name strings
 
 .. py:function:: admm_stoch_subproblem_scenario_names_creator(admm_subproblem_names, stoch_scenario_names)
@@ -323,8 +323,8 @@ Extending to Stochastic ADMM
 
 To support ``--stoch-admm``, additionally implement:
 
-1. ``admm_subproblem_names_creator(num_admm_subproblems)``
-2. ``stoch_scenario_names_creator(num_stoch_scens)``
+1. ``admm_subproblem_names_creator(cfg)``
+2. ``stoch_scenario_names_creator(cfg)``
 3. ``admm_stoch_subproblem_scenario_names_creator(admm_subproblem_names, stoch_scenario_names)``
 4. ``split_admm_stoch_subproblem_scenario_name(name)``
 
@@ -336,9 +336,6 @@ to determine both which ADMM subproblem and which stochastic scenario to build.
 Your ``consensus_vars_creator`` returns ``(variable_name, stage)`` tuples
 instead of plain strings.
 
-Register ``num_admm_subproblems`` and ``num_stoch_scens`` in your
-``inparser_adder`` so users can specify them on the command line.
-
 
 .. _admm_bundling:
 
@@ -361,7 +358,7 @@ Currently, full bundling is required: ``--scenarios-per-bundle`` must equal
        --module-name stoch_distr --stoch-admm \
        --num-admm-subproblems 2 --num-stoch-scens 4 \
        --default-rho 10 --max-iterations 50 --solver-name cplex \
-       --lagrangian --scenarios-per-bundle 4
+       --lagrangian --scenarios-per-bundle 4 --xhatxbar
 
 With ``--num-admm-subproblems 2`` and ``--scenarios-per-bundle 4``, PH sees
 only 2 bundles (one per subproblem) instead of 8 virtual scenarios.
@@ -432,14 +429,9 @@ Argument                                    Domain       Description
 ==========================================  ===========  ============================================
 ``--admm``                                  bool         Enable deterministic ADMM decomposition
 ``--stoch-admm``                            bool         Enable stochastic ADMM decomposition
-``--num-admm-subproblems``                  int          Number of ADMM subproblems (stoch-admm)
-``--num-stoch-scens``                       int          Number of stochastic scenarios (stoch-admm)
 ``--scenarios-per-bundle``                  int          Bundle stochastic scenarios (stoch-admm only)
 ==========================================  ===========  ============================================
 
 .. Note::
    For deterministic ADMM, the number of subproblems is given by ``--num-scens``,
-   which should be registered by the model's ``inparser_adder``.
-   For stochastic ADMM, ``--num-admm-subproblems`` and ``--num-stoch-scens``
-   may be registered by either the model's ``inparser_adder`` or automatically
-   by the generic framework (if not already present).
+   which should be registered by the model's ``inparser_adder``.
diff --git a/doc/src/generic_cylinders.rst b/doc/src/generic_cylinders.rst
@@ -127,6 +127,29 @@ Spokes provide bounds and heuristic solutions. Enable them with flags:
 
 See :ref:`Spokes` for details on each spoke type.
 
+Multistage Options
+-------------------
+
+For multistage problems (three or more stages), the model's
+``inparser_adder`` should register ``branching_factors`` via
+``cfg.multistage()`` or ``cfg.add_branching_factors()``.
+
+``--stage2-ef-solver-name``
+  Solver to use for forming second-stage EFs during xhat evaluation.
+  When set, the ``xhatshuffle`` spoke forms an EF for each second-stage
+  node, fixes the first-stage nonants, and solves.  The number of ranks
+  allocated to the xhatshuffle spoke must be an integer multiple of the
+  number of second-stage nodes.
+
+Example for the hydro model (three stages, branching factors 3 3):
+
+.. code-block:: bash
+
+    mpiexec -np 3 python -m mpi4py mpisppy/generic_cylinders.py \
+        --module-name hydro --solver-name cplex --max-iterations 100 \
+        --default-rho 1 --lagrangian --xhatshuffle --rel-gap 0.001 \
+        --branching-factors "3 3" --stage2-ef-solver-name cplex
+
 ADMM Decomposition
 -------------------
 

diff --git a/doc/src/spokes.rst b/doc/src/spokes.rst
@@ -121,16 +121,27 @@ for integer variables.
 stage2ef
 ~~~~~~~~
 
-An option for ``xhatshufflelooper_bounder`` is under development 
-for multistage problems that creates an EF for each second stage nodes by
-fixing the first stage nonanticipative variables.  This code requires
-that the number of ranks allocated to the ``xhatshufflelooper_bounder``
-is an integer multiple of the number of second stage nodes. Here is a 
-hint about how to to use it in a driver:
+An option for ``xhatshufflelooper_bounder`` for multistage problems that
+creates an EF for each second stage node by fixing the first stage
+nonanticipative variables.  This code requires that the number of ranks
+allocated to the ``xhatshufflelooper_bounder`` is an integer multiple of
+the number of second stage nodes.
+
+When using ``generic_cylinders.py``, this is enabled with the
+``--stage2-ef-solver-name`` flag:
+
+.. code-block:: bash
+
+    mpiexec -np 3 python -m mpi4py mpisppy/generic_cylinders.py \
+        --module-name hydro --solver-name cplex --max-iterations 100 \
+        --default-rho 1 --lagrangian --xhatshuffle --rel-gap 0.001 \
+        --branching-factors "3 3" --stage2-ef-solver-name cplex
+
+In a custom driver, set the spoke options directly:
 
 ::
 
-    xhatshuffle_spoke["opt_kwargs"]["options"]["stage2EFsolvern"] = solver_name
+    xhatshuffle_spoke["opt_kwargs"]["options"]["stage2_ef_solver_name"] = solver_name
     xhatshuffle_spoke["opt_kwargs"]["options"]["branching_factors"] = branching_factors
 
 An example is shown in ``examples.hydro.hydro_cylinders.py`` (this particular example

diff --git a/examples/generic_cylinders.bash b/examples/generic_cylinders.bash
@@ -151,7 +151,7 @@ mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name far
 # try a simple Hydro
 pwd
 cd hydro
-mpiexec -np 3 python -m mpi4py ../../mpisppy/generic_cylinders.py --module-name hydro --solver-name ${SOLVER} --max-iterations 100 --default-rho 1 --lagrangian --xhatshuffle --rel-gap 0.001 --branching-factors "3 3" --stage2EFsolvern ${SOLVER}
+mpiexec -np 3 python -m mpi4py ../../mpisppy/generic_cylinders.py --module-name hydro --solver-name ${SOLVER} --max-iterations 100 --default-rho 1 --lagrangian --xhatshuffle --rel-gap 0.001 --branching-factors "3 3" --stage2-ef-solver-name ${SOLVER}
 cd ..
 
 

diff --git a/examples/generic_tester.py b/examples/generic_tester.py
@@ -186,7 +186,7 @@ def do_one(dirname, modname, np, argstring, xhat_baseline_dir=None, tol=1e-6):
 
 hydroa = ("--max-iterations 100 --default-rho 1 "
           "--lagrangian --xhatshuffle --rel-gap 0.001 --branching-factors '3 3' "
-          f"--stage2EFsolvern {solver_name} --solver-name={solver_name}")
+          f"--stage2-ef-solver-name {solver_name} --solver-name={solver_name}")
 #rebaseline_xhat("hydro", "hydro", 3, hydroa, "test_data/hydroa_baseline")
 do_one("hydro", "hydro", 3, hydroa, xhat_baseline_dir="test_data/hydroa_baseline")
 
@@ -227,7 +227,7 @@ def do_one(dirname, modname, np, argstring, xhat_baseline_dir=None, tol=1e-6):
 hydroa_rc = ("--max-iterations 100 --default-rho 1 "
           "--reduced-costs --xhatshuffle --rel-gap 0.001 --branching-factors '3 3' "
           "--rc-fixer --reduced-costs-rho --reduced-costs-rho-multiplier=1.0 "
-          f"--stage2EFsolvern {solver_name} --solver-name={solver_name}")
+          f"--stage2-ef-solver-name {solver_name} --solver-name={solver_name}")
 #rebaseline_xhat("hydro", "hydro", 3, hydroa, "test_data/hydroa_baseline")
 do_one("hydro", "hydro", 3, hydroa_rc, xhat_baseline_dir = "test_data/hydroa_baseline")
 

diff --git a/examples/hydro/cfg.bash b/examples/hydro/cfg.bash
@@ -3,8 +3,8 @@
 
 SOLVERNAME=cplex
 
-mpiexec --oversubscribe --np 3 python -m mpi4py hydro_cylinders_config.py --branching-factors "3 3" --max-iterations=100 --default-rho=1 --xhatshuffle --lagrangian --solver-name=${SOLVERNAME} --stage2EFsolvern=${SOLVERNAME}
+mpiexec --oversubscribe --np 3 python -m mpi4py hydro_cylinders_config.py --branching-factors "3 3" --max-iterations=100 --default-rho=1 --xhatshuffle --lagrangian --solver-name=${SOLVERNAME} --stage2-ef-solver-name=${SOLVERNAME}
 #--tee-rank0-solves
 
 # including this option will cause the upper bounder to solve the EF since there are only three ranks in total.
-#--stage2EFsolvern=${SOLVERNAME}
+#--stage2-ef-solver-name=${SOLVERNAME}
diff --git a/examples/hydro/demo.bash b/examples/hydro/demo.bash
@@ -2,7 +2,7 @@
 
 SOLVERNAME=cplex
 
-mpiexec --oversubscribe --np 3 python -m mpi4py hydro_cylinders.py --branching-factors "3 3" --max-iterations=100 --default-rho=1 --xhatshuffle --lagrangian --solver-name=${SOLVERNAME} --stage2EFsolvern=${SOLVERNAME}
+mpiexec --oversubscribe --np 3 python -m mpi4py hydro_cylinders.py --branching-factors "3 3" --max-iterations=100 --default-rho=1 --xhatshuffle --lagrangian --solver-name=${SOLVERNAME} --stage2-ef-solver-name=${SOLVERNAME}
 
 # including this option will cause the upper bounder to solve the EF since there are only three ranks in total.
-#--stage2EFsolvern=${SOLVERNAME}
+#--stage2-ef-solver-name=${SOLVERNAME}
diff --git a/examples/hydro/hydro.py b/examples/hydro/hydro.py
@@ -250,10 +250,6 @@ def scenario_denouement(rank, scenario_name, scenario):
 #=========
 def inparser_adder(cfg):
     cfg.multistage()
-    cfg.add_to_config(name ="stage2EFsolvern",
-                      description="Solver to use for xhatlooper stage2ef option (default None)",
-                      domain = str,
-                      default=None)
     cfg.add_to_config(name ="hydro_data_path",
                       description="Path to hydro data (seldom used; default None)",
                       domain = str,

diff --git a/examples/hydro/hydro_cylinders.py b/examples/hydro/hydro_cylinders.py
@@ -31,11 +31,6 @@ def _parse_args():
     cfg.lagrangian_args()
     cfg.xhatspecific_args()
 
-    cfg.add_to_config(name ="stage2EFsolvern",
-                         description="Solver to use for xhatlooper stage2ef option (default None)",
-                         domain = str,
-                         default=None)
-
     cfg.parse_command_line("hydro_cylinders")
     return cfg
 
@@ -96,9 +91,9 @@ def main():
     if xhatshuffle:
         list_of_spoke_dict.append(xhatshuffle_spoke)
 
-    if cfg.stage2EFsolvern is not None:
-        assert xhatshuffle is not None, "xhatshuffle is required for stage2EFsolvern"
-        xhatshuffle_spoke["opt_kwargs"]["options"]["stage2EFsolvern"] = cfg["stage2EFsolvern"]
+    if cfg.stage2_ef_solver_name is not None:
+        assert xhatshuffle is not None, "xhatshuffle is required for stage2_ef_solver_name"
+        xhatshuffle_spoke["opt_kwargs"]["options"]["stage2_ef_solver_name"] = cfg["stage2_ef_solver_name"]
         xhatshuffle_spoke["opt_kwargs"]["options"]["branching_factors"] = cfg["branching_factors"]
 
     wheel = WheelSpinner(hub_dict, list_of_spoke_dict)

diff --git a/examples/run_all.py b/examples/run_all.py
@@ -247,7 +247,7 @@ def do_one_mmw(dirname, modname, runefstring, npyfile, mmwargstring):
     do_one("hydro", "hydro_cylinders.py", 3,
            "--branching-factors \'3 3\' --max-iterations=100 "
            "--default-rho=1 --xhatshuffle --lagrangian "
-           "--solver-name={} --stage2EFsolvern={}".format(solver_name, solver_name))
+           "--solver-name={} --stage2-ef-solver-name={}".format(solver_name, solver_name))
 
     do_one("hydro", "hydro_cylinders_pysp.py", 3,
            "--max-iterations=100 "

diff --git a/examples/run_more.py b/examples/run_more.py
@@ -164,7 +164,7 @@ def time_one(ID, dirname, progname, np, argstring):
        f"--num-scens 3 --crops-multiplier=1  --EF-solver-name={solver_name} "
        "--BPL-c0 25 --BPL-eps 100 --confidence-level 0.95 --BM-vs-BPL BPL")
 
-# --- Hydro without stage2EFsolvern (also covered by generic_tester) ---
+# --- Hydro without stage2_ef_solver_name (also covered by generic_tester) ---
 
 do_one("hydro", "hydro_cylinders.py", 3,
        "--branching-factors \"3 3\" --max-iterations=100 "

diff --git a/examples/stoch_distr/stoch_distr.py b/examples/stoch_distr/stoch_distr.py
@@ -262,28 +262,28 @@ def consensus_vars_creator(admm_subproblem_names, stoch_scenario_name, inter_reg
     return consensus_vars
 
 
-def stoch_scenario_names_creator(num_stoch_scens):
+def stoch_scenario_names_creator(cfg):
     """Creates the name of every stochastic scenario.
 
     Args:
-        num_stoch_scens (int): number of stochastic scenarios
+        cfg: config object
 
     Returns:
         list (str): the list of stochastic scenario names
     """
-    return [f"StochasticScenario{i+1}" for i in range(num_stoch_scens)]
+    return [f"StochasticScenario{i+1}" for i in range(cfg.num_stoch_scens)]
 
 
-def admm_subproblem_names_creator(num_admm_subproblems):
+def admm_subproblem_names_creator(cfg):
     """Creates the name of every admm subproblem.
 
     Args:
-        num_subproblems (int): number of admm subproblems
+        cfg: config object
 
     Returns:
         list (str): the list of admm subproblem names
     """
-    return [f"Region{i+1}" for i in range(num_admm_subproblems)]
+    return [f"Region{i+1}" for i in range(cfg.num_admm_subproblems)]
 
 
 def combining_names(admm_subproblem_name,stoch_scenario_name):

diff --git a/examples/stoch_distr/stoch_distr_admm_cylinders.py b/examples/stoch_distr/stoch_distr_admm_cylinders.py
@@ -62,8 +62,8 @@ def _count_cylinders(cfg):
 def _make_admm(cfg, n_cylinders, verbose=None):
     options = {}
 
-    admm_subproblem_names = stoch_distr.admm_subproblem_names_creator(cfg.num_admm_subproblems)
-    stoch_scenario_names = stoch_distr.stoch_scenario_names_creator(num_stoch_scens=cfg.num_stoch_scens)
+    admm_subproblem_names = stoch_distr.admm_subproblem_names_creator(cfg)
+    stoch_scenario_names = stoch_distr.stoch_scenario_names_creator(cfg)
     all_admm_stoch_subproblem_scenario_names = stoch_distr.admm_stoch_subproblem_scenario_names_creator(admm_subproblem_names,stoch_scenario_names)
 
     split_admm_stoch_subproblem_scenario_name = stoch_distr.split_admm_stoch_subproblem_scenario_name

diff --git a/mpisppy/cylinders/xhatshufflelooper_bounder.py b/mpisppy/cylinders/xhatshufflelooper_bounder.py
@@ -40,13 +40,13 @@ def try_scenario_dict(self, xhat_scenario_dict):
         """ wrapper for _try_one"""
         snamedict = xhat_scenario_dict
 
-        stage2EFsolvern = self.opt.options.get("stage2EFsolvern", None)
+        stage2_ef_solver_name = self.opt.options.get("stage2_ef_solver_name", None)
         branching_factors = self.opt.options.get("branching_factors", None)  # for stage2ef
         obj = self.xhatter._try_one(snamedict,
                                     solver_options = self.solver_options,
                                     verbose=False,
                                     restore_nonants=True,
-                                    stage2EFsolvern=stage2EFsolvern,
+                                    stage2_ef_solver_name=stage2_ef_solver_name,
                                     branching_factors=branching_factors)
         def _vb(msg):
             if self.verbose and self.opt.cylinder_rank == 0:

diff --git a/mpisppy/extensions/xhatbase.py b/mpisppy/extensions/xhatbase.py
@@ -40,7 +40,7 @@ def __init__(self, opt):
 
      #**********
     def _try_one(self, snamedict, solver_options=None, verbose=False,
-                 restore_nonants=True, stage2EFsolvern=None, branching_factors=None):
+                 restore_nonants=True, stage2_ef_solver_name=None, branching_factors=None):
         """ try the scenario named sname in self.opt.local_scenarios
        Args:
             snamedict (dict): key: scenario tree non-leaf name, val: scen name
@@ -50,7 +50,7 @@ def _try_one(self, snamedict, solver_options=None, verbose=False,
             restore_nonants (bool): if True, restores the nonants to their original
                                     values in all scenarios. If False, leaves the
                                     nonants as they are in the tried scenario
-            stage2EFsolvern: use this for EFs based on second stage nodes for multi-stage
+            stage2_ef_solver_name: use this for EFs based on second stage nodes for multi-stage
             branching_factors (list): list of branching factors for stage2ef
        NOTE: The solve loop is with fixed nonants so W and rho do not
              matter to the optimization. When we want to check the obj
@@ -88,7 +88,7 @@ def _try_one(self, snamedict, solver_options=None, verbose=False,
                       .format(src_rank))
                 print("root comm size={}".format(self.comms["ROOT"].size))
                 raise
-        elif stage2EFsolvern is None:  # regular multi-stage
+        elif stage2_ef_solver_name is None:  # regular multi-stage
             # assemble parts and put it in xhats
             # send to ranks in the comm or receive ANY_SOURCE
             # (for closest do allreduce with loc operator) rank
@@ -180,8 +180,8 @@ def _try_one(self, snamedict, solver_options=None, verbose=False,
             for ndn2, sdict in local_2ndns.items():  # ndn2 will be a the node name
                 wxbarutils.fix_ef_ROOT_nonants(self._EFs[ndn2], xhats["ROOT"])
                 # solve EF
-                solver = pyo.SolverFactory(stage2EFsolvern)
-                if 'persistent' in stage2EFsolvern:
+                solver = pyo.SolverFactory(stage2_ef_solver_name)
+                if 'persistent' in stage2_ef_solver_name:
                     solver.set_instance(self._EFs[ndn2], symbolic_solver_labels=True)
                     results = solver.solve(tee=Tee)
                 else:

diff --git a/mpisppy/generic/admm.py b/mpisppy/generic/admm.py
@@ -31,6 +31,9 @@ def admm_args(cfg):
         cfg.add_to_config("num_stoch_scens",
                           description="Number of stochastic scenarios (stoch-admm only)",
                           domain=int, default=None)
+
+    cfg.add_branching_factors()
+    cfg.add_stage2_ef_solver_name_arg()
 
 
 def _count_cylinders(cfg):
@@ -134,8 +137,8 @@ def setup_stoch_admm(module, cfg, n_cylinders):
         tuple: (scenario_creator, scenario_creator_kwargs,
                 all_scenario_names, all_nodenames)
     """
-    admm_subproblem_names = module.admm_subproblem_names_creator(cfg.num_admm_subproblems)
-    stoch_scenario_names = module.stoch_scenario_names_creator(cfg.num_stoch_scens)
+    admm_subproblem_names = module.admm_subproblem_names_creator(cfg)
+    stoch_scenario_names = module.stoch_scenario_names_creator(cfg)
     all_names = module.admm_stoch_subproblem_scenario_names_creator(
         admm_subproblem_names, stoch_scenario_names)
 
@@ -181,8 +184,8 @@ def setup_stoch_admm_with_bundles(module, cfg, n_cylinders):
     """
     from mpisppy.utils.admm_bundler import AdmmBundler
 
-    admm_subproblem_names = module.admm_subproblem_names_creator(cfg.num_admm_subproblems)
-    stoch_scenario_names = module.stoch_scenario_names_creator(cfg.num_stoch_scens)
+    admm_subproblem_names = module.admm_subproblem_names_creator(cfg)
+    stoch_scenario_names = module.stoch_scenario_names_creator(cfg)
 
     scenario_creator_kwargs = module.kw_creator(cfg)
     stoch_scenario_name = stoch_scenario_names[0]