Merge branch 'main' into sep_rho

Author: bknueven

Diff for: mpisppy/extensions/extension.py

@@ -93,13 +93,18 @@ def post_solve_loop(self):
         pass
 
     def pre_iter0(self):
-        ''' Method called at the end of PH_Prep().
-            When this method is called, all scenarios have been created, and
+        ''' When this method is called, all scenarios have been created, and
             the dual/prox terms have been attached to the objective, but the
             solvers have not yet been created.
         '''
         pass
 
+    def iter0_post_solver_creation(self):
+        ''' When this method is called, PH iteration 0 has been initiated and 
+            all solver objects have been created.
+        '''
+        pass
+
     def post_iter0(self):
         ''' Method called after the first PH iteration.
             When this method is called, one call to solve_loop() has been
@@ -192,6 +197,10 @@ def pre_iter0(self):
         for lobject in self.extdict.values():
             lobject.pre_iter0()
 
+    def iter0_post_solver_creation(self):
+        for lobject in self.extdict.values():
+            lobject.iter0_post_solver_creation()        
+
     def post_iter0(self):
         for lobject in self.extdict.values():
             lobject.post_iter0()

Diff for: mpisppy/phbase.py

@@ -772,8 +772,7 @@ def PH_Prep(
 
         Note:
             This function constructs an Extension object if one was specified
-            at the time the PH object was created. It also calls the
-            `pre_iter0` method of the Extension object.
+            at the time the PH object was created. 
         """
 
         self.attach_Ws_and_prox()
@@ -832,9 +831,10 @@ def Iter0(self):
                 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"]
@@ -853,6 +853,9 @@ def _vb(msg):
         global_toc("Creating solvers")
         self._create_solvers()
 
+        if (self.extensions is not None):
+            self.extobject.iter0_post_solver_creation()        
+
         teeme = ("tee-rank0-solves" in self.options
                  and self.options['tee-rank0-solves']
                  and self.cylinder_rank == 0