Merge pull request #3561 from eslickj/sens

blnicho · web-flow · commit 926c91a3dcbc · 2025-07-02T23:11:57.000-06:00
Add some pynumero-based sensitivity functions
diff --git a/pyomo/contrib/sensitivity_toolbox/pynumero.py b/pyomo/contrib/sensitivity_toolbox/pynumero.py
@@ -0,0 +1,129 @@
+#  ___________________________________________________________________________
+#
+#  Pyomo: Python Optimization Modeling Objects
+#  Copyright (c) 2008-2025
+#  National Technology and Engineering Solutions of Sandia, LLC
+#  Under the terms of Contract DE-NA0003525 with National Technology and
+#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
+#  rights in this software.
+#  This software is distributed under the 3-clause BSD License.
+#  ___________________________________________________________________________
+
+from pyomo.common.dependencies import numpy as np
+from pyomo.common.dependencies import scipy, attempt_import
+
+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]
+from pyomo.common.collections import ComponentSet, ComponentMap
+
+
+def _coo_reorder_cols(mat, remap):
+    """Change the order of columns in a COO matrix. The main use of this is
+    to reorder variables in the Jacobian matrix. This changes the matrix in
+    place.
+
+    Parameters
+    ----------
+    mat: scipy.sparse.coo_matrix
+        Reorder the columns of this matrix
+    remap: dict
+        dictionary where keys are old column and value is new column, if a column
+        doesn't move, it doesn't need to be included.
+
+    Returns
+    -------
+    NoneType
+        None
+    """
+    for i in range(len(mat.data)):
+        try:
+            mat.col[i] = remap[mat.col[i]]
+        except KeyError:
+            pass  # it's fine if we don't move a col in remap
+
+
+def get_dsdp_dfdp(model, theta):
+    """Calculate the derivatives of the state variables (s) with respect to
+    parameters (p) (ds/dp), and the derivative of the objective function (f)
+    with respect to p (df/dp). The number of parameters in theta should be the
+    same as the number of degrees of freedom.
+
+    Parameters
+    ----------
+    model: pyomo.environ.Block | pyomo.contrib.pynumero.interfaces.PyomoNLP
+        Model to calculate sensitivity on. To retain the cached objects in
+        the pynumero interface, create a PyomoNLP first and pass it to this function.
+    theta: list
+        A list of parameters as pyomo.environ.VarData, the number of parameters
+        should be equal to the degrees of freedom.
+
+    Returns
+    -------
+    scipy.sparse.csc_matrix, csc_matrix, ComponentMap, ComponentMap
+        ds/dp (ns by np), df/dp (1 by np), row map, column map.
+        The column map maps Pyomo variables p to columns and the
+        row map maps Pyomo variables s to rows.
+    """
+    # Create a Pynumero NLP and get Jacobian
+    if isinstance(model, nlp.PyomoNLP):
+        m2 = model
+    else:
+        m2 = nlp.PyomoNLP(model)
+    J = m2.evaluate_jacobian_eq()
+    v_list = m2.get_pyomo_variables()
+    # Map variables to columns in J
+    mv_map = {id(v): i for i, v in enumerate(v_list)}
+    s_list = list(ComponentSet(v_list) - ComponentSet(theta))
+    ns = len(s_list)
+    np = len(theta)
+    col_remap = {mv_map[id(v)]: i for i, v in enumerate(s_list + theta)}
+    _coo_reorder_cols(J, remap=col_remap)
+    J = J.tocsc()
+    dB = -(
+        J
+        @ scipy.sparse.vstack(
+            (scipy.sparse.coo_matrix((ns, np)), scipy.sparse.identity(np))
+        ).tocsc()
+    )
+    # Calculate sensitivity matrix
+    dsdp = scipy.sparse.linalg.spsolve(J[:, range(ns)], dB)
+    # Get a map of state vars to columns
+    s_map = {id(v): i for i, v in enumerate(s_list)}
+    # Get the outputs we are interested in from the list of output vars
+    column_map = ComponentMap([(v, i) for i, v in enumerate(theta)])
+    row_map = ComponentMap([(v, i) for i, v in enumerate(s_list)])
+    dfdx = scipy.sparse.coo_matrix(m2.evaluate_grad_objective())
+    _coo_reorder_cols(dfdx, remap=col_remap)
+    dfdx = dfdx.tocsc()
+    dfdp = dfdx[0, :ns] @ dsdp + dfdx[0, ns:]
+    # return sensitivity of the outputs to p and component maps
+    return dsdp, dfdp, row_map, column_map
+
+
+def get_dydp(y_list, dsdp, row_map):
+    """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)
+    dsdp: csc_matrix
+        A sensitivity matrix calculated by get_dsdp_dfdp
+    row_map: ComponentMap
+        A row map from get_dsdp_dfdp
+
+    Returns
+    -------
+    csc_matrix, ComponentMap
+        dy/dp and a new row map with only y variables
+
+    """
+    new_row_map = ComponentMap()
+    for i, v in enumerate(y_list):
+        new_row_map[v] = i
+    rows = [row_map[v] for v in y_list]
+    dydp = dsdp[rows, :]
+    return dydp, new_row_map
diff --git a/pyomo/contrib/sensitivity_toolbox/tests/test_pynumero.py b/pyomo/contrib/sensitivity_toolbox/tests/test_pynumero.py
@@ -0,0 +1,107 @@
+#  ___________________________________________________________________________
+#
+#  Pyomo: Python Optimization Modeling Objects
+#  Copyright (c) 2008-2025
+#  National Technology and Engineering Solutions of Sandia, LLC
+#  Under the terms of Contract DE-NA0003525 with National Technology and
+#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
+#  rights in this software.
+#  This software is distributed under the 3-clause BSD License.
+#  ___________________________________________________________________________
+
+
+import pyomo.common.unittest as unittest
+from pyomo.common.dependencies import numpy as np, numpy_available
+from pyomo.common.dependencies import scipy_available, attempt_import
+
+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]
+import pyomo.contrib.sensitivity_toolbox.pynumero as pnsens
+from pyomo.contrib.pynumero.asl import AmplInterface
+
+if not scipy_available or not numpy_available:
+    raise unittest.SkipTest("scipy or numpy is not available")
+
+if not AmplInterface.available():
+    raise unittest.SkipTest("Pynumero needs the ASL extension to run NLP tests")
+
+
+class TestSeriesData(unittest.TestCase):
+    def test_dsdp_dfdp_pyomo(self):
+        m = pyo.ConcreteModel()
+        m.x1 = pyo.Var(initialize=200)
+        m.x2 = pyo.Var(initialize=5)
+        m.p1 = pyo.Var(initialize=10)
+        m.p2 = pyo.Var(initialize=5)
+        m.obj = pyo.Objective(
+            expr=m.x1 * m.p1 + m.x2 * m.x2 * m.p2 + m.p1 * m.p2, sense=pyo.minimize
+        )
+        m.c1 = pyo.Constraint(expr=m.x1 == 2 * m.p1**2)
+        m.c2 = pyo.Constraint(expr=m.x2 == m.p2)
+        theta = [m.p1, m.p2]
+
+        dsdp, dfdp, rmap, cmap = pnsens.get_dsdp_dfdp(m, theta)
+
+        # Since x1 = p1
+        np.testing.assert_almost_equal(dsdp[rmap[m.x1], cmap[m.p1]], 40.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x1], cmap[m.p2]], 0.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x2], cmap[m.p1]], 0.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x2], cmap[m.p2]], 1.0)
+
+        # if x1 = 2 * p1 and x2 = p2 then
+        #   df/dp1 = 6 p1**2 + p2 = 45.0
+        #   df/dp2 = 3 p2 + p1 = 85.0
+        np.testing.assert_almost_equal(dfdp[0, cmap[m.p1]], 605.0)
+        np.testing.assert_almost_equal(dfdp[0, cmap[m.p2]], 85.0)
+
+    def test_dsdp_dfdp_pyomo_nlp(self):
+        m = pyo.ConcreteModel()
+        m.x1 = pyo.Var(initialize=200)
+        m.x2 = pyo.Var(initialize=5)
+        m.p1 = pyo.Var(initialize=10)
+        m.p2 = pyo.Var(initialize=5)
+        m.obj = pyo.Objective(
+            expr=m.x1 * m.p1 + m.x2 * m.x2 * m.p2 + m.p1 * m.p2, sense=pyo.minimize
+        )
+        m.c1 = pyo.Constraint(expr=m.x1 == 2 * m.p1**2)
+        m.c2 = pyo.Constraint(expr=m.x2 == m.p2)
+        theta = [m.p1, m.p2]
+
+        m2 = nlp.PyomoNLP(m)
+        dsdp, dfdp, rmap, cmap = pnsens.get_dsdp_dfdp(m2, theta)
+
+        # Since x1 = p1
+        assert dsdp.shape == (2, 2)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x1], cmap[m.p1]], 40.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x1], cmap[m.p2]], 0.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x2], cmap[m.p1]], 0.0)
+        np.testing.assert_almost_equal(dsdp[rmap[m.x2], cmap[m.p2]], 1.0)
+
+        # if x1 = 2 * p1 and x2 = p2 then
+        #   df/dp1 = 6 p1**2 + p2 = 45.0
+        #   df/dp2 = 3 p2 + p1 = 85.0
+        assert dfdp.shape == (1, 2)
+        np.testing.assert_almost_equal(dfdp[0, cmap[m.p1]], 605.0)
+        np.testing.assert_almost_equal(dfdp[0, cmap[m.p2]], 85.0)
+
+    def test_dydp_pyomo(self):
+        m = pyo.ConcreteModel()
+        m.x1 = pyo.Var(initialize=200)
+        m.x2 = pyo.Var(initialize=5)
+        m.p1 = pyo.Var(initialize=10)
+        m.p2 = pyo.Var(initialize=5)
+        m.obj = pyo.Objective(
+            expr=m.x1 * m.p1 + m.x2 * m.x2 * m.p2 + m.p1 * m.p2, sense=pyo.minimize
+        )
+        m.c1 = pyo.Constraint(expr=m.x1 == 2 * m.p1**2)
+        m.c2 = pyo.Constraint(expr=m.x2 == 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)
+
+        np.testing.assert_almost_equal(dydp[rmap[m.x2], cmap[m.p1]], 0.0)
+        np.testing.assert_almost_equal(dydp[rmap[m.x2], cmap[m.p2]], 1.0)
+        assert dydp.shape == (1, 2)
