updates

watertap/core/plugins/solvers.py

@@ -30,6 +30,7 @@
 
 import pyomo.environ as pyo
 from pyomo.common.collections import Bunch
+from pyomo.common.dependencies import attempt_import
 from pyomo.common.modeling import unique_component_name
 from pyomo.contrib.pynumero.asl import AmplInterface
 from pyomo.contrib.pynumero.interfaces.pyomo_nlp import PyomoNLP
@@ -42,6 +43,8 @@
 )
 from idaes.logger import getLogger
 
+IPython, IPython_available = attempt_import("IPython")
+
 _log = getLogger("watertap.core")
 
 _pyomo_nl_writer_log = logging.getLogger("pyomo.repn.plugins.nl_writer")
@@ -372,3 +375,100 @@ def _constraint_autoscale_large_jac(
             # update the scaled jacobian
             jac_scaled[i, j] = jac_scaled[i, j] * sc
     return None
+
+
+class _BaseDebugSolverWrapper:
+
+    # defined by the derived class,
+    # created on the fly
+    _base_solver = None
+    _debug_solver_name = None
+
+    def __init__(self, **kwds):
+
+        kwds["name"] = self._debug_solver_name
+        self.options = Bunch()
+        if kwds.get("options") is not None:
+            for key in kwds["options"]:
+                setattr(self.options, key, kwds["options"][key])
+
+        self._value_cache = pyo.ComponentMap()
+
+    def restore_initial_values(self, blk):
+        for var in blk.component_data_objects(pyo.Var, descend_into=True):
+            var.set_value(self._value_cache[var], skip_validation=True)
+
+    def _cache_initial_values(self, blk):
+        for v in blk.component_data_objects(pyo.Var, descend_into=True):
+            self._value_cache[v] = v.value
+
+    def solve(self, blk, *args, **kwds):
+
+        if not IPython_available:
+            raise ImportError(f"The DebugSolverWrapper requires ipython.")
+
+        solver = pyo.SolverFactory(self._base_solver)
+
+        for k, v in self.options.items():
+            solver.options[k] = v
+
+        self._cache_initial_values(blk)
+
+        try:
+            results = solver.solve(blk, *args, **kwds)
+        except:
+            results = None
+        if results is not None and pyo.check_optimal_termination(results):
+            return results
+
+        # prevent circular imports
+        from watertap.core.util import model_debug_mode
+
+        # deactivate the model debug mode so we don't
+        # nest this environment within itself
+        model_debug_mode.deactivate()
+
+        self.restore_initial_values(blk)
+        debug = self
+
+        # else there was a problem
+        print(f"\nSolver debugging mode: the block {blk.name} failed to solve.\n")
+        print(f"{blk.name} can be called as `blk` in debugging mode.\n")
+        print(f"The solver {solver.name} is available in the variable `solver`.\n")
+        print(f"The initial values before the failed solve have been stored.\n")
+        print(
+            f"You can restore these initial values at anytime by calling `debug.restore_initial_values(blk)`\n."
+        )
+        print(
+            f"The model has been loaded into an IDAES DiagnosticsToolbox instance called `dt`.\n"
+        )
+        print(
+            "WARNING: If you ran your python file in an interactive window, this debugging mode will not work as intended. Be sure to run your python file in a terminal.\n"
+        )
+        from idaes.core.util.model_diagnostics import DiagnosticsToolbox
+
+        dt = DiagnosticsToolbox(blk)
+        IPython.embed(colors="neutral")
+
+        # activate the model debug mode
+        # to keep the state the same
+        model_debug_mode.activate()
+
+        return results
+
+
+def create_debug_solver_wrapper(solver_name):
+
+    assert pyo.SolverFactory(solver_name).available()
+
+    debug_solver_name = f"debug-solver-wrapper-{solver_name}"
+
+    @pyo.SolverFactory.register(
+        debug_solver_name,
+        doc=f"Debug solver wrapper for {solver_name}",
+    )
+    class DebugSolverWrapper(_BaseDebugSolverWrapper):
+        _base_solver = solver_name
+        _debug_solver_name = debug_solver_name
+
+    return debug_solver_name