Merge pull request #429 from Pyomo/newiter0callback

Adding iter0_post_solver_creation callback to extensions object
Pyomo · Sep 20, 2024 · 3de100e · 3de100e
2 parents d836482 + abb8e36
commit 3de100e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/mpisppy/extensions/extension.py b/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

diff --git a/mpisppy/phbase.py b/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