From 9658599b6ec96db90a2fe798fd395b1ec00770cd Mon Sep 17 00:00:00 2001
From: bknueven <30801372+bknueven@users.noreply.github.com>
Date: Thu, 6 Feb 2025 14:39:19 -0700
Subject: [PATCH] Ensure `interval_initializer` always restores bounds (#1555)

* add failing test

* add fix for infeasible models
---
 watertap/core/util/initialization.py          | 61 +++++++++++--------
 .../core/util/tests/test_initialization.py    | 36 +++++++++++
 2 files changed, 70 insertions(+), 27 deletions(-)

diff --git a/watertap/core/util/initialization.py b/watertap/core/util/initialization.py
index ebc9bba266..f246860dba 100644
--- a/watertap/core/util/initialization.py
+++ b/watertap/core/util/initialization.py
@@ -162,31 +162,38 @@ def interval_initializer(
     for v in blk.component_data_objects(Var, active=True, descend_into=True):
         bound_cache[v] = v.bounds
 
-    fbbt(blk, feasibility_tol=feasibility_tol, deactivate_satisfied_constraints=False)
-
-    for v, bounds in bound_cache.items():
-        if v.value is None:
-            logger.info(
-                f"variable {v.name} has no initial value: setting to {default_initial_value}"
-            )
-            v.set_value(default_initial_value, skip_validation=True)
-        if v.lb is not None:
-            if v.lb == v.ub:
-                logger.debug(f"setting {v.name} to derived value {v.value}")
-                v.set_value(v.lb, skip_validation=True)
-                continue
-            if v.value < v.lb:
-                logger.debug(
-                    f"projecting {v.name} at value {v.value} onto derived lower bound {v.lb}"
+    try:
+        fbbt(
+            blk, feasibility_tol=feasibility_tol, deactivate_satisfied_constraints=False
+        )
+
+        for v, bounds in bound_cache.items():
+            if v.value is None:
+                logger.info(
+                    f"variable {v.name} has no initial value: setting to {default_initial_value}"
                 )
-                v.set_value(v.lb, skip_validation=True)
-        if v.ub is not None:
-            if v.value > v.ub:
-                logger.debug(
-                    f"projecting {v.name} at value {v.value} onto derived upper bound {v.ub}"
-                )
-                v.set_value(v.ub, skip_validation=True)
-
-    for v, bounds in bound_cache.items():
-        # restore bounds to original
-        v.bounds = bounds
+                v.set_value(default_initial_value, skip_validation=True)
+            if v.lb is not None:
+                if v.lb == v.ub:
+                    logger.debug(f"setting {v.name} to derived value {v.value}")
+                    v.set_value(v.lb, skip_validation=True)
+                    continue
+                if v.value < v.lb:
+                    logger.debug(
+                        f"projecting {v.name} at value {v.value} onto derived lower bound {v.lb}"
+                    )
+                    v.set_value(v.lb, skip_validation=True)
+            if v.ub is not None:
+                if v.value > v.ub:
+                    logger.debug(
+                        f"projecting {v.name} at value {v.value} onto derived upper bound {v.ub}"
+                    )
+                    v.set_value(v.ub, skip_validation=True)
+
+    except:
+        raise
+
+    finally:
+        # restore the bounds before leaving this function
+        for v, bounds in bound_cache.items():
+            v.bounds = bounds
diff --git a/watertap/core/util/tests/test_initialization.py b/watertap/core/util/tests/test_initialization.py
index 73e77a70e2..5ce0f2b43c 100644
--- a/watertap/core/util/tests/test_initialization.py
+++ b/watertap/core/util/tests/test_initialization.py
@@ -179,3 +179,39 @@ def test_interval_initializer(self, m):
         assert m.y.ub == None
         assert m.z.lb == None
         assert m.z.ub == None
+
+    @pytest.fixture(scope="class")
+    def m_infeas(self):
+
+        # This is the same model used in the pyomo fbbt test at
+        # https://github.com/Pyomo/pyomo/blob/0e749d0c993df960af6cde0e775bef7cab6e2568/pyomo/contrib/fbbt/tests/test_fbbt.py#L957C9-L966C32
+
+        m = ConcreteModel()
+        m.x = Var(bounds=(-3, 3))
+        m.y = Var(bounds=(0, None))
+        m.z = Var()
+        m.c = ConstraintList()
+        m.c.add(m.x + m.y >= -1)
+        m.c.add(m.x + m.y <= -2)
+        m.c.add(m.y - m.x * m.z <= 2)
+        m.c.add(m.y - m.x * m.z >= -2)
+        m.c.add(m.x + m.z == 1)
+
+        return m
+
+    @pytest.mark.unit
+    def test_interval_initializer_failure_restores_bounds(self, m_infeas):
+        m = m_infeas
+        feasibility_tol = 1.0e-6
+        try:
+            interval_initializer(m, feasibility_tol=feasibility_tol)
+        except:
+            pass
+
+        # Assert the restored bounds
+        assert m.x.lb == -3.0
+        assert m.x.ub == 3.0
+        assert m.y.lb == 0.0
+        assert m.y.ub == None
+        assert m.z.lb == None
+        assert m.z.ub == None