def miditer(self):
self.iter_since_last_check += 1
ib = self.opt.spcomm.BestInnerBound
if ib != self.cur_ib:
self.cur_ib = ib
self.iter_at_cur_ib = 1
elif self.cur_ib is not None and math.isfinite(self.cur_ib):
self.iter_at_cur_ib += 1
ob = self.opt.spcomm.BestOuterBound
if self.cur_ob is not None and math.isclose(ob, self.cur_ob):
ob_new = False
else:
self.cur_ob = ob
ob_new = True
if not self.any_cuts:
if self.opt.spcomm.new_cuts:
self.any_cuts = True
## if its the second time or more with this IB, we'll only check
## if the last improved the OB, or if the OB is new itself (from somewhere else)
check = (self.check_bound_iterations is not None) and self.any_cuts and ( \
(self.iter_at_cur_ib == self.check_bound_iterations) or \
(self.iter_at_cur_ib > self.check_bound_iterations and ob_new) or \
((self.iter_since_last_check%self.check_bound_iterations == 0) and self.opt.spcomm.new_cuts))
# if there hasn't been OB movement, check every so often if we have new cuts
if check:
global_toc(f"Attempting to update Best Bound with CrossScenarioExtension")
self._check_bound()
self.opt.spcomm.new_cuts = False
self.iter_since_last_check = 0
def is_converged(self):
## might as well get a bound, in this case
if self.opt._PHIter == 1:
self.BestOuterBound = self.OuterBoundUpdate(self.opt.trivial_bound)
if not self.has_innerbound_spokes:
if self.opt._PHIter == 1:
logger.warning("PHHub cannot compute convergence without "
"inner bound spokes.")
## you still want to output status, even without inner bounders configured
if self.global_rank == 0:
self.screen_trace()
return False
if not self.has_outerbound_spokes:
if self.opt._PHIter == 1:
global_toc("Without outer bound spokes, no progress "
"will be made on the Best Bound")
## log some output
if self.global_rank == 0:
self.screen_trace()
return self.determine_termination()
def hub_finalize(self):
if self.has_outerbound_spokes:
self.receive_outerbounds()
if self.has_innerbound_spokes:
self.receive_innerbounds()
if self.global_rank == 0:
self.print_init = True
global_toc(f"Statistics at termination", True)
self.screen_trace()
def _determine_innerbound_winner(self):
if self.spcomm.global_rank == 0:
if self.spcomm.last_ib_idx is None:
best_strata_rank = -1
global_toc("No incumbent solution available to write!")
else:
best_strata_rank = self.spcomm.last_ib_idx
else:
best_strata_rank = None
best_strata_rank = self.spcomm.fullcomm.bcast(best_strata_rank, root=0)
return (self.spcomm.strata_rank == best_strata_rank)
def screen_trace(self):
current_iteration = self.current_iteration()
abs_gap, rel_gap = self.compute_gaps()
best_solution = self.BestInnerBound
best_bound = self.BestOuterBound
update_source = self.get_update_string()
if self.print_init:
row = f'{"Iter.":>5s} {" "} {"Best Bound":>14s} {"Best Incumbent":>14s} {"Rel. Gap":>12s} {"Abs. Gap":>14s}'
global_toc(row, True)
self.print_init = False
row = f"{current_iteration:5d} {update_source} {best_bound:14.4f} {best_solution:14.4f} {rel_gap*100:12.3f}% {abs_gap:14.4f}"
global_toc(row, True)
self.clear_latest_chars()
def determine_termination(self):
abs_gap_satisfied = False
rel_gap_satisfied = False
if hasattr(self,"options") and self.options is not None:
if "rel_gap" in self.options:
rel_gap = self.compute_gap(compute_relative=True)
rel_gap_satisfied = rel_gap <= self.options["rel_gap"]
if "abs_gap" in self.options:
abs_gap = self.compute_gap(compute_relative=False)
abs_gap_satisfied = abs_gap <= self.options["abs_gap"]
if abs_gap_satisfied:
global_toc(f"Terminating based on inter-cylinder absolute gap {abs_gap:12.4f}")
if rel_gap_satisfied:
global_toc(f"Terminating based on inter-cylinder relative gap {rel_gap*100:12.3f}%")
return abs_gap_satisfied or rel_gap_satisfied
def __init__(
self,
options,
all_scenario_names,
scenario_creator,
scenario_denouement=None,
all_nodenames=None,
mpicomm=None,
scenario_creator_kwargs=None,
extensions=None,
extension_kwargs=None,
PH_converger=None,
rho_setter=None,
variable_probability=None,
):
""" PHBase constructor. """
super().__init__(
options,
all_scenario_names,
scenario_creator,
scenario_denouement=scenario_denouement,
all_nodenames=all_nodenames,
mpicomm=mpicomm,
extensions=extensions,
extension_kwargs=extension_kwargs,
scenario_creator_kwargs=scenario_creator_kwargs,
variable_probability=variable_probability,
)
global_toc("Initializing PHBase")
# Note that options can be manipulated from outside on-the-fly.
# self.options (from super) will archive the original options.
self.options = options
self.options_check()
self.PH_converger = PH_converger
self.rho_setter = rho_setter
self.iter0_solver_options = options["iter0_solver_options"]
self.iterk_solver_options = options["iterk_solver_options"]
self.current_solver_options = self.iter0_solver_options
# flags to complete the invariant
self.convobject = None # PH converger
self.attach_xbars()
def determine_termination(self):
# return True if termination is indicated, otherwise return False
if not hasattr(self,"options") or self.options is None\
or ("rel_gap" not in self.options and "abs_gap" not in self.options\
and "max_stalled_iters" not in self.options):
return False # Nothing to see here folks...
# If we are still here, there is some option for termination
abs_gap, rel_gap = self.compute_gaps()
abs_gap_satisfied = False
rel_gap_satisfied = False
max_stalled_satisfied = False
if "rel_gap" in self.options and rel_gap <= self.options["rel_gap"]:
rel_gap_satisfied = True
if "abs_gap" in self.options and abs_gap <= self.options["abs_gap"]:
rel_gap_satisfied = True
if "max_stalled_iters" in self.options:
if abs_gap < self.last_gap: # liberal test (we could use an epsilon)
self.last_gap = abs_gap
self.stalled_iter_cnt = 0
else:
self.stalled_iter_cnt += 1
if self.stalled_iter_cnt >= self.options["max_stalled_iters"]:
max_stalled_satisfied = True
if abs_gap_satisfied:
global_toc(f"Terminating based on inter-cylinder absolute gap {abs_gap:12.4f}")
if rel_gap_satisfied:
global_toc(f"Terminating based on inter-cylinder relative gap {rel_gap*100:12.3f}%")
if max_stalled_satisfied:
global_toc(f"Terminating based on max-stalled-iters {self.stalled_iter_cnt}")
return abs_gap_satisfied or rel_gap_satisfied or max_stalled_satisfied
def iterk_loop(self):
""" Perform all PH iterations after iteration 0.
This function terminates if any of the following occur:
1. The maximum number of iterations is reached.
2. The user specifies a converger, and the `is_converged()` method of
that converger returns True.
3. The hub tells it to terminate.
4. The user does not specify a converger, and the default convergence
criteria are met (i.e. the convergence value falls below the
user-specified threshold).
Args: None
"""
verbose = self.options["verbose"]
have_extensions = self.extensions is not None
have_converger = self.PH_converger is not None
dprogress = self.options["display_progress"]
dtiming = self.options["display_timing"]
dconvergence_detail = self.options["display_convergence_detail"]
self.conv = None
max_iterations = int(self.options["PHIterLimit"])
for self._PHIter in range(1, max_iterations + 1):
iteration_start_time = time.time()
if dprogress:
global_toc(f"\nInitiating PH Iteration {self._PHIter}\n",
self.cylinder_rank == 0)
# Compute xbar
#global_toc('Rank: {} - Before Compute_Xbar'.format(self.cylinder_rank), True)
self.Compute_Xbar(verbose)
#global_toc('Rank: {} - After Compute_Xbar'.format(self.cylinder_rank), True)
# update the weights
self.Update_W(verbose)
#global_toc('Rank: {} - After Update_W'.format(self.cylinder_rank), True)
self.conv = self.convergence_diff()
#global_toc('Rank: {} - After convergence_diff'.format(self.cylinder_rank), True)
if have_extensions:
self.extobject.miditer()
# The hub object takes precedence
# over the converger, such that
# the spokes will always have the
# latest data, even at termination
if self.spcomm is not None:
self.spcomm.sync()
if self.spcomm.is_converged():
global_toc("Cylinder convergence", self.cylinder_rank == 0)
break
if have_converger:
if self.convobject.is_converged():
converged = True
global_toc(
"User-supplied converger determined termination criterion reached",
self.cylinder_rank == 0)
break
elif self.conv is not None:
if self.conv < self.options["convthresh"]:
converged = True
global_toc(
"Convergence metric=%f dropped below user-supplied threshold=%f"
% (self.conv, self.options["convthresh"]),
self.cylinder_rank == 0)
break
teeme = ("tee-rank0-solves" in self.options
and self.options["tee-rank0-solves"]
and self.cylinder_rank == 0)
self.solve_loop(solver_options=self.current_solver_options,
dtiming=dtiming,
gripe=True,
disable_pyomo_signal_handling=False,
tee=teeme,
verbose=verbose)
if have_extensions:
self.extobject.enditer()
if dprogress and self.cylinder_rank == 0:
print("")
print("After PH Iteration", self._PHIter)
print("Scaled PHBase Convergence Metric=", self.conv)
print("Iteration time: %6.2f" %
(time.time() - iteration_start_time))
print("Elapsed time: %6.2f" %
(time.perf_counter() - self.start_time))
if dconvergence_detail:
self.report_var_values_at_rank0(header="Convergence detail:")
else: # no break, (self._PHIter == max_iterations)
# NOTE: If we return for any other reason things are reasonably in-sync.
# due to the convergence check. However, here we return we'll be
# out-of-sync because of the solve_loop could take vasty different
# times on different threads. This can especially mess up finalization.
# As a guard, we'll put a barrier here.
self.mpicomm.Barrier()
global_toc(
"Reached user-specified limit=%d on number of PH iterations" %
max_iterations, self.cylinder_rank == 0)
def run(self, comm_world=None):
""" top level for the hub and spoke system
Args:
comm_world (MPI comm): the world for this hub-spoke system
"""
if self._ran:
raise RuntimeError("WheelSpinner can only be run once")
hub_dict = self.hub_dict
list_of_spoke_dict = self.list_of_spoke_dict
# Confirm that the provided dictionaries specifying
# the hubs and spokes contain the appropriate keys
if "hub_class" not in hub_dict:
raise RuntimeError(
"The hub_dict must contain a 'hub_class' key specifying "
"the hub class to use")
if "opt_class" not in hub_dict:
raise RuntimeError(
"The hub_dict must contain an 'opt_class' key specifying "
"the SPBase class to use (e.g. PHBase, etc.)")
if "hub_kwargs" not in hub_dict:
hub_dict["hub_kwargs"] = dict()
if "opt_kwargs" not in hub_dict:
hub_dict["opt_kwargs"] = dict()
for spoke_dict in list_of_spoke_dict:
if "spoke_class" not in spoke_dict:
raise RuntimeError(
"Each spoke_dict must contain a 'spoke_class' key "
"specifying the spoke class to use")
if "opt_class" not in spoke_dict:
raise RuntimeError(
"Each spoke_dict must contain an 'opt_class' key "
"specifying the SPBase class to use (e.g. PHBase, etc.)")
if "spoke_kwargs" not in spoke_dict:
spoke_dict["spoke_kwargs"] = dict()
if "opt_kwargs" not in spoke_dict:
spoke_dict["opt_kwargs"] = dict()
if comm_world is None:
comm_world = MPI.COMM_WORLD
n_spokes = len(list_of_spoke_dict)
# Create the necessary communicators
fullcomm = comm_world
strata_comm, cylinder_comm = _make_comms(n_spokes, fullcomm=fullcomm)
strata_rank = strata_comm.Get_rank()
cylinder_rank = cylinder_comm.Get_rank()
global_rank = fullcomm.Get_rank()
# Assign hub/spokes to individual ranks
if strata_rank == 0: # This rank is a hub
sp_class = hub_dict["hub_class"]
sp_kwargs = hub_dict["hub_kwargs"]
opt_class = hub_dict["opt_class"]
opt_kwargs = hub_dict["opt_kwargs"]
opt_dict = hub_dict
else: # This rank is a spoke
spoke_dict = list_of_spoke_dict[strata_rank - 1]
sp_class = spoke_dict["spoke_class"]
sp_kwargs = spoke_dict["spoke_kwargs"]
opt_class = spoke_dict["opt_class"]
opt_kwargs = spoke_dict["opt_kwargs"]
opt_dict = spoke_dict
# Create the appropriate opt object locally
opt_kwargs["mpicomm"] = cylinder_comm
opt = opt_class(**opt_kwargs)
# Create the SPCommunicator object (hub/spoke) with
# the appropriate SPBase object attached
if strata_rank == 0: # Hub
spcomm = sp_class(opt, fullcomm, strata_comm, cylinder_comm,
list_of_spoke_dict, **sp_kwargs)
else: # Spokes
spcomm = sp_class(opt, fullcomm, strata_comm, cylinder_comm,
**sp_kwargs)
# Create the windows, run main(), destroy the windows
spcomm.make_windows()
if strata_rank == 0:
spcomm.setup_hub()
global_toc("Starting spcomm.main()")
spcomm.main()
if strata_rank == 0: # If this is the hub
spcomm.send_terminate()
# Anything that's left to do
spcomm.finalize()
# to ensure the messages below are True
cylinder_comm.Barrier()
global_toc(
f"Hub algorithm {opt_class.__name__} complete, waiting for spoke finalization"
)
global_toc(f"Spoke {sp_class.__name__} finalized",
(cylinder_rank == 0 and strata_rank != 0))
fullcomm.Barrier()
## give the hub the chance to catch new values
spcomm.hub_finalize()
spcomm.free_windows()
global_toc("Windows freed")
self.spcomm = spcomm
self.opt_dict = opt_dict
self.global_rank = global_rank
self.strata_rank = strata_rank
self.cylinder_rank = cylinder_rank
if self.strata_rank == 0:
self.BestInnerBound = spcomm.BestInnerBound
self.BestOuterBound = spcomm.BestOuterBound
else: # the cylinder ranks don't track the inner / outer bounds
self.BestInnerBound = None
self.BestOuterBound = None
self._ran = True
def run(self, confidence_level=0.95, objective_gap=False):
# We get the MMW right term, then xhat, then the MMW left term.
#Compute the nonant xhat (the one used in the left term of MMW (9) ) using
# the first scenarios
############### get the parameters
start = self.start
scenario_denouement = self.refmodel.scenario_denouement
#Introducing batches otpions
num_batches = self.num_batches
bs = self.batch_size
batch_size = bs if (
bs is not None) else start #is None : take size_batch=num_scens
sample_options = self.options
#Some options are specific to 2-stage or multi-stage problems
if self.multistage:
sampling_BFs = ciutils.BFs_from_numscens(batch_size,
self.numstages)
#TODO: Change this to get a more logical way to compute BFs
batch_size = np.prod(sampling_BFs)
else:
sampling_BFs = None
sample_options['num_scens'] = batch_size
sample_options['_mpisppy_probability'] = 1 / batch_size
scenario_creator_kwargs = self.refmodel.kw_creator(sample_options)
sample_scen_creator = self.refmodel.scenario_creator
#Solver settings
solvername = self.options['EF_solver_name']
solver_options = self.options[
'EF_solver_options'] if 'EF_solver_options' in self.options else None
solver_options = remove_None(solver_options)
#Now we compute for each batch the whole Gn term from MMW (9)
G = np.zeros(num_batches) #the Gbar of MMW (10)
#we will compute the mean via a loop (to be parallelized ?)
zhats = [] #evaluation of xhat at each scenario
for i in range(num_batches):
scenstart = None if self.multistage else start
gap_options = {
'seed': start,
'BFs': sampling_BFs
} if self.multistage else None
scenario_names = self.refmodel.scenario_names_creator(
batch_size, start=scenstart)
estim = ciutils.gap_estimators(
self.xhat_one,
self.refmodelname,
solving_type=self.type,
scenario_names=scenario_names,
sample_options=gap_options,
ArRP=1,
scenario_creator_kwargs=scenario_creator_kwargs,
scenario_denouement=scenario_denouement,
solvername=solvername,
solver_options=solver_options,
objective_gap=objective_gap)
Gn = estim['G']
start = estim['seed']
# collect evaluation of xhat at all scenario
if objective_gap:
for zhat in estim['zhats']:
zhats.append(zhat)
if (self.verbose):
global_toc(
f"Gn={Gn} for the batch {i}") # Left term of LHS of (9)
G[i] = Gn
s_g = np.std(G) #Standard deviation of gap
Gbar = np.mean(G)
t_g = scipy.stats.t.ppf(confidence_level, num_batches - 1)
epsilon_g = t_g * s_g / np.sqrt(num_batches)
gap_inner_bound = Gbar + epsilon_g
gap_outer_bound = 0
if objective_gap == True:
zhat_bar = np.mean(zhats)
s_zhat = np.std(zhats) # Stanard deviation of objectives at xhat
t_zhat = scipy.stats.t.ppf(confidence_level, len(zhats) - 1)
epsilon_zhat = t_zhat * s_zhat / np.sqrt(len(zhats))
gap_inner_bound += zhat_bar + epsilon_zhat
gap_outer_bound += zhat_bar - epsilon_zhat
self.result = {
"gap_inner_bound": gap_inner_bound,
"gap_outer_bound": gap_outer_bound,
"Gbar": Gbar,
"std": s_g,
"Glist": G
}
if objective_gap:
self.result["zhat_bar"] = zhat_bar
self.result["std_zhat"] = s_zhat
return (self.result)
def run(self, maxit=200):
# Do everything as specified by the options (maxit is provided as a safety net).
refmodel = self.refmodel
mult = self.sample_size_ratio # used to set m_k= mult*n_k
scenario_denouement = refmodel.scenario_denouement if hasattr(
refmodel, "scenario_denouement") else None
#----------------------------Step 0 -------------------------------------#
#Initialization
k = 1
if self.stopping_criterion == "BM":
#Finding a constant used to compute nk
r = 2 #TODO : we could add flexibility here
j = np.arange(1, 1000)
if self.q is None:
s = sum(np.power(j, -self.p * np.log(j)))
else:
if self.q < 1:
raise RuntimeError("Parameter q should be greater than 1.")
s = sum(np.exp(-self.p * np.power(j, 2 * self.q / r)))
self.c = max(
1, 2 * np.log(s / (np.sqrt(2 * np.pi) *
(1 - self.confidence_level))))
lower_bound_k = self.sample_size(k, None, None, None)
#Computing xhat_1.
#We use sample_size_ratio*n_k observations to compute xhat_k
xhat_branching_factors = ciutils.scalable_branching_factors(
mult * lower_bound_k, self.options['branching_factors'])
mk = np.prod(xhat_branching_factors)
self.xhat_gen_options[
'start_seed'] = self.SeedCount #TODO: Maybe find a better way to manage seed
xhat_scenario_names = refmodel.scenario_names_creator(mk)
xgo = self.xhat_gen_options.copy()
xgo.pop("solver_options", None) # it will be given explicitly
xgo.pop("scenario_names", None) # it will be given explicitly
xhat_k = self.xhat_generator(xhat_scenario_names,
solver_options=self.solver_options,
**xgo)
self.SeedCount += sputils.number_of_nodes(xhat_branching_factors)
#----------------------------Step 1 -------------------------------------#
#Computing n_1 and associated scenario names
nk = np.prod(
ciutils.scalable_branching_factors(
lower_bound_k, self.options['branching_factors'])
) #To ensure the same growth that in the one-tree seqsampling
estimator_scenario_names = refmodel.scenario_names_creator(nk)
#Computing G_nk and s_k associated with xhat_1
Gk, sk = self._gap_estimators_with_independent_scenarios(
xhat_k, nk, estimator_scenario_names, scenario_denouement)
#----------------------------Step 2 -------------------------------------#
while (self.stop_criterion(Gk, sk, nk) and k < maxit):
#----------------------------Step 3 -------------------------------------#
k += 1
nk_m1 = nk #n_{k-1}
mk_m1 = mk
lower_bound_k = self.sample_size(k, Gk, sk, nk_m1)
#Computing m_k and associated scenario names
xhat_branching_factors = ciutils.scalable_branching_factors(
mult * lower_bound_k, self.options['branching_factors'])
mk = np.prod(xhat_branching_factors)
self.xhat_gen_options[
'start_seed'] = self.SeedCount #TODO: Maybe find a better way to manage seed
xhat_scenario_names = refmodel.scenario_names_creator(mk)
#Computing xhat_k
xgo = self.xhat_gen_options.copy()
xgo.pop("solver_options", None) # it will be given explicitly
xgo.pop("scenario_names", None) # it will be given explicitly
xhat_k = self.xhat_generator(xhat_scenario_names,
solver_options=self.solver_options,
**xgo)
#Computing n_k and associated scenario names
self.SeedCount += sputils.number_of_nodes(xhat_branching_factors)
nk = np.prod(
ciutils.scalable_branching_factors(
lower_bound_k, self.options['branching_factors'])
) #To ensure the same growth that in the one-tree seqsampling
nk += self.batch_size - nk % self.batch_size
estimator_scenario_names = refmodel.scenario_names_creator(nk)
Gk, sk = self._gap_estimators_with_independent_scenarios(
xhat_k, nk, estimator_scenario_names, scenario_denouement)
if (k % 10 == 0):
print(f"k={k}")
print(f"n_k={nk}")
#----------------------------Step 4 -------------------------------------#
if (k == maxit):
raise RuntimeError(
f"The loop terminated after {maxit} iteration with no acceptable solution"
)
T = k
final_xhat = xhat_k
if self.stopping_criterion == "BM":
upper_bound = self.h * sk + self.eps
elif self.stopping_criterion == "BPL":
upper_bound = self.eps
else:
raise RuntimeError("Only BM and BPL criterion are supported yet.")
CI = [0, upper_bound]
global_toc(f"G={Gk}")
global_toc(f"s={sk}")
global_toc(f"xhat has been computed with {nk*mult} observations.")
return {
"T": T,
"Candidate_solution": final_xhat,
"CI": CI,
}
def run(self,maxit=200):
refmodel = self.refmodel
mult = self.sample_size_ratio # used to set m_k= mult*n_k
#----------------------------Step 0 -------------------------------------#
#Initialization
k =1
#Computing the lower bound for n_1
if self.stopping_criterion == "BM":
#Finding a constant used to compute nk
r = 2 #TODO : we could add flexibility here
j = np.arange(1,1000)
if self.q is None:
s = sum(np.power(j,-self.p*np.log(j)))
else:
if self.q<1:
raise RuntimeError("Parameter q should be greater than 1.")
s = sum(np.exp(-self.p*np.power(j,2*self.q/r)))
self.c = max(1,2*np.log(s/(np.sqrt(2*np.pi)*(1-self.confidence_level))))
lower_bound_k = self.sample_size(k, None, None, None)
#Computing xhat_1.
#We use sample_size_ratio*n_k observations to compute xhat_k
if self.multistage:
xhat_BFs = ciutils.scalable_BFs(mult*lower_bound_k, self.options['BFs'])
mk = np.prod(xhat_BFs)
self.xhat_gen_options['start_seed'] = self.SeedCount #TODO: Maybe find a better way to manage seed
xhat_scenario_names = refmodel.scenario_names_creator(mk)
else:
mk = int(np.floor(mult*lower_bound_k))
xhat_scenario_names = refmodel.scenario_names_creator(mk, start=self.ScenCount)
self.ScenCount+=mk
xhat_k = self.xhat_generator(xhat_scenario_names,
solvername=self.solvername,
solver_options=self.solver_options,
**self.xhat_gen_options)
#----------------------------Step 1 -------------------------------------#
#Computing n_1 and associated scenario names
if self.multistage:
self.SeedCount += sputils.number_of_nodes(xhat_BFs)
gap_BFs = ciutils.scalable_BFs(lower_bound_k, self.options['BFs'])
nk = np.prod(gap_BFs)
estimator_scenario_names = refmodel.scenario_names_creator(nk)
sample_options = {'BFs':gap_BFs, 'seed':self.SeedCount}
else:
nk = self.ArRP *int(np.ceil(lower_bound_k/self.ArRP))
estimator_scenario_names = refmodel.scenario_names_creator(nk,
start=self.ScenCount)
sample_options = None
self.ScenCount+= nk
#Computing G_nkand s_k associated with xhat_1
self.options['num_scens'] = nk
scenario_creator_kwargs = self.refmodel.kw_creator(self.options)
scenario_denouement = refmodel.scenario_denouement if hasattr(refmodel, "scenario_denouement") else None
estim = ciutils.gap_estimators(xhat_k, self.refmodelname,
solving_type=self.solving_type,
scenario_names=estimator_scenario_names,
sample_options=sample_options,
ArRP=self.ArRP,
scenario_creator_kwargs=scenario_creator_kwargs,
scenario_denouement=scenario_denouement,
solvername=self.solvername,
solver_options=self.solver_options)
Gk,sk = estim['G'],estim['s']
if self.multistage:
self.SeedCount = estim['seed']
#----------------------------Step 2 -------------------------------------#
while( self.stop_criterion(Gk,sk,nk) and k<maxit):
#----------------------------Step 3 -------------------------------------#
k+=1
nk_m1 = nk #n_{k-1}
mk_m1 = mk
lower_bound_k = self.sample_size(k, Gk, sk, nk_m1)
#Computing m_k and associated scenario names
if self.multistage:
xhat_BFs = ciutils.scalable_BFs(mult*lower_bound_k, self.options['BFs'])
mk = np.prod(xhat_BFs)
self.xhat_gen_options['start_seed'] = self.SeedCount #TODO: Maybe find a better way to manage seed
xhat_scenario_names = refmodel.scenario_names_creator(mk)
else:
mk = int(np.floor(mult*lower_bound_k))
assert mk>= mk_m1, "Our sample size should be increasing"
if (k%self.kf_xhat==0):
#We use only new scenarios to compute xhat
xhat_scenario_names = refmodel.scenario_names_creator(mult*nk,
start=self.ScenCount)
self.ScenCount+= mk
else:
#We reuse the previous scenarios
xhat_scenario_names+= refmodel.scenario_names_creator(mult*(nk-nk_m1),
start=self.ScenCount)
self.ScenCount+= mk-mk_m1
#Computing xhat_k
xhat_k = self.xhat_generator(xhat_scenario_names,
solvername=self.solvername,
solver_options=self.solver_options,
**self.xhat_gen_options)
#Computing n_k and associated scenario names
if self.multistage:
self.SeedCount += sputils.number_of_nodes(xhat_BFs)
gap_BFs = ciutils.scalable_BFs(lower_bound_k, self.options['BFs'])
nk = np.prod(gap_BFs)
estimator_scenario_names = refmodel.scenario_names_creator(nk)
sample_options = {'BFs':gap_BFs, 'seed':self.SeedCount}
else:
nk = self.ArRP *int(np.ceil(lower_bound_k/self.ArRP))
assert nk>= nk_m1, "Our sample size should be increasing"
if (k%self.kf_Gs==0):
#We use only new scenarios to compute gap estimators
estimator_scenario_names = refmodel.scenario_names_creator(nk,
start=self.ScenCount)
self.ScenCount+=nk
else:
#We reuse the previous scenarios
estimator_scenario_names+= refmodel.scenario_names_creator((nk-nk_m1),
start=self.ScenCount)
self.ScenCount+= (nk-nk_m1)
sample_options = None
#Computing G_k and s_k
self.options['num_scens'] = nk
scenario_creator_kwargs = self.refmodel.kw_creator(self.options)
estim = ciutils.gap_estimators(xhat_k, self.refmodelname,
solving_type=self.solving_type,
scenario_names=estimator_scenario_names,
sample_options=sample_options,
ArRP=self.ArRP,
scenario_creator_kwargs=scenario_creator_kwargs,
scenario_denouement=scenario_denouement,
solvername=self.solvername,
solver_options=self.solver_options)
if self.multistage:
self.SeedCount = estim['seed']
Gk,sk = estim['G'],estim['s']
if (k%10==0) and global_rank==0:
print(f"k={k}")
print(f"n_k={nk}")
print(f"G_k={Gk}")
print(f"s_k={sk}")
#----------------------------Step 4 -------------------------------------#
if (k==maxit) :
raise RuntimeError(f"The loop terminated after {maxit} iteration with no acceptable solution")
T = k
final_xhat=xhat_k
if self.stopping_criterion == "BM":
upper_bound=self.h*sk+self.eps
elif self.stopping_criterion == "BPL":
upper_bound = self.eps
else:
raise RuntimeError("Only BM and BPL criterion are supported yet.")
CI=[0,upper_bound]
global_toc(f"G={Gk}")
global_toc(f"s={sk}")
global_toc(f"xhat has been computed with {nk*mult} observations.")
return {"T":T,"Candidate_solution":final_xhat,"CI":CI,}
def gap_estimators(xhat_one,
mname,
solving_type="EF-2stage",
scenario_names=None,
sample_options=None,
ArRP=1,
scenario_creator_kwargs={},
scenario_denouement=None,
solvername=None,
solver_options=None,
verbose=False,
objective_gap=False):
''' Given a xhat, scenario names, a scenario creator and options,
gap_estimators creates a scenario tree and the associatd estimators
G and s from §2 of [bm2011].
Returns G and s evaluated at xhat.
If ArRP>1, G and s are pooled, from a number ArRP of estimators,
computed with different scenario trees.
Parameters
----------
xhat_one : dict
A candidate first stage solution
mname: str
Name of the reference model, e.g. 'mpisppy.tests.examples.farmer'.
solving_type: str, optional
The way we solve the approximate problem. Can be "EF-2stage" (default)
or "EF-mstage".
scenario_names: list, optional
List of scenario names used to compute G_n and s_n. Default is None
Must be specified for 2 stage, but can be missing for multistage
sample_options: dict, optional
Only for multistage. Must contain a 'seed' and a 'branching_factors' attribute,
specifying the starting seed and the branching factors
of the scenario tree
ArRP:int,optional
Number of batches (we create a ArRP model). Default is 1 (one batch).
scenario_creator_kwargs: dict, optional
Additional arguments for scenario_creator. Default is {}
scenario_denouement: function, optional
Function to run after scenario creation. Default is None.
solvername : str, optional
Solver. Default is None
solver_options: dict, optional
Solving options. Default is None
verbose: bool, optional
Should it print the gap estimator ? Default is True
objective_gap: bool, optional
Returns a gap estimate around approximate objective value
branching_factors: list, optional
Only for multistage. List of branching factors of the sample scenario tree.
Returns
-------
G_k and s_k, gap estimator and associated standard deviation estimator.
'''
global_toc("Enter gap_estimators")
if solving_type not in ["EF-2stage", "EF-mstage"]:
raise RuntimeError(
"Only EF solve for the approximate problem is supported yet.")
else:
is_multi = (solving_type == "EF-mstage")
if is_multi:
try:
branching_factors = sample_options['branching_factors']
start = sample_options['seed']
except (TypeError, KeyError, RuntimeError):
raise RuntimeError(
'For multistage problems, sample_options must be a dict with branching_factors and seed attributes.'
)
else:
start = sputils.extract_num(scenario_names[0])
if ArRP > 1: #Special case : ArRP, G and s are pooled from r>1 estimators.
if is_multi:
raise RuntimeError(
"Pooled estimators are not supported for multistage problems yet."
)
n = len(scenario_names)
if (n % ArRP != 0):
raise RuntimeWarning("You put as an input a number of scenarios"+\
f" which is not a mutliple of {ArRP}.")
n = n - n % ArRP
G = []
s = []
for k in range(ArRP):
scennames = scenario_names[k * (n // ArRP):(k + 1) * (n // ArRP)]
tmp = gap_estimators(
xhat_one,
mname,
solvername=solvername,
scenario_names=scennames,
ArRP=1,
scenario_creator_kwargs=scenario_creator_kwargs,
scenario_denouement=scenario_denouement,
solver_options=solver_options,
solving_type=solving_type)
G.append(tmp['G'])
s.append(tmp['s'])
global_toc(f"ArRP {k} of {ArRP}")
#Pooling
G = np.mean(G)
s = np.linalg.norm(s) / np.sqrt(n // ArRP)
return {"G": G, "s": s, "seed": start}
#A1RP
#We start by computing the optimal solution to the approximate problem induced by our scenarios
if is_multi:
#Sample a scenario tree: this is a subtree, but starting from stage 1
samp_tree = sample_tree.SampleSubtree(
mname,
xhats=[],
root_scen=None,
starting_stage=1,
branching_factors=branching_factors,
seed=start,
options=scenario_creator_kwargs,
solvername=solvername,
solver_options=solver_options)
samp_tree.run()
start += sputils.number_of_nodes(branching_factors)
ama_object = samp_tree.ama
else:
#We use amalgamator to do it
ama_options = dict(scenario_creator_kwargs)
ama_options['start'] = start
ama_options['num_scens'] = len(scenario_names)
ama_options['EF_solver_name'] = solvername
ama_options['EF_solver_options'] = solver_options
ama_options[solving_type] = True
ama_object = ama.from_module(mname,
ama_options,
use_command_line=False)
ama_object.scenario_names = scenario_names
ama_object.verbose = False
ama_object.run()
start += len(scenario_names)
#Optimal solution of the approximate problem
zstar = ama_object.best_outer_bound
#Associated policies
xstars = sputils.nonant_cache_from_ef(ama_object.ef)
#Then, we evaluate the fonction value induced by the scenario at xstar.
if is_multi:
# Find feasible policies (i.e. xhats) for every non-leaf nodes
if len(samp_tree.ef._ef_scenario_names) > 1:
local_scenarios = {
sname: getattr(samp_tree.ef, sname)
for sname in samp_tree.ef._ef_scenario_names
}
else:
local_scenarios = {
samp_tree.ef._ef_scenario_names[0]: samp_tree.ef
}
xhats, start = sample_tree.walking_tree_xhats(
mname,
local_scenarios,
xhat_one['ROOT'],
branching_factors,
start,
scenario_creator_kwargs,
solvername=solvername,
solver_options=solver_options)
#Compute then the average function value with this policy
scenario_creator_kwargs = samp_tree.ama.kwargs
all_nodenames = sputils.create_nodenames_from_branching_factors(
branching_factors)
else:
#In a 2 stage problem, the only non-leaf is the ROOT node
xhats = xhat_one
all_nodenames = None
xhat_eval_options = {
"iter0_solver_options": None,
"iterk_solver_options": None,
"display_timing": False,
"solvername": solvername,
"verbose": False,
"solver_options": solver_options
}
ev = xhat_eval.Xhat_Eval(xhat_eval_options,
scenario_names,
ama_object.scenario_creator,
scenario_denouement,
scenario_creator_kwargs=scenario_creator_kwargs,
all_nodenames=all_nodenames)
#Evaluating xhat and xstar and getting the value of the objective function
#for every (local) scenario
zhat = ev.evaluate(xhats)
objs_at_xhat = ev.objs_dict
zstar = ev.evaluate(xstars)
objs_at_xstar = ev.objs_dict
eval_scen_at_xhat = []
eval_scen_at_xstar = []
scen_probs = []
for k, s in ev.local_scenarios.items():
eval_scen_at_xhat.append(objs_at_xhat[k])
eval_scen_at_xstar.append(objs_at_xstar[k])
scen_probs.append(s._mpisppy_probability)
scen_gaps = np.array(eval_scen_at_xhat) - np.array(eval_scen_at_xstar)
local_gap = np.dot(scen_gaps, scen_probs)
local_ssq = np.dot(scen_gaps**2, scen_probs)
local_prob_sqnorm = np.linalg.norm(scen_probs)**2
local_obj_at_xhat = np.dot(eval_scen_at_xhat, scen_probs)
local_estim = np.array(
[local_gap, local_ssq, local_prob_sqnorm, local_obj_at_xhat])
global_estim = np.zeros(4)
ev.mpicomm.Allreduce(local_estim, global_estim, op=mpi.SUM)
G, ssq, prob_sqnorm, obj_at_xhat = global_estim
if global_rank == 0 and verbose:
print(f"G = {G}")
sample_var = (ssq - G**2) / (1 - prob_sqnorm) #Unbiased sample variance
s = np.sqrt(sample_var)
use_relative_error = (np.abs(zstar) > 1)
G = correcting_numeric(G,
objfct=obj_at_xhat,
relative_error=use_relative_error)
if objective_gap:
if is_multi:
return {
"G": G,
"s": s,
"zhats": [zhat],
"zstars": [zstar],
"seed": start
}
else:
return {
"G": G,
"s": s,
"zhats": eval_scen_at_xhat,
"zstars": eval_scen_at_xstar,
"seed": start
}
else:
return {"G": G, "s": s, "seed": start}
def __init__(
self,
options,
all_scenario_names,
scenario_creator,
scenario_denouement=None,
all_nodenames=None,
mpicomm=None,
scenario_creator_kwargs=None,
variable_probability=None,
E1_tolerance=1e-5
):
# TODO add missing and private attributes (JP)
# TODO add a class attribute called ROOTNODENAME = "ROOT"
# TODO? add decorators to the class attributes
self.start_time = time.perf_counter()
self.options = options
self.all_scenario_names = all_scenario_names
self.scenario_creator = scenario_creator
self.scenario_denouement = scenario_denouement
self.comms = dict()
self.local_scenarios = dict()
self.local_scenario_names = list()
self.E1_tolerance = E1_tolerance # probs must sum to almost 1
self.names_in_bundles = None
self.scenarios_constructed = False
if all_nodenames is None:
self.all_nodenames = ["ROOT"]
elif "ROOT" in all_nodenames:
self.all_nodenames = all_nodenames
else:
raise RuntimeError("'ROOT' must be in the list of node names")
self.variable_probability = variable_probability
self.multistage = (len(self.all_nodenames) > 1)
# Set up MPI communicator and rank
if mpicomm is not None:
self.mpicomm = mpicomm
else:
self.mpicomm = MPI.COMM_WORLD
self.cylinder_rank = self.mpicomm.Get_rank()
self.n_proc = self.mpicomm.Get_size()
self.global_rank = MPI.COMM_WORLD.Get_rank()
global_toc("Initializing SPBase")
if self.n_proc > len(self.all_scenario_names):
raise RuntimeError("More ranks than scenarios")
# Call various initialization methods
if "branching_factors" in self.options:
self.branching_factors = self.options["branching_factors"]
else:
self.branching_factors = [len(self.all_scenario_names)]
self._calculate_scenario_ranks()
if "bundles_per_rank" in self.options and self.options["bundles_per_rank"] > 0:
self._assign_bundles()
self.bundling = True
else:
self.bundling = False
self._create_scenarios(scenario_creator_kwargs)
self._look_and_leap()
self._compute_unconditional_node_probabilities()
self._attach_nlens()
self._attach_nonant_indices()
self._attach_varid_to_nonant_index()
self._create_communicators()
self._verify_nonant_lengths()
self._set_sense()
self._use_variable_probability_setter()
## SPCommunicator object
self._spcomm = None
if args.batch_size == None:
args.batch_size = args.num_scens
refmodel = modelpath #Change this path to use a different model
options = {"EF-2stage": True,# 2stage vs. mstage
"start": False,
"EF_solver_name": args.solver_name,
"EF_solver_options": solver_options,
"num_scens": args.num_scens} #Are the scenario shifted by a start arg ?
#should we accept these as arguments?
num_batches = args.num_batches
batch_size = args.batch_size
mmw = mmw_ci.MMWConfidenceIntervals(refmodel, options, xhat, num_batches, batch_size=batch_size,
verbose=True)
if args.alpha == None:
print('\nNo alpha given, defaulting to alpha = 0.95. To provide an alpha try:\n')
print('python -m mpisppy.confidence_intervals.mmw_conf {} {} {} {} --alpha 0.97 --MMW-num-batches {} --MMW-batch-size {} --num-scens {}\
\n'.format(sys.argv[0], args.instance, args.xhatpath, args.solver_name,
args.num_batches, args.batch_size, args.num_scens))
alpha = 0.95
else:
alpha = float(args.alpha)
r = mmw.run(confidence_level=alpha, objective_gap = args.objective_gap)
global_toc(r)
def Iter0(self):
""" Create solvers and perform the initial PH solve (with no dual
weights or prox terms).
This function quits() if the scenario probabilities do not sum to one,
or if any of the scenario subproblems are infeasible. It also calls the
`post_iter0` method of any extensions, and uses the rho setter (if
present) after the inital solve.
Returns:
float:
The so-called "trivial bound", i.e., the objective value of the
stochastic program with the nonanticipativity constraints
removed.
"""
if (self.extensions is not None):
self.extobject.pre_iter0()
verbose = self.options["verbose"]
dprogress = self.options["display_progress"]
dtiming = self.options["display_timing"]
dconvergence_detail = self.options["display_convergence_detail"]
have_extensions = self.extensions is not None
have_converger = self.PH_converger is not None
def _vb(msg):
if verbose and self.cylinder_rank == 0:
print("(rank0)", msg)
self._PHIter = 0
self._save_original_nonants()
global_toc("Creating solvers")
self._create_solvers()
teeme = ("tee-rank0-solves" in self.options
and self.options['tee-rank0-solves']
and self.cylinder_rank == 0)
if self.options["verbose"]:
print("About to call PH Iter0 solve loop on rank={}".format(
self.cylinder_rank))
global_toc("Entering solve loop in PHBase.Iter0")
self.solve_loop(solver_options=self.current_solver_options,
dtiming=dtiming,
gripe=True,
tee=teeme,
verbose=verbose)
if self.options["verbose"]:
print("PH Iter0 solve loop complete on rank={}".format(
self.cylinder_rank))
self._update_E1() # Apologies for doing this after the solves...
if (abs(1 - self.E1) > self.E1_tolerance):
if self.cylinder_rank == 0:
print("ERROR")
print("Total probability of scenarios was ", self.E1)
print("E1_tolerance = ", self.E1_tolerance)
quit()
feasP = self.feas_prob()
if feasP != self.E1:
if self.cylinder_rank == 0:
print("ERROR")
print("Infeasibility detected; E_feas, E1=", feasP, self.E1)
quit()
"""
with open('mpi.out-{}'.format(rank), 'w') as fd:
for sname in self.local_scenario_names:
fd.write('*** {} ***\n'.format(sname))
"""
#global_toc('Rank: {} - Building and solving models 0th iteration'.format(rank), True)
#global_toc('Rank: {} - assigning rho'.format(rank), True)
if have_extensions:
self.extobject.post_iter0()
if self.rho_setter is not None:
if self.cylinder_rank == 0:
self._use_rho_setter(verbose)
else:
self._use_rho_setter(False)
converged = False
if have_converger:
# Call the constructor of the converger object
self.convobject = self.PH_converger(self)
#global_toc('Rank: {} - Before iter loop'.format(self.cylinder_rank), True)
self.conv = None
self.trivial_bound = self.Ebound(verbose)
if dprogress and self.cylinder_rank == 0:
print("")
print("After PH Iteration", self._PHIter)
print("Trivial bound =", self.trivial_bound)
print("PHBase Convergence Metric =", self.conv)
print("Elapsed time: %6.2f" %
(time.perf_counter() - self.start_time))
if dconvergence_detail:
self.report_var_values_at_rank0(header="Convergence detail:")
self.reenable_W_and_prox()
self.current_solver_options = self.options["iterk_solver_options"]
return self.trivial_bound
def run(self):
""" Top-level execution."""
if self.is_EF:
ef = sputils.create_EF(
self.scenario_names,
self.scenario_creator,
scenario_creator_kwargs=self.kwargs,
suppress_warnings=True,
)
solvername = self.solvername
solver = pyo.SolverFactory(solvername)
if hasattr(self, "solver_options") and (self.solver_options
is not None):
for option_key, option_value in self.solver_options.items():
if option_value is not None:
solver.options[option_key] = option_value
if self.verbose:
global_toc("Starting EF solve")
if 'persistent' in solvername:
solver.set_instance(ef, symbolic_solver_labels=True)
results = solver.solve(tee=False)
else:
results = solver.solve(
ef,
tee=False,
symbolic_solver_labels=True,
)
if self.verbose:
global_toc("Completed EF solve")
self.EF_Obj = pyo.value(ef.EF_Obj)
objs = sputils.get_objs(ef)
self.is_minimizing = objs[0].is_minimizing
#TBD : Write a function doing this
if self.is_minimizing:
self.best_outer_bound = results.Problem[0]['Lower bound']
self.best_inner_bound = results.Problem[0]['Upper bound']
else:
self.best_inner_bound = results.Problem[0]['Upper bound']
self.best_outer_bound = results.Problem[0]['Lower bound']
self.ef = ef
if 'write_solution' in self.options:
if 'first_stage_solution' in self.options['write_solution']:
sputils.write_ef_first_stage_solution(
self.ef,
self.options['write_solution']['first_stage_solution'])
if 'tree_solution' in self.options['write_solution']:
sputils.write_ef_tree_solution(
self.ef,
self.options['write_solution']['tree_solution'])
self.xhats = sputils.nonant_cache_from_ef(ef)
self.local_xhats = self.xhats #Every scenario is local for EF
self.first_stage_solution = {"ROOT": self.xhats["ROOT"]}
else:
self.ef = None
args = argparse.Namespace(**self.options)
#Create a hub dict
hub_name = find_hub(self.options['cylinders'], self.is_multi)
hub_creator = getattr(vanilla, hub_name + '_hub')
beans = {
"args": args,
"scenario_creator": self.scenario_creator,
"scenario_denouement": self.scenario_denouement,
"all_scenario_names": self.scenario_names,
"scenario_creator_kwargs": self.kwargs
}
if self.is_multi:
beans["all_nodenames"] = self.options["all_nodenames"]
hub_dict = hub_creator(**beans)
#Add extensions
if 'extensions' in self.options:
for extension in self.options['extensions']:
extension_creator = getattr(vanilla, 'add_' + extension)
hub_dict = extension_creator(hub_dict, args)
#Create spoke dicts
potential_spokes = find_spokes(self.options['cylinders'],
self.is_multi)
#We only use the spokes with an associated command line arg set to True
spokes = [
spoke for spoke in potential_spokes
if self.options['with_' + spoke]
]
list_of_spoke_dict = list()
for spoke in spokes:
spoke_creator = getattr(vanilla, spoke + '_spoke')
spoke_beans = copy.deepcopy(beans)
if spoke == "xhatspecific":
spoke_beans["scenario_dict"] = self.options[
"scenario_dict"]
spoke_dict = spoke_creator(**spoke_beans)
list_of_spoke_dict.append(spoke_dict)
spcomm, opt_dict = sputils.spin_the_wheel(hub_dict,
list_of_spoke_dict)
self.opt = spcomm.opt
self.cylinder_rank = self.opt.cylinder_rank
self.on_hub = ("hub_class" in opt_dict)
if self.on_hub: # we are on a hub rank
self.best_inner_bound = spcomm.BestInnerBound
self.best_outer_bound = spcomm.BestOuterBound
#NOTE: We do not get bounds on every rank, only on hub
# This should change if we want to use cylinders for MMW
if 'write_solution' in self.options:
if 'first_stage_solution' in self.options['write_solution']:
sputils.write_spin_the_wheel_first_stage_solution(
spcomm, opt_dict,
self.options['write_solution']['first_stage_solution'])
if 'tree_solution' in self.options['write_solution']:
sputils.write_spin_the_wheel_tree_solution(
spcomm, opt_dict,
self.options['write_solution']['tree_solution'])
if self.on_hub: #we are on a hub rank
a_sname = self.opt.local_scenario_names[0]
root = self.opt.local_scenarios[a_sname]._mpisppy_node_list[0]
self.first_stage_solution = {
"ROOT":
[pyo.value(var) for var in root.nonant_vardata_list]
}
self.local_xhats = sputils.local_nonant_cache(spcomm)
def spin_the_wheel(hub_dict, list_of_spoke_dict, comm_world=None):
""" top level for the hub and spoke system
Args:
hub_dict(dict): controls hub creation
list_of_spoke_dict(list dict): controls creation of spokes
comm_world (MPI comm): the world for this hub-spoke system
Returns:
spcomm (Hub or Spoke object): the object that did the work (windowless)
opt_dict (dict): the dictionary that controlled creation for this rank
NOTE: the return is after termination; the objects are provided for query.
"""
if not haveMPI:
raise RuntimeError("spin_the_wheel called, but cannot import mpi4py")
# Confirm that the provided dictionaries specifying
# the hubs and spokes contain the appropriate keys
if "hub_class" not in hub_dict:
raise RuntimeError(
"The hub_dict must contain a 'hub_class' key specifying "
"the hub class to use")
if "opt_class" not in hub_dict:
raise RuntimeError(
"The hub_dict must contain an 'opt_class' key specifying "
"the SPBase class to use (e.g. PHBase, etc.)")
if "hub_kwargs" not in hub_dict:
hub_dict["hub_kwargs"] = dict()
if "opt_kwargs" not in hub_dict:
hub_dict["opt_kwargs"] = dict()
for spoke_dict in list_of_spoke_dict:
if "spoke_class" not in spoke_dict:
raise RuntimeError(
"Each spoke_dict must contain a 'spoke_class' key "
"specifying the spoke class to use")
if "opt_class" not in spoke_dict:
raise RuntimeError(
"Each spoke_dict must contain an 'opt_class' key "
"specifying the SPBase class to use (e.g. PHBase, etc.)")
if "spoke_kwargs" not in spoke_dict:
spoke_dict["spoke_kwargs"] = dict()
if "opt_kwargs" not in spoke_dict:
spoke_dict["opt_kwargs"] = dict()
if comm_world is None:
comm_world = MPI.COMM_WORLD
n_spokes = len(list_of_spoke_dict)
# Create the necessary communicators
fullcomm = comm_world
strata_comm, cylinder_comm = make_comms(n_spokes, fullcomm=fullcomm)
strata_rank = strata_comm.Get_rank()
cylinder_rank = cylinder_comm.Get_rank()
global_rank = fullcomm.Get_rank()
# Assign hub/spokes to individual ranks
if strata_rank == 0: # This rank is a hub
sp_class = hub_dict["hub_class"]
sp_kwargs = hub_dict["hub_kwargs"]
opt_class = hub_dict["opt_class"]
opt_kwargs = hub_dict["opt_kwargs"]
opt_dict = hub_dict
else: # This rank is a spoke
spoke_dict = list_of_spoke_dict[strata_rank - 1]
sp_class = spoke_dict["spoke_class"]
sp_kwargs = spoke_dict["spoke_kwargs"]
opt_class = spoke_dict["opt_class"]
opt_kwargs = spoke_dict["opt_kwargs"]
opt_dict = spoke_dict
# Create the appropriate opt object locally
opt_kwargs["mpicomm"] = cylinder_comm
opt = opt_class(**opt_kwargs)
# Create the SPCommunicator object (hub/spoke) with
# the appropriate SPBase object attached
if strata_rank == 0: # Hub
spcomm = sp_class(opt, fullcomm, strata_comm, cylinder_comm,
list_of_spoke_dict, **sp_kwargs)
else: # Spokes
spcomm = sp_class(opt, fullcomm, strata_comm, cylinder_comm,
**sp_kwargs)
# Create the windows, run main(), destroy the windows
spcomm.make_windows()
if strata_rank == 0:
spcomm.setup_hub()
global_toc("Starting spcomm.main()")
spcomm.main()
if strata_rank == 0: # If this is the hub
spcomm.send_terminate()
# Anything that's left to do
spcomm.finalize()
global_toc("Hub algorithm complete, waiting for termination barrier")
fullcomm.Barrier()
## give the hub the chance to catch new values
spcomm.hub_finalize()
spcomm.free_windows()
global_toc("Windows freed")
return spcomm, opt_dict
