Merge branch 'main' into fix_xhatshuffle

Author: bknueven

doc/src/properbundles.rst

@@ -68,16 +68,22 @@ there are two modules that have most of the support for proper bundles:
   - ``mpisppy.utils.proper_bundler.py`` has wrappers for cylinder programs
 
 
-Multistage and notes
---------------------
+Multistage
+----------
 
-At the time of this writing, multi-stage proper, pickled bundles is a
-little bit beyond the bleeding edge.  The idea is that bundles are
-formed and then saved as dill pickle files for rapid retrieval. The
+The most flexible way to create proper bundles is to write
+your own problem-specific code to do it. The
 file ``aircond_cylinders.py`` in the aircond example directory
 provides an example.  The latter part of the ``allways.bash`` script
 demonstrates how to run it.
 
+There is support for multi-stage bundles in mpi-sppy, but the scenario
+probabilities must be uniform and the bundles must span the same number
+of entire second stage nodes.
+
+Notes
+-----
+
 Pickled bundles are clearly useful for algorithm tuning and algorithm
 experimentation. In some, but not all, settings they can also improve
 wall-clock performance for a single optimization run. The pickler

examples/generic_cylinders.bash

@@ -4,6 +4,21 @@
 SOLVER="cplex"
 SPB=1
 
+echo "^^^ Multi-stage AirCond ^^^"
+mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name mpisppy.tests.examples.aircond --branching-factors "3 3 3" --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --lagrangian --xhatxbar --rel-gap 0.01 --solution-base-name aircond_nonants
+# --xhatshuffle --stag2EFsolvern
+
+echo "^^^ Multi-stage AirCond, pickle  the scenarios ^^^"
+mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name mpisppy.tests.examples.aircond --branching-factors "3 3 3" --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --lagrangian --xhatxbar --rel-gap 0.01 --solution-base-name aircond_nonants --pickle-scenarios-dir aircond/pickles
+
+echo "^^^ Multi-stage AirCond, bundle the scenarios ^^^"
+mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name mpisppy.tests.examples.aircond --branching-factors "3 3 3" --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --lagrangian --xhatxbar --rel-gap 0.01 --solution-base-name aircond_nonants --scenarios-per-bundle 9
+
+echo "^^^ Multi-stage AirCond, bundle the scenarios and write ^^^"
+mpiexec -np 3 python -m mpi4py ../mpisppy/generic_cylinders.py --module-name mpisppy.tests.examples.aircond --branching-factors "3 3 3" --solver-name ${SOLVER} --max-iterations 10 --max-solver-threads 4 --default-rho 1 --lagrangian --xhatxbar --rel-gap 0.01 --solution-base-name aircond_nonants --pickle-scenarios-dir aircond/pickles --scenarios-per-bundle 9
+
+#### HEY! check on error messages for bad bundle sizes
+
 echo "^^^ write scenario lp and nonant json files ^^^"
 cd sizes
 python ../../mpisppy/generic_cylinders.py --module-name sizes --num-scens 3 --default-rho 1 --solver-name ${SOLVER} --max-iterations 0 --scenario-lpfiles
@@ -20,8 +35,6 @@ echo "^^^ unpickle the sizes bundles and write the lp and nonant files ^^^"
 cd sizes
 python ../../mpisppy/generic_cylinders.py --module-name sizes --num-scens 10 --default-rho 1 --solver-name ${SOLVER} --max-iterations 0 --scenario-lpfiles --unpickle-bundles-dir sizes_pickles --scenarios-per-bundle 5
 cd ..
-echo "xxxx Early exit. xxxx"
-exit
 
 echo "^^^ pickle the scenarios ^^^"
 cd farmer

examples/generic_tester.py

@@ -163,6 +163,12 @@ def do_one(dirname, modname, np, argstring, xhat_baseline_dir=None, tol=1e-6):
 #rebaseline_xhat("hydro", "hydro", 3, hydroa, "test_data/hydroa_baseline")
 do_one("hydro", "hydro", 3, hydroa, xhat_baseline_dir="test_data/hydroa_baseline")
 
+# write hydro bundles for at least some testing of multi-stage proper bundles
+# (just looking for smoke)
+hydro_wr = ("--pickle-bundles-dir hydro_pickles --scenarios-per-bundle 3"
+            "--branching-factors '3 3' ")
+do_one("hydro", "hydro", 3, hydro_wr, xhat_baseline_dir=None)
+
 # write, then read, pickled scenarios
 print("starting write/read pickled scenarios")
 farmer_wr = "--pickle-scenarios-dir farmer_pickles --crops-mult 2 --num-scens 10"

mpisppy/generic_cylinders.py

@@ -96,8 +96,6 @@ def _name_lists(module, cfg, bundle_wrapper=None):
         num_scens = np.prod(cfg.branching_factors)
         assert not cfg.xhatshuffle or cfg.get("stage2EFsolvern") is not None,\
             "For now, stage2EFsolvern is required for multistage xhat"
-        assert cfg.scenarios_per_bundle is None, "proper bundles in generic_cylinders does not yet support multistage"
-
     else:
         all_nodenames = None
         num_scens = cfg.num_scens

mpisppy/spopt.py

@@ -189,8 +189,9 @@ def _vb(msg):
         if Ag is not None:
             assert not disable_pyomo_signal_handling, "Not thinking about agnostic APH yet"
             kws = {"s": s, "solve_keyword_args": solve_keyword_args, "gripe": gripe, "tee": tee, "need_solution": need_solution}
-            Ag.callout_agnostic(kws)
+            Ag.callout_agnostic(kws)  # not going to use the return values
         else:
+            #  didcallout = False  (returned true by the callout, but not used)
             try:
                 results = s._solver_plugin.solve(s,
                                                  **solve_keyword_args,

mpisppy/tests/examples/aircond.py

@@ -319,17 +319,8 @@ def scenario_creator(sname, **kwargs):
 
     #Constructing the nodes used by the scenario
     model._mpisppy_node_list = MakeNodesforScen(model, nodenames, branching_factors)
-    model._mpisppy_probability = 1 / np.prod(branching_factors)
-    """
-    from mpisppy import MPI
-
-    comm = MPI.COMM_WORLD
-    rank = comm.Get_rank()
-    if rank == 0:
-        with open("efmodel.txt", "w") as fileh:
-            model.pprint(fileh)
-    quit()
-    """
+    #model._mpisppy_probability = 1 / np.prod(branching_factors)
+    model._mpisppy_probability = "uniform"
     return(model)
 
 
@@ -427,7 +418,12 @@ def kw_creator(cfg, optionsin=None):
     create the key word arguments for the creator fucntion(s)
     args:
         cfg (Config object): probably has parsed values
-        optionsin (dict): programmatic options"""
+        optionsin (dict): programmatic options
+    returns:
+        kwargs (dict): the keyword args
+    side-effect:
+        checks and/or adds cfg.num_scens 
+    """
     # use an empty dict instead of None
     options = optionsin if optionsin is not None else dict()
     if "kwargs" in options:
@@ -443,7 +439,7 @@ def _kwarg(option_name, default = None, arg_name=None):
             return
         # if not in the options, see if it is availalbe globally
         aname = option_name if arg_name is None else arg_name
-        retval = getattr(cfg, aname) if hasattr(cfg, aname) else None
+        retval = cfg.get(aname)
         retval = default if retval is None else retval
         kwargs[option_name] = retval
 
@@ -457,7 +453,21 @@ def _kwarg(option_name, default = None, arg_name=None):
         raise ValueError(f"kw_creator called, but no value given for start_ups, options={options}")
     if kwargs["start_seed"] is None:
         raise ValueError(f"kw_creator called, but no value given for start_seed, options={options}")
-
+    if kwargs["branching_factors"] is not None:
+        BFs = kwargs["branching_factors"]
+        ns = cfg.get("num_scens")
+        if BFs is not None:
+            if ns is None:
+                cfg.add_and_assign("num_scens",
+                                   description="Number of scenarios",
+                                   domain=int,
+                                   value=np.prod(BFs),
+                                   default=None
+                                   )
+            else:
+                pass  # we cannot do the assert because of tree sampling
+                #assert ns == np.prod(BFs),\
+                #f"num_scens != prod(BFs); {ns=}, {BFs=}"
     return kwargs
 
 

mpisppy/tests/test_pickle_bundle.py

@@ -6,7 +6,7 @@
 # All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for
 # full copyright and license information.
 ###############################################################################
-# Provide some test for pickled bundles
+# Provide a test for special-purpose (aircond only)multi-stage pickled bundles
 """
 
 """

mpisppy/utils/proper_bundler.py

@@ -10,6 +10,7 @@
 # TBD: we should consider restructuring this and moving the capability to spbase
 
 import os
+import numpy as np
 import mpisppy.utils.sputils as sputils
 import mpisppy.utils.pickle_bundle as pickle_bundle
 
@@ -19,8 +20,12 @@
 #  - read a bundle
 #  - make a bundle 
 #  - make a bundle and write it
-# Multi-stage (as of Sept 2024) not supported in class or generic_cylinders
-#  You need to do something clever like in aircondB
+# Multi-stage
+#  Is not very special because all we do is make sure that
+#  bundles cover entire second-stage nodes so the new bundled
+#  problem is a two-stage problem no matter how many stages
+#  were in the original problem. As of Dec 2024, support
+#  is provided only when there are branching factors.
 # NOTE:: the caller needs to make sure it is two stage
 #        the caller needs to worry about what is in what rank
 #        (local_scenarios might have bundle names, e.g.)
@@ -50,14 +55,28 @@ def scenario_names_creator(self, num_scens, start=None, cfg=None):
         return cfg.model.scenario_names_creator(num_scens, start=start)
 
     def bundle_names_creator(self, num_buns, start=None, cfg=None):
+
+        def _multistage_check(bunsize):
+            # returns bunBFs as a side-effect
+            BFs = cfg.branching_factors
+            beyond2size = np.prod(BFs[1:])
+            if bunsize % beyond2size!= 0:
+                raise RuntimeError(f"Bundles must consume the same number of entire second stage nodes: {beyond2size=} {bunsize=}")
+            # we need bunBFs for EF formulation
+            self.bunBFs = [bunsize // beyond2size] + BFs[1:]
+        
         # start refers to the bundle number; bundles are always zero-based
         if start is None:
             start = 0
         assert cfg is not None, "ProperBundler needs cfg for bundle names"
         assert cfg.get("num_scens") is not None
         assert cfg.get("scenarios_per_bundle") is not None
-        assert cfg.num_scens % cfg.scenarios_per_bundle == 0
+        assert cfg.num_scens % cfg.scenarios_per_bundle == 0, "Bundles must consume the same number of entire second stage nodes: {cfg.num_scens=} {bunsize=}"
         bsize = cfg.scenarios_per_bundle  # typing aid
+        if cfg.get("branching_factors") is not None:
+            _multistage_check(bsize)
+        else:
+            self.bunBFs = None
         # We need to know if scenarios (not bundles) are one-based.
         inum = sputils.extract_num(self.module.scenario_names_creator(1)[0])
         names = [f"Bundle_{bn*bsize+inum}_{(bn+1)*bsize-1+inum}" for bn in range(start+num_buns)]
@@ -80,42 +99,50 @@ def scenario_creator(self, sname, **kwargs):
         cfg = kwargs["cfg"]
         if "scen" in sname or "Scen" in sname:
             # In case the user passes in kwargs from scenario_creator_kwargs.
-            return self.module.scenario_creator(sname, {**self.original_kwargs, **kwargs})
+            return self.module.scenario_creator(sname, **{**self.original_kwargs, **kwargs})
 
         elif "Bundle" in sname and cfg.get("unpickle_bundles_dir") is not None:
             fname = os.path.join(cfg.unpickle_bundles_dir, sname+".pkl")
             bundle = pickle_bundle.dill_unpickle(fname)
             return bundle
         elif "Bundle" in sname and cfg.get("unpickle_bundles_dir") is None:
-            # this is also the branch for proper_no_files
             # If we are still here, we have to create the bundle.
             firstnum = int(sname.split("_")[1])  # sname is a bundle name
             lastnum = int(sname.split("_")[2])
             # snames are scenario names
             snames = self.module.scenario_names_creator(lastnum-firstnum+1,
                                                         firstnum)
+            kws = self.original_kwargs
+            if self.bunBFs is not None:
+                # The original scenario creator needs to handle these
+                kws["branching_factors"] = self.bunBFs
+
             # We are assuming seeds are managed by the *scenario* creator.
             bundle = sputils.create_EF(snames, self.module.scenario_creator,
-                                       scenario_creator_kwargs=self.original_kwargs,
+                                       scenario_creator_kwargs=kws,
                                        EF_name=sname,
                                        suppress_warnings=True,                                       
                                        nonant_for_fixed_vars = False)
 
             nonantlist = [v for idx, v in bundle.ref_vars.items() if idx[0] =="ROOT"]
             # if the original scenarios were uniform, this needs to be also
             # (EF formation will recompute for the bundle if uniform)
+            # Get an arbitrary scenario.
             scen = self.module.scenario_creator(snames[0], **self.original_kwargs)
             if scen._mpisppy_probability == "uniform":
                 bprob = "uniform"
             else:
-                bprob = bundle._mpisppy_probability 
+                raise RuntimeError("Proper bundles created by proper_bundle.py require uniform probability (consider creating problem-specific bundles)")
+                bprob = bundle._mpisppy_probability
             sputils.attach_root_node(bundle, 0, nonantlist)
             bundle._mpisppy_probability = bprob
+
+            if len(scen._mpisppy_node_list) > 1 and self.bunBFs is None:
+                raise RuntimeError("You are creating proper bundles for a\n"
+                      "multi-stage problem, but without cfg.branching_factors.\n"
+                      "We need branching factors and all bundles must cover\n"
+                      "the same number of entire second stage nodes.\n"
+                      )
             return bundle
         else:
             raise RuntimeError (f"Scenario name does not have scen or Bundle: {sname}")
-
-
-
-
-