diff --git a/doc/optimizers/SNOPT_options.yaml b/doc/optimizers/SNOPT_options.yaml index 0c740234..65240285 100644 --- a/doc/optimizers/SNOPT_options.yaml +++ b/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`` diff --git a/pyoptsparse/__init__.py b/pyoptsparse/__init__.py index 5ab4a539..c122965d 100644 --- a/pyoptsparse/__init__.py +++ b/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 diff --git a/pyoptsparse/pySNOPT/pySNOPT.py b/pyoptsparse/pySNOPT/pySNOPT.py index a7d7c168..1a39fb2b 100644 --- a/pyoptsparse/pySNOPT/pySNOPT.py +++ b/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) diff --git a/tests/test_hs015.py b/tests/test_hs015.py index 88f8779c..27ee8fd3 100644 --- a/tests/test_hs015.py +++ b/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]}