Adds ONNX Surrogate support from OMLT

idaes/core/surrogate/keras_surrogate.py

@@ -17,23 +17,21 @@
 # pylint: disable=missing-class-docstring
 # pylint: disable=missing-function-docstring
 
-from enum import Enum
 import json
 import os.path
 
-import numpy as np
 import pandas as pd
 
 from pyomo.common.dependencies import attempt_import
 
-from idaes.core.surrogate.base.surrogate_base import SurrogateBase
 from idaes.core.surrogate.sampling.scaling import OffsetScaler
 
+from idaes.core.surrogate.omlt_base_surrogate_class import OMLTSurrogate
+
 keras, keras_available = attempt_import("tensorflow.keras")
 omlt, omlt_available = attempt_import("omlt")
 
 if omlt_available:
-    from omlt import OmltBlock, OffsetScaling
     from omlt.neuralnet.nn_formulation import (
         FullSpaceSmoothNNFormulation,
         ReducedSpaceSmoothNNFormulation,
@@ -45,7 +43,7 @@
         from omlt.io import load_keras_sequential
 
 
-class KerasSurrogate(SurrogateBase):
+class KerasSurrogate(OMLTSurrogate):
     def __init__(
         self,
         keras_model,
@@ -87,52 +85,11 @@ def __init__(
             input_labels=input_labels,
             output_labels=output_labels,
             input_bounds=input_bounds,
+            input_scaler=input_scaler,
+            output_scaler=output_scaler,
         )
-
-        # make sure we are using the standard scaler
-        if (
-            input_scaler is not None
-            and not isinstance(input_scaler, OffsetScaler)
-            or output_scaler is not None
-            and not isinstance(output_scaler, OffsetScaler)
-        ):
-            raise NotImplementedError("KerasSurrogate only supports the OffsetScaler.")
-
-        # check that the input labels match
-        if input_scaler is not None and input_scaler.expected_columns() != input_labels:
-            raise ValueError(
-                "KerasSurrogate created with input_labels that do not match"
-                " the expected columns in the input_scaler.\n"
-                "input_labels={}\n"
-                "input_scaler.expected_columns()={}".format(
-                    input_labels, input_scaler.expected_columns()
-                )
-            )
-
-        # check that the output labels match
-        if (
-            output_scaler is not None
-            and output_scaler.expected_columns() != output_labels
-        ):
-            raise ValueError(
-                "KerasSurrogate created with output_labels that do not match"
-                " the expected columns in the output_scaler.\n"
-                "output_labels={}\n"
-                "output_scaler.expected_columns()={}".format(
-                    output_labels, output_scaler.expected_columns()
-                )
-            )
-
-        self._input_scaler = input_scaler
-        self._output_scaler = output_scaler
         self._keras_model = keras_model
 
-    class Formulation(Enum):
-        FULL_SPACE = 1
-        REDUCED_SPACE = 2
-        RELU_BIGM = 3
-        RELU_COMPLEMENTARITY = 4
-
     def populate_block(self, block, additional_options=None):
         """
         Method to populate a Pyomo Block with the keras model constraints.
@@ -149,37 +106,7 @@ def populate_block(self, block, additional_options=None):
         formulation = additional_options.pop(
             "formulation", KerasSurrogate.Formulation.FULL_SPACE
         )
-        offset_inputs = np.zeros(self.n_inputs())
-        factor_inputs = np.ones(self.n_inputs())
-        offset_outputs = np.zeros(self.n_outputs())
-        factor_outputs = np.ones(self.n_outputs())
-        if self._input_scaler:
-            offset_inputs = self._input_scaler.offset_series()[
-                self.input_labels()
-            ].to_numpy()
-            factor_inputs = self._input_scaler.factor_series()[
-                self.input_labels()
-            ].to_numpy()
-        if self._output_scaler:
-            offset_outputs = self._output_scaler.offset_series()[
-                self.output_labels()
-            ].to_numpy()
-            factor_outputs = self._output_scaler.factor_series()[
-                self.output_labels()
-            ].to_numpy()
-
-        # build the OMLT scaler object
-        omlt_scaling = OffsetScaling(
-            offset_inputs=offset_inputs,
-            factor_inputs=factor_inputs,
-            offset_outputs=offset_outputs,
-            factor_outputs=factor_outputs,
-        )
-
-        # omlt takes *scaled* input bounds as a dictionary with int keys
-        input_bounds = dict(enumerate(self.input_bounds().values()))
-        scaled_input_bounds = omlt_scaling.get_scaled_input_expressions(input_bounds)
-        scaled_input_bounds = {i: tuple(bnd) for i, bnd in scaled_input_bounds.items()}
+        omlt_scaling, scaled_input_bounds = self.generate_omlt_scaling_objecets()
 
         net = load_keras_sequential(
             self._keras_model,
@@ -201,32 +128,7 @@ def populate_block(self, block, additional_options=None):
                 "KerasSurrogate.populate_block. Please pass a valid "
                 "formulation.".format(formulation)
             )
-        block.nn = OmltBlock()
-        block.nn.build_formulation(
-            formulation_object,
-        )
-
-        # input/output variables need to be constrained to be equal
-        # auto-created variables that come from OMLT.
-        input_idx_by_label = {s: i for i, s in enumerate(self._input_labels)}
-        input_vars_as_dict = block.input_vars_as_dict()
-
-        @block.Constraint(self._input_labels)
-        def input_surrogate_ties(m, input_label):
-            return (
-                input_vars_as_dict[input_label]
-                == block.nn.inputs[input_idx_by_label[input_label]]
-            )
-
-        output_idx_by_label = {s: i for i, s in enumerate(self._output_labels)}
-        output_vars_as_dict = block.output_vars_as_dict()
-
-        @block.Constraint(self._output_labels)
-        def output_surrogate_ties(m, output_label):
-            return (
-                output_vars_as_dict[output_label]
-                == block.nn.outputs[output_idx_by_label[output_label]]
-            )
+        self.populate_block_with_net(block, formulation_object)
 
     def evaluate_surrogate(self, inputs):
         """

idaes/core/surrogate/omlt_base_surrogate_class.py

@@ -0,0 +1,187 @@
+#################################################################################
+# The Institute for the Design of Advanced Energy Systems Integrated Platform
+# Framework (IDAES IP) was produced under the DOE Institute for the
+# Design of Advanced Energy Systems (IDAES).
+#
+# Copyright (c) 2018-2023 by the software owners: The Regents of the
+# University of California, through Lawrence Berkeley National Laboratory,
+# National Technology & Engineering Solutions of Sandia, LLC, Carnegie Mellon
+# University, West Virginia University Research Corporation, et al.
+# All rights reserved.  Please see the files COPYRIGHT.md and LICENSE.md
+# for full copyright and license information.
+#################################################################################
+"""
+Interface for importing ONNX models into IDAES
+"""
+# TODO: Missing docstrings
+# pylint: disable=missing-class-docstring
+# pylint: disable=missing-function-docstring
+
+from enum import Enum
+import numpy as np
+
+from pyomo.common.dependencies import attempt_import
+
+from idaes.core.surrogate.base.surrogate_base import SurrogateBase
+from idaes.core.surrogate.sampling.scaling import OffsetScaler
+
+keras, keras_available = attempt_import("tensorflow.keras")
+omlt, omlt_available = attempt_import("omlt")
+
+if omlt_available:
+    from omlt import OmltBlock, OffsetScaling
+
+
+class OMLTSurrogate(SurrogateBase):
+    def __init__(
+        self,
+        input_labels,
+        output_labels,
+        input_bounds,
+        input_scaler=None,
+        output_scaler=None,
+    ):
+        """
+        Standard SurrogateObject for surrogates based on Keras models.
+        Utilizes the OMLT framework for importing Keras models to IDAES.
+
+        Contains methods to both populate a Pyomo Block with constraints
+        representing the surrogate and to evaluate the surrogate a set of user
+        provided points.
+
+        This constructor should only be used when first creating the surrogate within IDAES.
+        Once created, this object can be stored to disk using save_to_folder and loaded
+        with load_from_folder
+
+        Args:
+           onnx_model: Onnx model file to be loaded.
+           input_labels: list of str
+              The ordered list of labels corresponding to the inputs in the keras model
+           output_labels: list of str
+              The ordered list of labels corresponding to the outputs in the keras model
+           input_bounds: None of dict of tuples
+              Keys correspond to each of the input labels and values are the tuples of
+              bounds (lb, ub)
+           input_scaler: None or OffsetScaler
+              The scaler to be used for the inputs. If None, then no scaler is used
+           output_scaler: None of OffsetScaler
+              The scaler to be used for the outputs. If None, then no scaler is used
+        """
+        super().__init__(
+            input_labels=input_labels,
+            output_labels=output_labels,
+            input_bounds=input_bounds,
+        )
+
+        # make sure we are using the standard scaler
+        if (
+            input_scaler is not None
+            and not isinstance(input_scaler, OffsetScaler)
+            or output_scaler is not None
+            and not isinstance(output_scaler, OffsetScaler)
+        ):
+            raise NotImplementedError("KerasSurrogate only supports the OffsetScaler.")
+
+        # check that the input labels match
+        if input_scaler is not None and input_scaler.expected_columns() != input_labels:
+            raise ValueError(
+                "KerasSurrogate created with input_labels that do not match"
+                " the expected columns in the input_scaler.\n"
+                "input_labels={}\n"
+                "input_scaler.expected_columns()={}".format(
+                    input_labels, input_scaler.expected_columns()
+                )
+            )
+
+        # check that the output labels match
+        if (
+            output_scaler is not None
+            and output_scaler.expected_columns() != output_labels
+        ):
+            raise ValueError(
+                "KerasSurrogate created with output_labels that do not match"
+                " the expected columns in the output_scaler.\n"
+                "output_labels={}\n"
+                "output_scaler.expected_columns()={}".format(
+                    output_labels, output_scaler.expected_columns()
+                )
+            )
+
+        self._input_scaler = input_scaler
+        self._output_scaler = output_scaler
+
+    class Formulation(Enum):
+        FULL_SPACE = 1
+        REDUCED_SPACE = 2
+        RELU_BIGM = 3
+        RELU_COMPLEMENTARITY = 4
+
+    def generate_omlt_scaling_objecets(self):
+        offset_inputs = np.zeros(self.n_inputs())
+        factor_inputs = np.ones(self.n_inputs())
+        offset_outputs = np.zeros(self.n_outputs())
+        factor_outputs = np.ones(self.n_outputs())
+        if self._input_scaler:
+            offset_inputs = self._input_scaler.offset_series()[
+                self.input_labels()
+            ].to_numpy()
+            factor_inputs = self._input_scaler.factor_series()[
+                self.input_labels()
+            ].to_numpy()
+        if self._output_scaler:
+            offset_outputs = self._output_scaler.offset_series()[
+                self.output_labels()
+            ].to_numpy()
+            factor_outputs = self._output_scaler.factor_series()[
+                self.output_labels()
+            ].to_numpy()
+
+        omlt_scaling = OffsetScaling(
+            offset_inputs=offset_inputs,
+            factor_inputs=factor_inputs,
+            offset_outputs=offset_outputs,
+            factor_outputs=factor_outputs,
+        )
+
+        # omlt takes *scaled* input bounds as a dictionary with int keys
+        input_bounds = dict(enumerate(self.input_bounds().values()))
+        scaled_input_bounds = omlt_scaling.get_scaled_input_expressions(input_bounds)
+        scaled_input_bounds = {i: tuple(bnd) for i, bnd in scaled_input_bounds.items()}
+        return omlt_scaling, scaled_input_bounds
+
+    def populate_block_with_net(self, block, formulation_object):
+        """
+        Method to populate a Pyomo Block with the omlt model constraints and build its formulation.
+
+        Args:
+           block: Pyomo Block component
+              The block to be populated with variables and/or constraints.
+           formulation_object: omlt loaded network formulation
+        """
+
+        block.nn = OmltBlock()
+        block.nn.build_formulation(
+            formulation_object,
+        )
+
+        # input/output variables need to be constrained to be equal
+        # auto-created variables that come from OMLT.
+        input_idx_by_label = {s: i for i, s in enumerate(self._input_labels)}
+        input_vars_as_dict = block.input_vars_as_dict()
+
+        @block.Constraint(self._input_labels)
+        def input_surrogate_ties(m, input_label):
+            return (
+                input_vars_as_dict[input_label]
+                == block.nn.inputs[input_idx_by_label[input_label]]
+            )
+
+        output_idx_by_label = {s: i for i, s in enumerate(self._output_labels)}
+        output_vars_as_dict = block.output_vars_as_dict()
+
+        @block.Constraint(self._output_labels)
+        def output_surrogate_ties(m, output_label):
+            return (
+                output_vars_as_dict[output_label]
+                == block.nn.outputs[output_idx_by_label[output_label]]
+            )