1010# TBD: we should consider restructuring this and moving the capability to spbase
1111
1212import os
13+ import numpy as np
1314import mpisppy .utils .sputils as sputils
1415import mpisppy .utils .pickle_bundle as pickle_bundle
1516
1920# - read a bundle
2021# - make a bundle
2122# - make a bundle and write it
22- # Multi-stage (as of Sept 2024) not supported in class or generic_cylinders
23- # You need to do something clever like in aircondB
23+ # Multi-stage
24+ # Is not very special because all we do is make sure that
25+ # bundles cover entire second-stage nodes so the new bundled
26+ # problem is a two-stage problem no matter how many stages
27+ # were in the original problem. As of Dec 2024, support
28+ # is provided only when there are branching factors.
2429# NOTE:: the caller needs to make sure it is two stage
2530# the caller needs to worry about what is in what rank
2631# (local_scenarios might have bundle names, e.g.)
@@ -50,14 +55,28 @@ def scenario_names_creator(self, num_scens, start=None, cfg=None):
5055 return cfg .model .scenario_names_creator (num_scens , start = start )
5156
5257 def bundle_names_creator (self , num_buns , start = None , cfg = None ):
58+
59+ def _multistage_check (bunsize ):
60+ # returns bunBFs as a side-effect
61+ BFs = cfg .branching_factors
62+ beyond2size = np .prod (BFs [1 :])
63+ if bunsize % beyond2size != 0 :
64+ raise RuntimeError (f"Bundles must consume the same number of entire second stage nodes: { beyond2size = } { bunsize = } )
65+ # we need bunBFs for EF formulation
66+ self .bunBFs = [bunsize // beyond2size ] + BFs [1 :]
67+
5368 # start refers to the bundle number; bundles are always zero-based
5469 if start is None :
5570 start = 0
5671 assert cfg is not None , "ProperBundler needs cfg for bundle names"
5772 assert cfg .get ("num_scens" ) is not None
5873 assert cfg .get ("scenarios_per_bundle" ) is not None
59- assert cfg .num_scens % cfg .scenarios_per_bundle == 0
74+ assert cfg .num_scens % cfg .scenarios_per_bundle == 0 , "Bundles must consume the same number of entire second stage nodes: {cfg.num_scens=} {bunsize=}"
6075 bsize = cfg .scenarios_per_bundle # typing aid
76+ if cfg .get ("branching_factors" ) is not None :
77+ _multistage_check (bsize )
78+ else :
79+ self .bunBFs = None
6180 # We need to know if scenarios (not bundles) are one-based.
6281 inum = sputils .extract_num (self .module .scenario_names_creator (1 )[0 ])
6382 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):
8099 cfg = kwargs ["cfg" ]
81100 if "scen" in sname or "Scen" in sname :
82101 # In case the user passes in kwargs from scenario_creator_kwargs.
83- return self .module .scenario_creator (sname , {** self .original_kwargs , ** kwargs })
102+ return self .module .scenario_creator (sname , ** {** self .original_kwargs , ** kwargs })
84103
85104 elif "Bundle" in sname and cfg .get ("unpickle_bundles_dir" ) is not None :
86105 fname = os .path .join (cfg .unpickle_bundles_dir , sname + ".pkl" )
87106 bundle = pickle_bundle .dill_unpickle (fname )
88107 return bundle
89108 elif "Bundle" in sname and cfg .get ("unpickle_bundles_dir" ) is None :
90- # this is also the branch for proper_no_files
91109 # If we are still here, we have to create the bundle.
92110 firstnum = int (sname .split ("_" )[1 ]) # sname is a bundle name
93111 lastnum = int (sname .split ("_" )[2 ])
94112 # snames are scenario names
95113 snames = self .module .scenario_names_creator (lastnum - firstnum + 1 ,
96114 firstnum )
115+ kws = self .original_kwargs
116+ if self .bunBFs is not None :
117+ # The original scenario creator needs to handle these
118+ kws ["branching_factors" ] = self .bunBFs
119+
97120 # We are assuming seeds are managed by the *scenario* creator.
98121 bundle = sputils .create_EF (snames , self .module .scenario_creator ,
99- scenario_creator_kwargs = self . original_kwargs ,
122+ scenario_creator_kwargs = kws ,
100123 EF_name = sname ,
101124 suppress_warnings = True ,
102125 nonant_for_fixed_vars = False )
103126
104127 nonantlist = [v for idx , v in bundle .ref_vars .items () if idx [0 ] == "ROOT" ]
105128 # if the original scenarios were uniform, this needs to be also
106129 # (EF formation will recompute for the bundle if uniform)
130+ # Get an arbitrary scenario.
107131 scen = self .module .scenario_creator (snames [0 ], ** self .original_kwargs )
108132 if scen ._mpisppy_probability == "uniform" :
109133 bprob = "uniform"
110134 else :
111- bprob = bundle ._mpisppy_probability
135+ raise RuntimeError ("Proper bundles created by proper_bundle.py require uniform probability (consider creating problem-specific bundles)" )
136+ bprob = bundle ._mpisppy_probability
112137 sputils .attach_root_node (bundle , 0 , nonantlist )
113138 bundle ._mpisppy_probability = bprob
139+
140+ if len (scen ._mpisppy_node_list ) > 1 and self .bunBFs is None :
141+ raise RuntimeError ("You are creating proper bundles for a\n "
142+ "multi-stage problem, but without cfg.branching_factors.\n "
143+ "We need branching factors and all bundles must cover\n "
144+ "the same number of entire second stage nodes.\n "
145+ )
114146 return bundle
115147 else :
116148 raise RuntimeError (f"Scenario name does not have scen or Bundle: { sname } )
117-
118-
119-
120-
121-
0 commit comments