Add GNN-Based Predictor with DAG Preprocessing

pyproject.toml

@@ -44,6 +44,8 @@ dependencies = [
     "numpy>=1.24; python_version >= '3.11'",
     "numpy>=1.22",
     "numpy>=1.22,<2; sys_platform == 'darwin' and 'x86_64' in platform_machine and python_version < '3.13'",  # Restrict numpy v2 for macOS x86 since it is not supported anymore since torch v2.3.0
+    "optuna>=4.5.0",
+    "torch-geometric>=2.6.1",
     "torch>=2.7.1,<2.8.0; sys_platform == 'darwin' and 'x86_64' in platform_machine and python_version < '3.13'",  # Restrict torch v2.3.0 for macOS x86 since it is not supported anymore.
     "typing-extensions>=4.1",  # for `assert_never`
 ]
@@ -164,9 +166,15 @@ implicit_reexport = true
 # recent versions of `gym` are typed, but stable-baselines3 pins a very old version of gym.
 # qiskit is not yet marked as typed, but is typed mostly.
 # the other libraries do not have type stubs.
-module = ["qiskit.*", "joblib.*", "sklearn.*", "matplotlib.*", "gymnasium.*", "mqt.bench.*", "sb3_contrib.*", "bqskit.*", "qiskit_ibm_runtime.*", "networkx.*", "stable_baselines3.*"]
+module = ["qiskit.*", "joblib.*", "sklearn.*", "matplotlib.*", "gymnasium.*", "mqt.bench.*", "sb3_contrib.*", "bqskit.*", "qiskit_ibm_runtime.*", "networkx.*", "stable_baselines3.*", "torch", "torch.*", "torch_geometric", "torch_geometric.*", "optuna.*"]
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = ["mqt.predictor.ml.*"]
+disallow_subclassing_any = false
+
+
+
 [tool.ruff]
 line-length = 120
 extend-include = ["*.ipynb"]
@@ -245,6 +253,7 @@ wille = "wille"
 anc = "anc"
 aer = "aer"
 fom = "fom"
+TPE = "TPE"
 
 [tool.repo-review]
 ignore = ["GH200"]

src/mqt/predictor/_version.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2023 - 2025 Chair for Design Automation, TUM
+# Copyright (c) 2025 Munich Quantum Software Company GmbH
+# All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+#
+# Licensed under the MIT License
+
+# file generated by setuptools-scm
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__commit_id__",
+    "__version__",
+    "__version_tuple__",
+    "commit_id",
+    "version",
+    "version_tuple",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    VERSION_TUPLE = tuple[int | str, ...]
+    COMMIT_ID = str | None
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = "2.3.1.dev6+g1d835bd4c"
+__version_tuple__ = version_tuple = (2, 3, 1, "dev6", "g1d835bd4c")
+
+__commit_id__ = commit_id = None

src/mqt/predictor/hellinger/utils.py

@@ -132,12 +132,12 @@ def calc_device_specific_features(
     return np.array(list(feature_dict.values()))
 
 
-def get_hellinger_model_path(device: Target) -> Path:
+def get_hellinger_model_path(device: Target, gnn: bool = False) -> Path:
     """Returns the path to the trained model folder resulting from the machine learning training."""
-    training_data_path = Path(str(resources.files("mqt.predictor"))) / "ml" / "training_data"
+    training_data_path = Path(str(resources.files("mqt.predictor"))) / "ml" / "training_data" / "trained_model"
     model_path = (
-        training_data_path
-        / "trained_model"
-        / ("trained_hellinger_distance_regressor_" + device.description + ".joblib")
+        (training_data_path / ("trained_hellinger_distance_regressor_gnn_" + device.description + ".pth"))
+        if gnn
+        else (training_data_path / ("trained_hellinger_distance_regressor_" + device.description + ".joblib"))
     )
     return Path(model_path)

src/mqt/predictor/ml/__init__.py

@@ -13,4 +13,9 @@
 from mqt.predictor.ml import helper
 from mqt.predictor.ml.predictor import Predictor, predict_device_for_figure_of_merit, setup_device_predictor
 
-__all__ = ["Predictor", "helper", "predict_device_for_figure_of_merit", "setup_device_predictor"]
+__all__ = [
+    "Predictor",
+    "helper",
+    "predict_device_for_figure_of_merit",
+    "setup_device_predictor",
+]

src/mqt/predictor/ml/gnn.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2023 - 2025 Chair for Design Automation, TUM
+# Copyright (c) 2025 Munich Quantum Software Company GmbH
+# All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+#
+# Licensed under the MIT License
+
+"""This module contains the GNN module for graph neural networks."""
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+from torch_geometric.nn import SAGEConv, global_mean_pool
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,  # on 3.10+ prefer collections.abc
+    )
+
+    from torch_geometric.data import Data
+
+
+class GraphConvolutionSage(nn.Module):
+    """Graph convolutional layer using SAGEConv."""
+
+    def __init__(
+        self,
+        in_feats: int,
+        hidden_dim: int,
+        num_resnet_layers: int,
+        *,
+        conv_activation: Callable[..., torch.Tensor] = functional.leaky_relu,
+        conv_act_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        """A flexible SageConv graph classification model.
+
+        Args:
+          in_feats:        dimensionality of node features
+          hidden_dim:      output size of SageConv
+          num_resnet_layers:  how many SageConv layers (with residuals) to stack after the SageConvs
+          mlp_units:       list of units for each layer of the final MLP
+          conv_activation: activation fn after each graph layer
+          conv_act_kwargs: extra kwargs for conv_activation
+          final_activation: activation applied to the final scalar output
+        """
+        super().__init__()
+        self.conv_activation = conv_activation
+        self.conv_act_kwargs = conv_act_kwargs or {}
+
+        # --- GRAPH ENCODER ---
+        self.convs = nn.ModuleList()
+        # 1) Convolution not in residual configuration
+        # Possible to generalize the code
+        self.convs.append(SAGEConv(in_feats, hidden_dim))
+        self.convs.append(SAGEConv(hidden_dim, hidden_dim))
+
+        for _ in range(num_resnet_layers):
+            self.convs.append(SAGEConv(hidden_dim, hidden_dim))
+
+    def forward(self, data: Data) -> torch.Tensor:
+        """Forward function that allows to elaborate the input graph."""
+        x, edge_index, batch = data.x, data.edge_index, data.batch
+        # 1) Graph stack with residuals
+        for i, conv in enumerate(self.convs):
+            x_new = conv(x, edge_index)
+            x_new = self.conv_activation(x_new, **self.conv_act_kwargs)
+            # the number 2 is set because two convolution without residual configuration are applied
+            # and then all the others are in residual configuration
+            x = x_new if i < 2 else x + x_new
+
+        # 2) Global pooling
+        return global_mean_pool(x, batch)
+
+        # 3) MLP head
+
+
+class GNN(nn.Module):
+    """Architecture composed by a Graph Convolutional part with Sage Convolution module and followed by a MLP."""
+
+    def __init__(
+        self,
+        in_feats: int,
+        hidden_dim: int,
+        num_resnet_layers: int,
+        mlp_units: list[int],
+        *,
+        conv_activation: Callable[..., torch.Tensor] = functional.leaky_relu,
+        conv_act_kwargs: dict[str, Any] | None = None,
+        mlp_activation: Callable[..., torch.Tensor] = functional.leaky_relu,
+        mlp_act_kwargs: dict[str, Any] | None = None,
+        classes: list[str] | None = None,
+        output_dim: int = 1,
+    ) -> None:
+        """Init class for the GNN.
+
+        Arguments:
+            in_feats: dimension of input features of the node
+            hidden_dim: dimension of hidden output channels of the Convolutional part
+            num_resnet_layers: number of residual layers
+            mlp_units: list of units for each layer of the final MLP
+            conv_activation: activation fn after each graph layer
+            conv_act_kwargs: extra kwargs for conv_activation.
+            mlp_activation: activation fn after each MLP layer
+            mlp_act_kwargs: extra kwargs for mlp_activation.
+            output_dim: dimension of the output, default is 1 for regression tasks
+            classes: list of class names for classification tasks
+        """
+        # ─────────────────────────────────────────────────────────────────────────
+        # Suppress torch-geometric "plugin" import warnings (torch-scatter, etc.)
+        warnings.filterwarnings(
+            "ignore",
+            message=r"An issue occurred while importing 'torch-scatter'.*",
+            category=UserWarning,
+            module=r"torch_geometric.typing",
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=r"An issue occurred while importing 'torch-spline-conv'.*",
+            category=UserWarning,
+            module=r"torch_geometric.typing",
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=r"An issue occurred while importing 'torch-sparse'.*",
+            category=UserWarning,
+            module=r"torch_geometric.typing",
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=r"An issue occurred while importing 'torch-geometric'.*",
+            category=UserWarning,
+        )
+
+        warnings.filterwarnings(
+            "ignore",
+            message=r".*'type_params' parameter of 'typing\._eval_type'.*",
+            category=DeprecationWarning,
+        )
+
+        super().__init__()
+        # Convolutional part
+        self.graph_conv = GraphConvolutionSage(
+            in_feats, hidden_dim, num_resnet_layers, conv_activation=conv_activation, conv_act_kwargs=conv_act_kwargs
+        )
+
+        # MLP architecture
+        self.mlp_activation = mlp_activation
+        self.mlp_act_kwargs = mlp_act_kwargs or {}
+        self.classes = classes
+        self.fcs = nn.ModuleList()
+        last_dim = hidden_dim
+        for out_dim in mlp_units:
+            self.fcs.append(nn.Linear(last_dim, out_dim))
+            last_dim = out_dim
+        self.out = nn.Linear(last_dim, output_dim)
+
+    def forward(self, data: Data) -> torch.Tensor:
+        """Forward function that allows to elaborate the input graph.
+
+        Arguments:
+            data: The input graph data.
+        """
+        # apply the convolution
+        x = self.graph_conv(data)
+        # Apply the MLP
+        for fc in self.fcs:
+            x = self.mlp_activation(fc(x), **self.mlp_act_kwargs)
+        return self.out(x)