diff --git a/pyomo/contrib/sensitivity_toolbox/pynumero.py b/pyomo/contrib/sensitivity_toolbox/pynumero.py new file mode 100644 index 00000000000..c139aca14fc --- /dev/null +++ b/pyomo/contrib/sensitivity_toolbox/pynumero.py @@ -0,0 +1,128 @@ +# ___________________________________________________________________________ +# +# 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 + +import pyomo.environ as pyo +import pyomo.contrib.pynumero.interfaces.pyomo_nlp as nlp +from pyomo.common.collections import ComponentSet, ComponentMap + + +def _coo_reorder_cols(mat, remap): + """Change the order of columns is a COO matrix. The main use of this is + to reorder variables in the Jacobian matrix. This changes the matrix in + place. This work 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 columns + 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, if you think you make want to + retain the cached objects in the pynumero interface, you can 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 new file mode 100644 index 00000000000..96959e4388e --- /dev/null +++ b/pyomo/contrib/sensitivity_toolbox/tests/test_pynumero.py @@ -0,0 +1,105 @@ +# ___________________________________________________________________________ +# +# 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 + +import pyomo.environ as pyo +import pyomo.contrib.pynumero.interfaces.pyomo_nlp as nlp +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)