Sensitivity calculation for named expressions

pyomo/contrib/sensitivity_toolbox/pynumero.py

@@ -17,6 +17,7 @@
 # Use attempt_import here due to unguarded NumPy import in this file
 nlp = attempt_import('pyomo.contrib.pynumero.interfaces.pyomo_nlp')[0]
 from pyomo.common.collections import ComponentSet, ComponentMap
+import pyomo.core.expr.calculus.derivatives as derivatives
 
 
 def _coo_reorder_cols(mat, remap):
@@ -111,28 +112,56 @@ def get_dsdp_dfdp(model, theta):
     return dsdp, dfdp, row_map, column_map
 
 
-def get_dydp(y_list, dsdp, row_map):
+def get_dydp(y_list, dsdp, row_map, column_map=None):
     """Reduce the sensitivity matrix from get_dsdp_dfdp to only
     a specified set of state variables of interest.
 
     Parameters
     ----------
     y_list: list
-        A list of state variables of interest (a subset of s)
+        A list of state variables or named expressions
     dsdp: csc_matrix
         A sensitivity matrix calculated by get_dsdp_dfdp
     row_map: ComponentMap
         A row map from get_dsdp_dfdp
+    column_map: ComponentMap
+        A column map from get_dsdp_dfdp, only needed if y_list
+        contains expressions
 
     Returns
     -------
     csc_matrix, ComponentMap
         dy/dp and a new row map with only y variables
 
     """
+    j = 0
     new_row_map = ComponentMap()
+    expr_row_map = ComponentMap()
+    rows = [None] * len(y_list)
     for i, v in enumerate(y_list):
         new_row_map[v] = i
-    rows = [row_map[v] for v in y_list]
-    dydp = dsdp[rows, :]
+        if v not in row_map:
+            expr_row_map[v] = j
+            j += 1
+        else:
+            rows[i] = dsdp[row_map[v], :]
+
+    if j > 0:
+        if column_map is None:
+            raise ValueError(
+                "A column_map must be provided if y_list contains named expressions."
+            )
+        wrt_list = [s for s in row_map]
+        dedx = [None] * j
+        for v, i in expr_row_map.items():
+            dedx[i] = scipy.sparse.csc_matrix(
+                derivatives.differentiate(v, wrt_list=wrt_list)
+            )
+        dedx = scipy.sparse.vstack(dedx)
+        ns = len(row_map) - len(column_map)
+        dedp = dedx[:, :ns] @ dsdp[:ns, :] + dedx[:, ns:]
+        for v, i in expr_row_map.items():
+            rows[new_row_map[v]] = dedp[i, :]
+
+    dydp = scipy.sparse.vstack(rows)
     return dydp, new_row_map

pyomo/contrib/sensitivity_toolbox/tests/test_pynumero.py

@@ -17,7 +17,7 @@
 import pyomo.environ as pyo
 
 # Use attempt_import here due to unguarded NumPy import in this file
-nlp = attempt_import('pyomo.contrib.pynumero.interfaces.pyomo_nlp')[0]
+nlp = attempt_import("pyomo.contrib.pynumero.interfaces.pyomo_nlp")[0]
 import pyomo.contrib.sensitivity_toolbox.pynumero as pnsens
 from pyomo.contrib.pynumero.asl import AmplInterface
 
@@ -101,11 +101,31 @@ def test_dydp_pyomo(self):
         )
         m.c1 = pyo.Constraint(expr=m.x1 == 2 * m.p1**2)
         m.c2 = pyo.Constraint(expr=m.x2 == m.p2)
+
+        m.e1 = pyo.Expression(expr=2 * m.p1**2)
+        m.e2 = pyo.Expression(expr=m.p2)
+        m.e3 = pyo.Expression(expr=m.x1 * m.p1 + m.x2 * m.x2 * m.p2 + m.p1 * m.p2)
+
         theta = [m.p1, m.p2]
         dsdp, dfdp, rmap, cmap = pnsens.get_dsdp_dfdp(m, theta)
 
-        dydp, rmap = pnsens.get_dydp([m.x2], dsdp, rmap)
+        dydp, rmap = pnsens.get_dydp([m.e1, m.x1, m.e2, m.x2, m.e3], dsdp, rmap, cmap)
 
+        assert dydp.shape == (5, 2)
+        np.testing.assert_almost_equal(dydp[rmap[m.x1], cmap[m.p1]], 40.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.e1], cmap[m.p1]], 40.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.x1], cmap[m.p2]], 0.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.e1], cmap[m.p2]], 0.0)
         np.testing.assert_almost_equal(dydp[rmap[m.x2], cmap[m.p1]], 0.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.e2], cmap[m.p1]], 0.0)
         np.testing.assert_almost_equal(dydp[rmap[m.x2], cmap[m.p2]], 1.0)
-        assert dydp.shape == (1, 2)
+        np.testing.assert_almost_equal(dydp[rmap[m.e2], cmap[m.p2]], 1.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.e3], cmap[m.p1]], 605.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.e3], cmap[m.p2]], 85.0)
+
+        # make sure the rows are in the order of y_list
+        assert rmap[m.e1] == 0
+        assert rmap[m.x1] == 1
+        assert rmap[m.e2] == 2
+        assert rmap[m.x2] == 3
+        assert rmap[m.e3] == 4