Added restartDict to snstop (#404)

* save restartDict after every major iteration

* added restartDict as snstop callback argument

* keep the final restartDict

* added example to docs

* added test

* test major iterations instead of function evaluations

* bumped minor version

* added snstop arguments option

* added work arrays save option

* added error for unknown snstop arg

* updated docs

* addressed Ella's comments

* minor version bump

* updated comment

---------

Co-authored-by: Marco Mangano <[email protected]>

doc/optimizers/SNOPT_options.yaml

@@ -100,7 +100,37 @@ Return work arrays:
     These arrays can be used to hot start a subsequent optimization.
     The SNOPT option 'Sticky parameters' will also be automatically set to 'Yes' to facilitate the hot start.
 
+Work arrays save file:
+  desc: >
+    This option is unique to the Python wrapper.
+    The SNOPT work arrays will be pickled and saved to this file after each major iteration.
+    This file is useful if you want to restart an optimization that did not exit cleanly.
+    If None, the work arrays are not saved.
+
 snSTOP function handle:
   desc: >
     This option is unique to the Python wrapper.
     A function handle can be supplied which is called at the end of each major iteration.
+    The following is an example of a callback function that saves the restart dictionary
+    to a different file after each major iteration.
+
+      .. code-block:: python
+
+        def snstopCallback(iterDict, restartDict):
+            # Get the major iteration number
+            nMajor = iterDict["nMajor"]
+
+            # Save the restart dictionary
+            writePickle(f"restart_{nMajor}.pickle", restartDict)
+
+            return 0
+
+snSTOP arguments:
+  desc: |
+    This option is unique to the Python wrapper.
+    It specifies a list of arguments that will be passed to the snSTOP function handle.
+    ``iterDict`` is always passed as an argument.
+    Additional arguments are passed in the same order as this list.
+    The possible values are
+
+    - ``restartDict``

pyoptsparse/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.11.3"
+__version__ = "2.12.0"
 
 from .pyOpt_history import History
 from .pyOpt_variable import Variable

pyoptsparse/pySNOPT/pySNOPT.py

@@ -12,7 +12,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 # External modules
-from baseclasses.utils import CaseInsensitiveSet
+from baseclasses.utils import CaseInsensitiveSet, writePickle
 import numpy as np
 from numpy import ndarray
 from pkg_resources import parse_version
@@ -60,7 +60,9 @@ def __init__(self, raiseError=True, options: Dict = {}):
             {
                 "Save major iteration variables",
                 "Return work arrays",
+                "Work arrays save file",
                 "snSTOP function handle",
+                "snSTOP arguments",
             }
         )
 
@@ -118,7 +120,9 @@ def _getDefaultOptions() -> Dict[str, Any]:
             "Total real workspace": [int, None],
             "Save major iteration variables": [list, []],
             "Return work arrays": [bool, False],
+            "Work arrays save file": [(type(None), str), None],
             "snSTOP function handle": [(type(None), type(lambda: None)), None],
+            "snSTOP arguments": [list, []],
         }
         return defOpts
 
@@ -667,12 +671,39 @@ def _snstop(self, ktcond, mjrprtlvl, minimize, n, nncon, nnobj, ns, itn, nmajor,
                 if "funcs" in self.cache.keys():
                     iterDict["funcs"].update(self.cache["funcs"])
 
+        # Create the restart dictionary to be passed to snstop_handle
+        restartDict = {
+            "cw": cw,
+            "iw": iw,
+            "rw": rw,
+            "xs": x,  # x is the same as xs; we call it x here to be consistent with the SNOPT subroutine snSTOP
+            "hs": hs,
+            "pi": pi,
+        }
+
+        workArraysSave = self.getOption("Work arrays save file")
+        if workArraysSave is not None:
+            # Save the restart dictionary
+            writePickle(workArraysSave, restartDict)
+
         # perform callback if requested
         snstop_handle = self.getOption("snSTOP function handle")
         if snstop_handle is not None:
+
+            # Get the arguments to pass in to snstop_handle
+            # iterDict is always included
+            snstopArgs = [iterDict]
+            for snstopArg in self.getOption("snSTOP arguments"):
+                if snstopArg == "restartDict":
+                    snstopArgs.append(restartDict)
+                else:
+                    raise Error(f"Received unknown snSTOP argument {snstopArg}. "
+                                + "Please see 'snSTOP arguments' option in the pyOptSparse documentation "
+                                + "under 'SNOPT'.")
+
             if not self.storeHistory:
                 raise Error("snSTOP function handle must be used with storeHistory=True")
-            iabort = snstop_handle(iterDict)
+            iabort = snstop_handle(*snstopArgs)
             # write iterDict again if anything was inserted
             if self.storeHistory and callCounter is not None:
                 self.hist.write(callCounter, iterDict)

tests/test_hs015.py

@@ -1,9 +1,11 @@
 """Test solution of problem HS15 from the Hock & Schittkowski collection"""
 
 # Standard Python modules
+import os
 import unittest
 
 # External modules
+from baseclasses.utils import readPickle, writePickle
 import numpy as np
 from parameterized import parameterized
 
@@ -193,6 +195,90 @@ def test_snopt_snstop(self):
         # we should get 70/74
         self.assert_inform_equal(sol, optInform=74)
 
+    def test_snopt_snstop_restart(self):
+        pickleFile = "restart.pickle"
+
+        def my_snstop_restart(iterDict, restartDict):
+            # Save the restart dictionary
+            writePickle(pickleFile, restartDict)
+
+            # Exit after 5 major iterations
+            if iterDict["nMajor"] == 5:
+                return 1
+
+            return 0
+
+        # Run the optimization for 5 major iterations
+        self.optName = "SNOPT"
+        self.setup_optProb()
+        optOptions = {
+            "snSTOP function handle": my_snstop_restart,
+            "snSTOP arguments": ["restartDict"],
+        }
+        sol = self.optimize(optOptions=optOptions, storeHistory=True)
+
+        # Check that the optimization exited with 74
+        self.assert_inform_equal(sol, optInform=74)
+
+        # Read the restart dictionary pickle file saved by snstop
+        restartDict = readPickle(pickleFile)
+
+        # Now optimize again but using the restart dictionary
+        self.setup_optProb()
+        opt = OPT(
+            self.optName,
+            options={
+                "Start": "Hot",
+                "Verify level": -1,
+                "snSTOP function handle": my_snstop_restart,
+                "snSTOP arguments": ["restartDict"],
+            },
+        )
+        histFile = "restart.hst"
+        sol = opt(self.optProb, sens=self.sens, storeHistory=histFile, restartDict=restartDict)
+
+        # Check that the optimization converged in fewer than 5 more major iterations
+        self.assert_solution_allclose(sol, 1e-12)
+        self.assert_inform_equal(sol, optInform=1)
+
+        # Delete the pickle and history files
+        os.remove(pickleFile)
+        os.remove(histFile)
+
+    def test_snopt_work_arrays_save(self):
+        # Run the optimization for 5 major iterations
+        self.optName = "SNOPT"
+        self.setup_optProb()
+        pickleFile = "work_arrays_save.pickle"
+        optOptions = {
+            "snSTOP function handle": self.my_snstop,
+            "Work arrays save file": pickleFile,
+        }
+        sol = self.optimize(optOptions=optOptions, storeHistory=True)
+
+        # Read the restart dictionary pickle file saved by snstop
+        restartDict = readPickle(pickleFile)
+
+        # Now optimize again but using the restart dictionary
+        self.setup_optProb()
+        opt = OPT(
+            self.optName,
+            options={
+                "Start": "Hot",
+                "Verify level": -1,
+            },
+        )
+        histFile = "work_arrays_save.hst"
+        sol = opt(self.optProb, sens=self.sens, storeHistory=histFile, restartDict=restartDict)
+
+        # Check that the optimization converged
+        self.assert_solution_allclose(sol, 1e-12)
+        self.assert_inform_equal(sol, optInform=1)
+
+        # Delete the pickle and history files
+        os.remove(pickleFile)
+        os.remove(histFile)
+
     def test_snopt_failed_initial(self):
         def failed_fun(x_dict):
             funcs = {"obj": 0.0, "con": [np.nan, np.nan]}