Solving the qubit maping problem with machine learning

import numpy as np
import torch
from torch_geometric.data import Data
def build_circuit_graph (circuit_data):
    # -----------------------------
    # STEP 1: Extract and scale single-qubit gate counts
    # Each function returns a vector of size (n_qubits,)
    # representing how often each gate was applied per qubit
    # -----------------------------
    rz_counts_scaled = get_single_qubit_counts(circuit_data, "rz")
    sx_counts_scaled = get_single_qubit_counts(circuit_data, "sx")
    x_counts_scaled = get_single_qubit_counts(circuit_data, "x")
    # -----------------------------
    # STEP 2: Extract and scale two-qubit gate counts
    # These describe interactions between pairs of qubits
    # -----------------------------
    ecr_scaled = get_two_qubit_counts(circuit_data,'ecr')
    # -----------------------------
    # STEP 3: Build node feature matrix
    # Each node = a logical qubit
    # Features = [rz count, sx count, x count]
    # Shape: (n_qubits, 3)
    # -----------------------------
    X =torch.tensor(
    list(zip(rz_counts_scaled, sx_counts_scaled, x_counts_scaled)),
    dtype=torch.float
    )
    # -----------------------------
    # STEP 4: Build graph connectivity (edges)
    # edge_index: which qubits are connected
    # edge_attr: weight/feature of each edge (scaled gate counts)
    # -----------------------------
    ecr_index, ecr_attr = edge_info(ecr_scaled)
    # -----------------------------
    # STEP 5: Create PyTorch Geometric graph object
    # x -> node features
    # edge_index -> connectivity (COO format)
    # edge_attr -> edge features
    # -----------------------------
    G_circuit = Data(
    x=X,
    edge_index=ecr_index,
    edge_attr=ecr_attr,
    )
    return G_circuit
def get_single_qubit_counts(data, gate_name):
    # Number of logical qubits in the circuit
    n = data["n_logical_qubits"]
    # Try to retrieve counts for the given gate type
    counts = data.get("single_qubit_counts", {}).get(gate_name, None)
    # If no data exists, return a zero vector
    if counts is None:
        return np.zeros(n)
    else:
    # Otherwise normalize counts across qubits (z-score normalization)
        return standard_scaling_across_qubits(np.array(counts))
def get_two_qubit_counts(data, gate_name):
    # Safely retrieve list of two-qubit gate interactions
    counts = data.get("two_qubit_counts", {}).get(gate_name, [])
    # If no interactions exist, return empty list
    if not counts:
        return counts
    else:
    # Otherwise scale the edge weights
      return scaled_two_qubit_list(counts)
def standard_scaling_across_qubits(vector):
    # Compute mean and standard deviation for normalization
    mean = vector.mean()
    std = vector.std()
    # Avoid division by zero if all values are identical
    if std == 0 :
        vector_scaled = vector
    else:
        vector_scaled = (vector - mean)/std
    return vector_scaled
def scaled_two_qubit_list(two_data):
    # Extract raw edge weights from dictionary list
    values = np.array([d['value']for d in two_data], dtype=float)
    # Normalize weights
    scaled_values = standard_scaling_across_qubits(values)
    # Reattach scaled values back into original edge dictionaries
    two_gate_scaled = [
    {**d, 'value':v_scaled}
    for d, v_scaled in zip(two_data, scaled_values)
    ]
    return two_gate_scaled
def edge_info(gate_dict):
    # Lists to store graph connectivity
    edge_index = []
    edge_attr = []
    # Convert list of edge dictionaries into tensor format
    for gate in gate_dict:
        i,j,w = gate["row"], gate["col"], gate["value"]
        edge_index.append([i,j])
        edge_attr.append([w])
    # Convert to PyTorch tensors
    # edge_index shape: [2, num_edges]
    edge_index = torch.tensor(edge_index, dtype=torch.long).T
    edge_attr = torch.tensor(edge_attr, dtype=torch.float)
    return edge_index, edge_attr

Graph_circuit = build_circuit_graph(ds.circuits[26])

import matplotlib.pyplot as plt
import networkx as nx
from torch_geometric.utils import to_networkx
# Convert PyG Data object to NetworkX graph
G = to_networkx(
    Graph_circuit,
    edge_attrs=["edge_attr"],  # keep CX counts
    node_attrs=["x"],          # keep node features
    to_undirected=False
)
edge_labels = {
    (u, v): f"{d['edge_attr'][0]:.2f}"
    for u, v, d in G.edges(data=True)
}
nx.draw(
    G,
    nx.spring_layout(G, seed=42),
    with_labels=True,
    node_size=800,
)
nx.draw_networkx_edge_labels(
    G,
    nx.spring_layout(G, seed=42),
    edge_labels=edge_labels
)
plt.show()

def build_hardware_graph(hardware_data):
    # Scaling single qubit properties and gate
    T1_scaled = standard_scaling_across_qubits(np.array(hardware_data['T1']))
    T2_scaled = standard_scaling_across_qubits(np.array(hardware_data['T2']))
    readout_error_scaled = standard_scaling_across_qubits(np.array(hardware_data['readout_error']))
    sx_error_scaled = standard_scaling_across_qubits(np.array(hardware_data['sx_error']))
    x_error_scaled = standard_scaling_across_qubits(np.array(hardware_data['x_error']))
    #Scaling two qubit gate properties
    ecr_error_scaled = scaled_two_qubit_list(hardware_data['ecr'])
    #Build the node feature matrix
    X = torch.tensor(
    list(zip(T1_scaled, T2_scaled, readout_error_scaled, sx_error_scaled, x_error_scaled)),
    dtype=torch.float
    )
    #Build the edge index and attributes matrix
    cx_index, cx_attr = edge_info(ecr_error_scaled)
    #Build the graph
    G_hardware = Data(
    x=X,
    edge_index=cx_index,
    edge_attr=cx_attr,
    )
    return G_hardware

Graph_hardware = build_hardware_graph(ds.hardware[26])

G_hardware = to_networkx(
    Graph_hardware,
    edge_attrs=["edge_attr"],  # keep CX counts
    node_attrs=["x"],          # keep node features
    to_undirected=False
)

plt.figure(figsize=(10, 10))
pos = nx.spring_layout(
    G_hardware,
    seed=42,
    k=1 / np.sqrt(len(G_hardware)),
    iterations=200
)
nx.draw_networkx_edges(G_hardware, pos, alpha=0.2, width=0.5)
nx.draw_networkx_nodes(G_hardware, pos, node_size=40)
plt.axis("off")
plt.savefig("graph_hardware.png", dpi=300, bbox_inches="tight")
plt.show()

import torch, torch.nn as nn, torch.nn.functional as F
from torch_geometric.data import Batch
from torch_geometric.nn import NNConv
from torch.utils.data import Dataset, DataLoader, random_split
from scipy.optimize import linear_sum_assignment

class QubitMapping_dot_Model(nn.Module):
    def __init__(self, node_c_dim, node_h_dim, edge_c_dim, edge_h_dim, hidden_dim):
        super().__init__()
        # === GNN for the LOGICAL CIRCUIT graph ===
        # These layers learn embeddings for logical qubits based on circuit structure
        # First layer: edge-conditioned convolution
        # Edge features (e.g., gate types) generate weights dynamically
        self.nn_c1 = make_edge_mlp(edge_c_dim, node_c_dim * hidden_dim)
        self.gnn_c1 = NNConv(node_c_dim, hidden_dim, nn=self.nn_c1, aggr='add')
        # Second layer: refines embeddings
        self.nn_c2 = make_edge_mlp(edge_c_dim, hidden_dim * hidden_dim)
        self.gnn_c2 = NNConv(hidden_dim, hidden_dim, nn=self.nn_c2, aggr='add')
        # === GNN for the HARDWARE graph ===
        # These layers learn embeddings for physical qubits based on hardware connectivity
        self.nn_h1 = make_edge_mlp(edge_h_dim, node_h_dim * hidden_dim)
        self.gnn_h1 = NNConv(node_h_dim, hidden_dim, nn=self.nn_h1, aggr='add')
        self.nn_h2 = make_edge_mlp(edge_h_dim, hidden_dim * hidden_dim)
        self.gnn_h2 = NNConv(hidden_dim, hidden_dim, nn=self.nn_h2, aggr='add')
    def forward(self, G_circuit, G_hardware):
        # === Step 1: Compute embeddings for LOGICAL qubits ===
        # Each node becomes a vector summarizing its role in the circuit
        x_c = F.relu(self.gnn_c1(
            G_circuit.x,              # node features
            G_circuit.edge_index,    # connectivity
            G_circuit.edge_attr      # edge features (e.g., gate info)
        ))
        h_L = F.relu(self.gnn_c2(
            x_c,
            G_circuit.edge_index,
            G_circuit.edge_attr
        ))
        # === Step 2: Compute embeddings for PHYSICAL qubits ===
        # Each node becomes a vector describing hardware properties
        x_h = F.relu(self.gnn_h1(
            G_hardware.x,
            G_hardware.edge_index,
            G_hardware.edge_attr
        ))
        h_P = F.relu(self.gnn_h2(
            x_h,
            G_hardware.edge_index,
            G_hardware.edge_attr
        ))
        # === Step 3: Compute compatibility between logical and physical qubits ===
        prob_list = []
        # Number of graphs in the batch
        num_graphs = G_circuit.batch.max().item() + 1
        # Loop over each graph in the batch
        for i in range(num_graphs):
            # --- Select nodes belonging to graph i ---
            # (because batching mixes all graphs together)
            mask_c = (G_circuit.batch == i)
            mask_h = (G_hardware.batch == i)
            # Extract embeddings for this specific graph
            h_L_i = h_L[mask_c]  # logical qubits
            h_P_i = h_P[mask_h]  # physical qubits
            # --- Compute compatibility matrix ---
            # Dot product measures how well each pair matches
            # Shape: [num_logical_qubits, num_physical_qubits]
            S_i = torch.matmul(h_L_i, h_P_i.T)
            if self.training:
                # During training:
                # return raw scores (used in loss function)
                prob_i = S_i
                prob_list.append(prob_i)
            else:
                # During inference:
                # find the BEST one-to-one assignment
                # Move to CPU + numpy for Hungarian algorithm
                S_cpu = S_i.detach().cpu().numpy()
                # Hungarian algorithm finds assignment maximizing total score
                row_ind, col_ind = linear_sum_assignment(-S_cpu)
                # Build mapping: logical qubit -> physical qubit
                mapping_i = torch.zeros(h_L_i.size(0), dtype=torch.long)
                mapping_i[row_ind] = torch.tensor(col_ind, dtype=torch.long)
                prob_list.append(mapping_i)
        # === Step 4: Merge results from all graphs ===
        if self.training:
            # Concatenate score matrices
            # Shape: [total_logical_nodes_in_batch, num_physical_nodes_per_graph]
            prob = torch.cat(prob_list, dim=0)
            return prob
        else:
            # Concatenate mappings
            # Shape: [total_logical_nodes_in_batch]
            mapping = torch.cat(prob_list, dim=0)
            return mapping

def make_edge_mlp(in_dim, out_dim, hid_dim=32):
    """
    Create a small multilayer perceptron (MLP) used to process edge features.

    Parameters
    ----------
    in_dim : int
        Dimension of the input edge feature vector.
    out_dim : int
        Dimension of the output representation (e.g., message, weight,
        or embedding associated with an edge).
    hid_dim : int, optional
        Number of neurons in the hidden layer (default: 32).

    Returns
    -------
    torch.nn.Sequential
        Two-layer feed-forward neural network:
            Linear(in_dim -> hid_dim)
            ReLU activation
            Linear(hid_dim -> out_dim)

    Notes
    -----
    In graph neural networks, this MLP is typically applied to edge
    attributes to generate edge embeddings or edge-dependent messages
    that are used during the message-passing/convolution step.
    """
    return nn.Sequential(
        nn.Linear(in_dim, hid_dim),
        nn.ReLU(),
        nn.Linear(hid_dim, out_dim)
    )

loss_fn = nn.CrossEntropyLoss()

def batch_graph_metrics_training(logit, qu_map, batch):
    """
    Compute per-graph metrics in a single pass:
      - circuit-averaged node accuracy
      - exact-mapping accuracy
    Args:
        logit: [N_total_nodes, N_P] predicted logit
        qu_map: [N_total_nodes] ground-truth physical qubits
        batch: [N_total_nodes] graph indices
    Returns:
        node_acc_sum: sum of node-level accuracies per graph
        exact_sum: number of perfectly mapped graphs
        num_graphs: number of graphs in the batch
    """
    pred = logit.argmax(dim=1)
    num_graphs = batch.max().item() + 1
    node_acc_sum = 0.0
    exact_sum = 0
    for g in range(num_graphs):
        mask = (batch == g)
        correct_nodes = (pred[mask] == qu_map[mask]).sum().item()
        n_nodes = mask.sum().item()
        node_acc_sum += correct_nodes / n_nodes
        if correct_nodes == n_nodes:
            exact_sum += 1
    return node_acc_sum, exact_sum, num_graphs
def batch_graph_metrics_val(pred, qu_map, batch):
    """
    Compute per-graph metrics in a single pass:
      - circuit-averaged node accuracy
      - exact-mapping accuracy
    Args:
        pred: [N_total_nodes] predicted mapping
        qu_map: [N_total_nodes] ground-truth physical qubits
        batch: [N_total_nodes] graph indices
    Returns:
        node_acc_sum: sum of node-level accuracies per graph
        exact_sum: number of perfectly mapped graphs
        num_graphs: number of graphs in the batch
    """
    num_graphs = batch.max().item() + 1
    node_acc_sum = 0.0
    exact_sum = 0
    for g in range(num_graphs):
        mask = (batch == g)
        correct_nodes = (pred[mask] == qu_map[mask]).sum().item()
        n_nodes = mask.sum().item()
        node_acc_sum += correct_nodes / n_nodes
        if correct_nodes == n_nodes:
            exact_sum += 1
    return node_acc_sum, exact_sum, num_graphs

def train_loop(model, optimizer, loss_fn, loader_train, loader_valid, n_epoch, device):
    """
    Training and validation loop for qubit mapping model.
    Returns:
        loss_hist_train, accuracy_hist_train, exact_accuracy_hist_train,
        accuracy_hist_valid, exact_accuracy_hist_valid
    """
    loss_hist_train = [0.0] * n_epoch
    accuracy_hist_train = [0.0] * n_epoch
    exact_accuracy_hist_train = [0.0] * n_epoch
    accuracy_hist_valid = [0.0] * n_epoch
    exact_accuracy_hist_valid = [0.0] * n_epoch
    for epoch in range(n_epoch):
        # ------------ TRAIN ------------------
        model.train()
        total_loss = 0.0
        total_acc = 0.0
        total_exact = 0
        total_graphs = 0
        for G_c, G_h,  qu_map in loader_train:
            G_c = G_c.to(device)
            G_h = G_h.to(device)
            qu_map = qu_map.to(device)
            optimizer.zero_grad()
            logit = model(G_c, G_h)
            loss = loss_fn(logit, qu_map)
            loss.backward()
            optimizer.step()
            #Metrics
            acc_sum, exact_sum, n_graphs = batch_graph_metrics_training(
                logit, qu_map, G_c.batch
            )
            total_loss += loss.item() * n_graphs
            total_acc += acc_sum
            total_exact += exact_sum
            total_graphs += n_graphs
        #Store epoch metrics
        loss_hist_train[epoch] = total_loss  / total_graphs
        accuracy_hist_train[epoch] = total_acc / total_graphs
        exact_accuracy_hist_train[epoch] = total_exact / total_graphs
        #VALIDATION
        model.eval()
        total_acc = 0.0
        total_exact = 0
        total_graphs = 0
        with torch.no_grad():
            for G_c_val, G_h_val, qu_map_val in loader_valid:
                G_c_val = G_c_val.to(device)
                G_h_val = G_h_val.to(device)
                qu_map_val = qu_map_val.to(device)
                mapping = model(G_c_val, G_h_val)
                acc_sum, exact_sum, n_graphs = batch_graph_metrics_val(
                mapping, qu_map_val, G_c_val.batch
                )
                total_acc += acc_sum
                total_exact += exact_sum
                total_graphs += n_graphs
        accuracy_hist_valid[epoch] = total_acc / total_graphs
        exact_accuracy_hist_valid[epoch] = total_exact / total_graphs
        print(
            f"Epoch {epoch+1:03d} | "
            f"Train Loss: {loss_hist_train[epoch]:.4f} | "
            f"Train Acc: {accuracy_hist_train[epoch]:.4f} | "
            f"Train Exact Acc: {exact_accuracy_hist_train[epoch]:.4f} | "
            f"Val Acc: {accuracy_hist_valid[epoch]:.4f} | "
            f"Val Exact Acc: {exact_accuracy_hist_valid[epoch]:.4f}"
        )
    return (loss_hist_train,
            accuracy_hist_train, exact_accuracy_hist_train,
            accuracy_hist_valid, exact_accuracy_hist_valid)