add graph neural networks

docs/notebooks.rst

@@ -14,6 +14,8 @@ github `page <https://github.com/cog-imperial/OMLT/tree/main/docs/notebooks/>`_.
 
 * `mnist_example_convolutional.ipynb <https://github.com/cog-imperial/OMLT/blob/main/docs/notebooks/neuralnet/mnist_example_convolutional.ipynb>`_ trains a convolutional neural network on MNIST and uses OMLT to find adversarial examples.
 
+* `graph_neural_network_formulation.ipynb <https://github.com/cog-imperial/OMLT/blob/main/docs/notebooks/graph_neural_network_formulation.ipynb>`_ transforms graph neural networks into OMLT and builds formulation to solve optimization problems.
+
 * `auto-thermal-reformer.ipynb <https://github.com/cog-imperial/OMLT/blob/main/docs/notebooks/neuralnet/auto-thermal-reformer.ipynb>`_ develops a neural network surrogate (using sigmoid activations) with data from a process model built using `IDAES-PSE <https://github.com/IDAES/idaes-pse>`_.
 
 * `auto-thermal-reformer-relu.ipynb <https://github.com/cog-imperial/OMLT/blob/main/docs/notebooks/neuralnet/auto-thermal-reformer-relu.ipynb>`_ develops a neural network surrogate (using ReLU activations) with data from a process model built using `IDAES-PSE <https://github.com/IDAES/idaes-pse>`_.

setup.cfg

@@ -88,6 +88,7 @@ testing =
     torchvision
     tqdm
     protobuf==3.20.3
+    torch_geometric
 
 testing_lean =
     setuptools

src/omlt/dependencies.py

@@ -4,4 +4,7 @@
 onnx, onnx_available = attempt_import("onnx")
 keras, keras_available = attempt_import("tensorflow.keras")
 
+torch, torch_available = attempt_import("torch")
+
+torch_geometric, torch_geometric_available = attempt_import("torch_geometric")
 lineartree, lineartree_available = attempt_import("lineartree")

src/omlt/io/__init__.py

@@ -1,4 +1,9 @@
-from omlt.dependencies import keras_available, onnx_available
+from omlt.dependencies import (
+    onnx_available,
+    keras_available,
+    torch_available,
+    torch_geometric_available,
+)
 
 if onnx_available:
     from omlt.io.onnx import (

src/omlt/io/torch_geometric/__init__.py

@@ -0,0 +1,8 @@
+from omlt.io.torch_geometric.torch_geometric_reader import (
+    load_torch_geometric_sequential,
+)
+
+from omlt.io.torch_geometric.build_gnn_formulation import (
+    gnn_with_fixed_graph,
+    gnn_with_non_fixed_graph,
+)

src/omlt/io/torch_geometric/build_gnn_formulation.py

@@ -0,0 +1,148 @@
+import numpy as np
+import pyomo.environ as pyo
+from omlt.neuralnet import FullSpaceNNFormulation
+from omlt.io.torch_geometric import load_torch_geometric_sequential
+
+
+def gnn_with_non_fixed_graph(
+    block,
+    nn,
+    N,
+    scaling_object=None,
+    scaled_input_bounds=None,
+    unscaled_input_bounds=None,
+):
+    """
+    Build formulation for a torch_geometric graph neural network model (built with Sequential).
+    Since the input graph is not fixed, the elements in adjacency matrix are decision variables.
+
+    Parameters
+    ----------
+    block : Block
+        the Pyomo block
+    nn : torch_geometric.model
+        A torch_geometric model that was built with Sequential
+    N : int
+        The number of nodes of input graph
+    scaling_object : instance of ScalingInterface or None
+        Provide an instance of a scaling object to use to scale iputs --> scaled_inputs
+        and scaled_outputs --> outputs. If None, no scaling is performed. See scaling.py.
+    scaled_input_bounds : dict or None
+        A dict that contains the bounds on the scaled variables (the
+        direct inputs to the neural network). If None, then no bounds
+        are specified or they are generated using unscaled bounds.
+    unscaled_input_bounds : dict or None
+        A dict that contains the bounds on the unscaled variables (the
+        direct inputs to the neural network). If specified the scaled_input_bounds
+        dictionary will be generated using the provided scaling object.
+        If None, then no bounds are specified.
+
+    Returns
+    -------
+    OmltBlock (formulated)
+    """
+
+    # build NetworkDefinition for nn
+    net = load_torch_geometric_sequential(
+        nn=nn,
+        N=N,
+        A=None,
+        scaling_object=scaling_object,
+        scaled_input_bounds=scaled_input_bounds,
+        unscaled_input_bounds=unscaled_input_bounds,
+    )
+
+    # define binary variables for adjacency matrix
+    block.A = pyo.Var(
+        pyo.Set(initialize=range(N)),
+        pyo.Set(initialize=range(N)),
+        within=pyo.Binary,
+    )
+    # assume that the self contribution always exists
+    for u in range(N):
+        block.A[u, u].fix(1)
+    # assume the adjacency matrix is always symmetric
+    block.symmetric_adjacency = pyo.ConstraintList()
+    for u in range(N):
+        for v in range(u + 1, N):
+            block.symmetric_adjacency.add((block.A[u, v] == block.A[v, u]))
+
+    # build formulation for GNN
+    block.build_formulation(FullSpaceNNFormulation(net))
+
+    return block
+
+
+def gnn_with_fixed_graph(
+    block,
+    nn,
+    N,
+    A,
+    scaling_object=None,
+    scaled_input_bounds=None,
+    unscaled_input_bounds=None,
+):
+    """
+    Build formulation for a torch_geometric graph neural network model (built with Sequential).
+    Given the adjacency matrix, the input graph structure is fixed.
+
+    Parameters
+    ----------
+    block : Block
+        the Pyomo block
+    nn : torch_geometric.model
+        A torch_geometric model that was built with Sequential
+    N : int
+        The number of nodes of input graph
+    A : matrix-like
+        The adjacency matrix of input graph
+    scaling_object : instance of ScalingInterface or None
+        Provide an instance of a scaling object to use to scale iputs --> scaled_inputs
+        and scaled_outputs --> outputs. If None, no scaling is performed. See scaling.py.
+    scaled_input_bounds : dict or None
+        A dict that contains the bounds on the scaled variables (the
+        direct inputs to the neural network). If None, then no bounds
+        are specified or they are generated using unscaled bounds.
+    unscaled_input_bounds : dict or None
+        A dict that contains the bounds on the unscaled variables (the
+        direct inputs to the neural network). If specified the scaled_input_bounds
+        dictionary will be generated using the provided scaling object.
+        If None, then no bounds are specified.
+
+    Returns
+    -------
+    OmltBlock (formulated)
+    """
+
+    # assume the adjacency matrix is always symmetric
+    assert np.array_equal(A, np.transpose(A))
+
+    # build NetworkDefinition for nn
+    net = load_torch_geometric_sequential(
+        nn=nn,
+        N=N,
+        A=A,
+        scaling_object=scaling_object,
+        scaled_input_bounds=scaled_input_bounds,
+        unscaled_input_bounds=unscaled_input_bounds,
+    )
+
+    # define binary variables for adjacency matrix
+    block.A = pyo.Var(
+        pyo.Set(initialize=range(N)),
+        pyo.Set(initialize=range(N)),
+        within=pyo.Binary,
+    )
+    # fix A using given values
+    for u in range(N):
+        for v in range(N):
+            block.A[u, v].fix(A[u, v])
+
+    # assume that the self contribution always exists
+    for u in range(N):
+        block.A[u, u].fix(1)
+
+    # build formulation for GNN
+    block.build_formulation(FullSpaceNNFormulation(net))
+
+    return block