Add GNN-Based Predictor with DAG Preprocessing #430

antotu · 2025-08-19T16:15:44Z

Description

This PR introduces a Graph Neural Network (GNN) as an alternative to the Random Forest model for predicting the best device to run a quantum circuit.
To support this, the preprocessing pipeline was redesigned: instead of manually extracting features from the circuit, the model now directly takes as input the Directed Acyclic Graph (DAG) representation of the quantum circuit.

🚀 Major Changes

Graph Neural Network Integration

Added a GNN model for predicting the target quantum device and estimating the Hellinger distance between output distributions.
Added a preprocessing method to transform quantum circuits into DAGs.
DAG representation captures gate dependencies and circuit topology for improved graph-based learning.
Integrated automated hyperparameter search with Optuna for tuning GNN performance.

🎯 Motivation

Previously, features were manually extracted from the quantum circuit, leading to loss of structural information.
This new method preserves the full circuit structure by representing it as a graph.
GNNs can exploit graph connectivity to make more accurate predictions.
Optuna ensures that GNN hyperparameters are efficiently optimized in a reproducible way.

🔧 Fixes and Enhancements

Transform input quantum circuits into DAGs, where each node is encoded as a numeric vector.
Integrated GNNs as an additional predictor in the pipeline.

📦 Dependency Updates

optuna>=4.5.0
torch-geometric>=2.6.1

Checklist:

The pull request only contains commits that are focused and relevant to this change.
I have added appropriate tests that cover the new/changed functionality.
I have updated the documentation to reflect these changes.
I have added entries to the changelog for any noteworthy additions, changes, fixes, or removals.
I have added migration instructions to the upgrade guide (if needed).
The changes follow the project's style guidelines and introduce no new warnings.
The changes are fully tested and pass the CI checks.
I have reviewed my own code changes.

…and test all done

src/mqt/predictor/_version.py

+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    VERSION_TUPLE = tuple[int | str, ...]


src/mqt/predictor/ml/helper.py

…into gnn-branch

…L e GNN

src/mqt/predictor/ml/helper.py

flowerthrower

Hey @antotu , thanks for your continued efforts!
I still didn't manage to get fully through, so here is another preliminary batch of feedback.

flowerthrower · 2025-08-27T07:22:46Z

src/mqt/predictor/ml/gnn.py

+        # 2) Global pooling
+        return global_mean_pool(x, batch)
+
+        # 3) MLP head


Suggested change

# 3) MLP head

tests/device_selection/test_predictor_ml.py

flowerthrower · 2025-08-27T07:32:51Z

pyproject.toml

+
+

Suggested change

flowerthrower · 2025-08-27T07:39:41Z

src/mqt/predictor/ml/predictor.py

+warnings.filterwarnings(
+    "ignore",
+    message=r"An issue occurred while importing 'torch-scatter'.*",
+    category=UserWarning,


pytorch.*:UserWarning should already be ignored through the filterwarnings in pyproject.toml. Please only add the additionally required ones there. The same goes for the other files, too. Thanks!

flowerthrower · 2025-08-27T07:46:14Z

src/mqt/predictor/ml/predictor.py

+        number_epochs: The number of epochs to train the GNN model. Defaults to 100.
+        number_trials: The number of trials to run for hyperparameter optimization for the GNN. Defaults to 50.
+        verbose: Whether to print verbose output during training GNN. Defaults to False.


Suggested change

number_epochs: The number of epochs to train the GNN model. Defaults to 100.

number_trials: The number of trials to run for hyperparameter optimization for the GNN. Defaults to 50.

verbose: Whether to print verbose output during training GNN. Defaults to False.

**gnn_kwargs: Forwarded to `Predictor.train_gnn_model` when `gnn=True`

(e.g., `number_epochs=100`, `number_trials=50`, `verbose=False`).

flowerthrower · 2025-08-27T07:47:49Z

src/mqt/predictor/ml/predictor.py

+    number_epochs: int = 100,
+    number_trials: int = 50,
+    verbose: bool = False,


Perhaps we can use a gnn_kwargs dictionary here to avoid cluttering the arguments with only GNN-specific things. It could also be useful in the future to add more hyperparameters if needed.

Co-authored-by: Patrick Hopf <[email protected]> Signed-off-by: Antonio Tudisco <[email protected]>

flowerthrower

This is just another batch of feedback. Thank you for integrating the requested changes so fast!

flowerthrower · 2025-08-27T09:05:25Z

src/mqt/predictor/ml/helper.py



+def create_dag(qc: QuantumCircuit) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Creates and returns the associate DAG of the quantum circuit.


Suggested change

"""Creates and returns the associate DAG of the quantum circuit.

"""Creates and returns the feature-annotated DAG of the quantum circuit.

flowerthrower · 2025-08-27T09:11:09Z

src/mqt/predictor/ml/helper.py


 from __future__ import annotations

+import math


Numpy is already imported and provides the same functionality. If I remember correctly, in a similar way, PyTorch provides these basic things too (perhaps we can further reduce imports here).

flowerthrower · 2025-08-27T10:10:41Z

src/mqt/predictor/ml/helper.py

+    return_arrays: bool = False,
+    verbose: bool = False,
+) -> tuple[float, dict[str, float], tuple[np.ndarray, np.ndarray] | None]:
+    """Evaluate the models.


Can we make the description a bit more detailed? Just so we know why this is necessary and that it is only required for the GNN models.

flowerthrower · 2025-08-27T10:13:39Z

src/mqt/predictor/ml/helper.py

+    restore_best: bool = True,
+    scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
+) -> None:
+    """Trains the model with optional early stopping on validation loss.


Suggested change

"""Trains the model with optional early stopping on validation loss.

"""Trains a GNN model with optional early stopping on validation loss.

src/mqt/predictor/ml/helper.py

flowerthrower · 2025-08-27T11:05:29Z

src/mqt/predictor/ml/predictor.py

        qc = QuantumCircuit.from_qasm_file(path_uncompiled_circuit / file)
-        feature_vec = create_feature_vector(qc)
-        training_sample = (feature_vec, target_label)
+        if not self.gnn:


Suggested change

if not self.gnn:

if self.gnn:

x, edge_index, number_of_gates = create_dag(qc)

y = torch.tensor([[dev.description for dev in self.devices].index(target_label)], dtype=torch.float)

training_sample = (x, y, edge_index, number_of_gates, target_label)

else:

feature_vec = create_feature_vector(qc)

training_sample = (feature_vec, target_label)

circuit_name = str(file).split(".")[0]

return training_sample, circuit_name, scores_list

flowerthrower · 2025-08-27T11:05:55Z

src/mqt/predictor/ml/predictor.py

+            feature_vec = create_feature_vector(qc)
+            training_sample = (feature_vec, target_label)
+            circuit_name = str(file).split(".")[0]
+            return training_sample, circuit_name, scores_list
+        x, edge_index, number_of_gates = create_dag(qc)


Suggested change

feature_vec = create_feature_vector(qc)

training_sample = (feature_vec, target_label)

circuit_name = str(file).split(".")[0]

return training_sample, circuit_name, scores_list

x, edge_index, number_of_gates = create_dag(qc)

flowerthrower · 2025-08-27T11:06:35Z

src/mqt/predictor/ml/predictor.py

+        self.devices_description = [dev.description for dev in self.devices]
+        y = self.devices_description.index(target_label)
+        print(target_label)
+        return Data(
+            x=x,
+            y=torch.tensor([y], dtype=torch.float),
+            circuit_name=circuit_name,
+            edge_index=edge_index,
+            target_label=target_label,  # torch.tensor([target_label], dtype=torch.float),
+            scores_list=scores_list,
+            num_nodes=number_of_gates,
+        )


Suggested change

self.devices_description = [dev.description for dev in self.devices]

y = self.devices_description.index(target_label)

print(target_label)

return Data(

x=x,

y=torch.tensor([y], dtype=torch.float),

circuit_name=circuit_name,

edge_index=edge_index,

target_label=target_label, # torch.tensor([target_label], dtype=torch.float),

scores_list=scores_list,

num_nodes=number_of_gates,

)

src/mqt/predictor/ml/predictor.py

flowerthrower · 2025-08-27T11:13:43Z

src/mqt/predictor/ml/predictor.py


        return mdl.best_estimator_

+    def _get_prepared_training_graphs(self) -> TrainingData:


With the changes above, we can drop this graph-specific method and instead use the _get_prepared_training_data with a slight modification when loading the graph-specific training data (if self.gnn: ...).

Co-authored-by: Patrick Hopf <[email protected]> Signed-off-by: Antonio Tudisco <[email protected]>

…-branch

src/mqt/predictor/ml/helper.py

+        train_loss = running_loss / max(1, total)
+        if scheduler is not None:
+            scheduler.step()
+        val_loss = float("inf")


…-branch

antotu added 2 commits August 19, 2025 13:47

added function related to training and for GNN, needed to define GNN …

12fad57

…and test all done

Added the gnn part, must be fine-tuned hyper-params, no test

be734dd

antotu marked this pull request as draft August 19, 2025 16:16

github-advanced-security bot found potential problems Aug 19, 2025

View reviewed changes

antotu and others added 25 commits August 19, 2025 18:25

Removed the barriers in the creation of the DAG

61c6824

🎨 pre-commit fixes

75875ff

coded tested and fixed, need to add a cross validation module

081651c

Merge branch 'gnn-branch' of https://github.com/antotu/predictor-gnn …

b82dc01

…into gnn-branch

fixed the problem of the predict_device_for_figure_of_merits

5ebd202

🎨 pre-commit fixes

857cd6f

Hellinger test done: success

6081f6b

Merge branch 'gnn-branch' of https://github.com/antotu/predictor-gnn …

7c54da6

…into gnn-branch

GNN predictor fixed with optuna and tested

bb4da24

🎨 pre-commit fixes

10bb52c

GNN predictor fixed with optuna and tested

06be0d6

Modified the tolm for running on the MacOS

ce990e3

Problems modified TPESampler and not TYPESampler

96ca75b

Problems modified TPESampler and not TYPESampler

a64a082

🎨 pre-commit fixes

f8c99b5

Problems modified TPESampler and not TYPESampler

e4e2742

Problems modified TPESampler and not TYPESampler

5784ff7

Test modified with number of epochs as parameter

7e17379

Eliminated trained model

082de05

Changed the test estimated hellinger for windows

5ed00a9

🎨 pre-commit fixes

e59a941

Changed the test estimated hellinger for windows

3a9f16c

Merge branch 'gnn-branch' of https://github.com/antotu/predictor-gnn …

c43ee01

…into gnn-branch

Changed the test estimated hellinger for windows

92eda99

Problem with windows solved eliminating warning

dc1aa55

antotu changed the title ~~Gnn branch~~ Add GNN-Based Predictor with DAG Preprocessing Aug 21, 2025

antotu and others added 6 commits August 25, 2025 10:37

Fixed the comments related to test hellinger distance and utils

8c77598

🎨 pre-commit fixes

dc0a824

Fixed modification also with pre-commit

2419952

Fixed modification also with pre-commit

5335241

Refactor the test ml predictor considering to join function related M…

96096a0

…L e GNN

Modified part of helper in order to solve problems code

4613012

github-advanced-security bot found potential problems Aug 26, 2025

View reviewed changes

src/mqt/predictor/ml/helper.py Fixed Show fixed Hide fixed

Pre-commit has substituted Wille in Will

1c728e2

flowerthrower reviewed Aug 27, 2025

View reviewed changes

antotu and others added 2 commits August 27, 2025 13:14

Update tests/device_selection/test_predictor_ml.py

c31cb46

Co-authored-by: Patrick Hopf <[email protected]> Signed-off-by: Antonio Tudisco <[email protected]>

🎨 pre-commit fixes

13cf0f4

flowerthrower reviewed Aug 27, 2025

View reviewed changes

antotu and others added 9 commits August 27, 2025 13:23

first round fixes

713343f

🎨 pre-commit fixes

17c6575

pre-commit fixes

169b00e

pre-commit fixes

2248081

Update src/mqt/predictor/ml/predictor.py

8f90b12

Co-authored-by: Patrick Hopf <[email protected]> Signed-off-by: Antonio Tudisco <[email protected]>

Update src/mqt/predictor/ml/predictor.py

74ec34b

Co-authored-by: Patrick Hopf <[email protected]> Signed-off-by: Antonio Tudisco <[email protected]>

🎨 pre-commit fixes

f99e17b

Partial modification

57b1a29

Merge branch 'gnn-branch' of github.com:antotu/predictor-gnn into gnn…

96232f0

…-branch

github-advanced-security bot found potential problems Aug 27, 2025

View reviewed changes

src/mqt/predictor/ml/helper.py Fixed Show fixed Hide fixed

🎨 pre-commit fixes

61965d8

github-advanced-security bot found potential problems Aug 27, 2025

View reviewed changes

antotu added 6 commits August 27, 2025 16:22

fixed comments repo

93f5414

Merge branch 'gnn-branch' of github.com:antotu/predictor-gnn into gnn…

95f5359

…-branch

Modified the gates accepted

4fb7112

Modified list

5ea1720

Fixed bug Swap and Cswap gates

312e6ea

Edit for saving memory GPU

156b7e6



		def create_dag(qc: QuantumCircuit) -> tuple[torch.Tensor, torch.Tensor, int]:
		"""Creates and returns the associate DAG of the quantum circuit.

	"""Creates and returns the associate DAG of the quantum circuit.
	"""Creates and returns the feature-annotated DAG of the quantum circuit.

	"""Trains the model with optional early stopping on validation loss.
	"""Trains a GNN model with optional early stopping on validation loss.


		return mdl.best_estimator_

		def _get_prepared_training_graphs(self) -> TrainingData:

Uh oh!

Add GNN-Based Predictor with DAG Preprocessing #430

Are you sure you want to change the base?

Add GNN-Based Predictor with DAG Preprocessing #430

Uh oh!

Conversation

antotu commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

🚀 Major Changes

Graph Neural Network Integration

🎯 Motivation

🔧 Fixes and Enhancements

📦 Dependency Updates

Checklist:

Uh oh!

Check warning

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

flowerthrower left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

flowerthrower left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Check warning

Uh oh!

Uh oh!

antotu commented Aug 19, 2025 •

edited

Loading