Merge pull request #360 from michaelbynum/value_bug

Ensure variables and parameters keep their values when constraints are removed

wntr/sim/aml/aml.py

@@ -153,6 +153,7 @@ def _decrement_var(self, var):
         if self._refcounts[var] == 0:
             cvar = self._var_cvar_map[var]
             var._c_obj = None
+            var._value = cvar.value
             del self._refcounts[var]
             del self._var_cvar_map[var]
             self._evaluator.remove_var(cvar)
@@ -162,6 +163,7 @@ def _decrement_param(self, p):
         if self._refcounts[p] == 0:
             cparam = self._param_cparam_map[p]
             p._c_obj = None
+            p._value = cparam.value
             del self._refcounts[p]
             del self._param_cparam_map[p]
             self._evaluator.remove_param(cparam)

wntr/tests/test_aml.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import wntr.sim.aml as aml
+from wntr.sim.solvers import NewtonSolver, SolverStatus
 
 
 def compare_evaluation(self, m, true_r, true_j):
@@ -17,6 +18,30 @@ def compare_evaluation(self, m, true_r, true_j):
             self.assertAlmostEqual(true_j[c][v], j[c.index, v.index], 10)
 
 
+class TestModel(unittest.TestCase):
+    def test_var_value_with_decrement(self):
+        m = aml.Model()
+        m.x = aml.Var()
+        m.p = aml.Param(val=1)
+        m.c = aml.Constraint(m.x - m.p)
+        m.set_structure()
+        opt = NewtonSolver({'TOL': 1e-8})
+        status, msg, num_iter = opt.solve(m)
+        self.assertEqual(status, SolverStatus.converged)
+        self.assertAlmostEqual(m.x.value, 1)
+        m.p.value = 2
+        status, msg, num_iter = opt.solve(m)
+        self.assertEqual(status, SolverStatus.converged)
+        self.assertAlmostEqual(m.x.value, 2)
+        del m.c
+        m.c = aml.Constraint(m.x**0.5 - m.p)
+        m.set_structure()
+        self.assertAlmostEqual(m.x.value, 2)  # this is the real test here
+        status, msg, num_iter = opt.solve(m)
+        self.assertEqual(status, SolverStatus.converged)
+        self.assertAlmostEqual(m.x.value, 4)
+
+
 class TestExpression(unittest.TestCase):
     def test_add(self):
         m = aml.Model()