{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":3004,"databundleVersionId":861823},{"sourceType":"datasetVersion","sourceId":242592,"datasetId":102285,"databundleVersionId":254413}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"*This notebook englobes the steps defined in the quantum self attention network (QSAN) from step 1 to step 9*","metadata":{}},{"cell_type":"markdown","source":"# Installing required lilbraries","metadata":{}},{"cell_type":"code","source":"pip install qiskit pylatexenc qiskit-aer qiskit-ibm-runtime pennylane","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:16:42.685953Z","iopub.execute_input":"2026-05-11T16:16:42.686213Z","iopub.status.idle":"2026-05-11T16:17:02.357958Z","shell.execute_reply.started":"2026-05-11T16:16:42.686185Z","shell.execute_reply":"2026-05-11T16:17:02.356788Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister\nfrom qiskit.quantum_info import Statevector\nfrom qiskit.visualization import circuit_drawer\nfrom IPython.display import display, Latex\nimport numpy as np\nfrom fractions import Fraction\nimport pennylane as qml\nfrom pennylane import numpy as np","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:02.360181Z","iopub.execute_input":"2026-05-11T16:17:02.360666Z","iopub.status.idle":"2026-05-11T16:17:07.622208Z","shell.execute_reply.started":"2026-05-11T16:17:02.360630Z","shell.execute_reply":"2026-05-11T16:17:07.621397Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1.Circuit Preparation ","metadata":{}},{"cell_type":"markdown","source":"# Initialization","metadata":{}},{"cell_type":"code","source":"def circuit_def(n, m):\n    input_wires = list(range(n * m))\n    query_wires = list(range(n * m, 2 * n * m))\n    key_wires = list(range(2 * n * m, 3 * n * m))\n    qls_size = m * sum(range(1, n + 1))\n    qls_wires = list(range(3 * n * m, 3 * n * m + qls_size))\n    total_wires = 3 * n * m + qls_size\n    return input_wires, query_wires, key_wires, qls_wires, total_wires","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.623842Z","iopub.execute_input":"2026-05-11T16:17:07.624592Z","iopub.status.idle":"2026-05-11T16:17:07.629395Z","shell.execute_reply.started":"2026-05-11T16:17:07.624559Z","shell.execute_reply":"2026-05-11T16:17:07.628500Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set parameters\nn = 2\nm = 2\ninput_wires, query_wires, key_wires, qls_wires, total_wires = circuit_def(n, m)\n\n# Define PennyLane device\ndev = qml.device('lightning.qubit', wires=total_wires)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.630145Z","iopub.execute_input":"2026-05-11T16:17:07.630362Z","iopub.status.idle":"2026-05-11T16:17:07.694745Z","shell.execute_reply.started":"2026-05-11T16:17:07.630345Z","shell.execute_reply":"2026-05-11T16:17:07.693799Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 1 Query Input","metadata":{}},{"cell_type":"code","source":"def apply_u(n, m, input_wires, params):\n \n    # Step 1: Encode with H and RY gates on input_wires\n    for i in range(n * m):\n        qml.Hadamard(wires=input_wires[i])\n    \n    # Apply RY gates with random angles\n    for i in range(n * m):\n        qml.RY(params[i], wires=input_wires[i])\n    \n    # Step 2: Entangle qubits within each m-qubit block\n    for j in range(n):  # For each block\n        for k in range(m - 1):  # Apply CNOT between consecutive qubits in the block\n            control = input_wires[j * m + k]\n            target = input_wires[j * m + k + 1]\n            qml.CNOT(wires=[control, target])\n    \n    # Add barrier (visual separator in PennyLane circuit drawing)\n    # qml.Barrier(wires=total_wires)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.696524Z","iopub.execute_input":"2026-05-11T16:17:07.697152Z","iopub.status.idle":"2026-05-11T16:17:07.701618Z","shell.execute_reply.started":"2026-05-11T16:17:07.697132Z","shell.execute_reply":"2026-05-11T16:17:07.700968Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 2 : Query Swap","metadata":{}},{"cell_type":"code","source":"def barbell_operation(n, m, input_wires, query_wires, key_wires, step='q'):\n\n    if step == 'q':\n        for r in range(m * n):\n            qml.SWAP(wires=[input_wires[r], query_wires[r]])  # Swap input[r] with query[r]\n        # qml.Barrier(wires=input_wires + query_wires)\n    \n    # Barbell operation for |K⟩ (Step 4)\n    if step == 'k':\n        for r in range(m * n):\n            qml.SWAP(wires=[input_wires[r], key_wires[r]])  # Swap input[r] with key[r]\n        # qml.Barrier(wires=input_wires + key_wires)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.702436Z","iopub.execute_input":"2026-05-11T16:17:07.702681Z","iopub.status.idle":"2026-05-11T16:17:07.717578Z","shell.execute_reply.started":"2026-05-11T16:17:07.702659Z","shell.execute_reply":"2026-05-11T16:17:07.716758Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 5 : QLS","metadata":{}},{"cell_type":"code","source":"def qls(n, m, query_wires, key_wires, qls_wires):\n\n    # Apply Toffoli gates between (q_i, k_j) pairs for j <= n-1-i\n    qls_idx = 0  # Track qls wire index\n    for q in range(n):  # For each query Q_i\n        for k in range(n - q):  # Compare with K_j where j <= n-1-i\n            for j in range(m):  # For each qubit in the m-dimensional input\n                q_idx = query_wires[q * m + j]  # Q_i's j-th wire\n                k_idx = key_wires[k * m + j]  # K_j's j-th wire\n                qml.Toffoli(wires=[q_idx, k_idx, qls_wires[qls_idx]])\n                qls_idx += 1\n    # qml.Barrier(wires=query_wires + key_wires + qls_wires)\n    \n    # Apply CNOT gates between consecutive positions within each m-qubit group in qls\n    qls_idx = 0  # Reset qls index\n    for q in range(n):  # For each query Q_i\n        for k in range(n - q):  # For each key K_j where j <= n-1-i\n            for j in range(m - 1):  # Within each m-qubit group, apply CNOT\n                qml.CNOT(wires=[qls_wires[qls_idx + j], qls_wires[qls_idx + j + 1]])\n            qls_idx += m  # Move to the next m-qubit group\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.718392Z","iopub.execute_input":"2026-05-11T16:17:07.718875Z","iopub.status.idle":"2026-05-11T16:17:07.740220Z","shell.execute_reply.started":"2026-05-11T16:17:07.718850Z","shell.execute_reply":"2026-05-11T16:17:07.739704Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 7 : Compression","metadata":{}},{"cell_type":"code","source":"def compression_operation(n, m, input_wires):\n    \n    # Apply CNOT gates for compression\n    for j in range(n - 1):  # Loop over consecutive pairs (j, j+1)\n        for i in range(m):  # For each qubit index i in the block\n            control = input_wires[j * m + i]       # w_{j,i}\n            target = input_wires[(j + 1) * m + i]  # w_{j+1,i}\n            qml.CNOT(wires=[control, target])\n    # qml.Barrier(wires=input_wires)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.740854Z","iopub.execute_input":"2026-05-11T16:17:07.741066Z","iopub.status.idle":"2026-05-11T16:17:07.754319Z","shell.execute_reply.started":"2026-05-11T16:17:07.741040Z","shell.execute_reply":"2026-05-11T16:17:07.753774Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Step 8 : Slicing Operation","metadata":{}},{"cell_type":"code","source":"def get_control_qubits(index, q0, m, n):\n    control_bits = []\n    num_comparisons = n - q0  \n    current_index = index\n    for j2 in range(num_comparisons):\n        control_bit = current_index + 2\n        control_bits.append(control_bit)\n        current_index = control_bit\n    return control_bits, current_index","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.755038Z","iopub.execute_input":"2026-05-11T16:17:07.755216Z","iopub.status.idle":"2026-05-11T16:17:07.773566Z","shell.execute_reply.started":"2026-05-11T16:17:07.755202Z","shell.execute_reply":"2026-05-11T16:17:07.772940Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def slicing_operation(n, m, query_wires, qls_wires):\n\n    target = n * m  # Start of query_wires (4)\n    index = 3 * m * n - 1  # Last index of key_wires (11)\n    for q in range(n):\n        control_bits, new_index = get_control_qubits(index, q, m, n)\n        # Reset operation (PennyLane lacks direct reset; use Identity as placeholder)\n        qml.StatePrep([1.0, 0.0], wires=query_wires[q*m])\n        # Multi-controlled X gate (combine control and target wires)\n        qml.MultiControlledX(wires=control_bits + [query_wires[q * m]])\n        # qml.Barrier(wires=query_wires + qls_wires)\n        index = new_index\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.774394Z","iopub.execute_input":"2026-05-11T16:17:07.774986Z","iopub.status.idle":"2026-05-11T16:17:07.787983Z","shell.execute_reply.started":"2026-05-11T16:17:07.774969Z","shell.execute_reply":"2026-05-11T16:17:07.787458Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Full Circuit","metadata":{}},{"cell_type":"markdown","source":"Example input for the circuit","metadata":{}},{"cell_type":"code","source":"import torch\nparams = [np.pi, np.pi/3, np.pi/5, np.pi/6, np.pi, np.pi/3, np.pi/5, np.pi/6, np.pi, np.pi/3,np.pi/5, np.pi/6]\n# Define a synthetic 4x4 image (e.g., resembling a simplified MNIST digit \"1\")\nimage_4x4 = np.array([\n    [0.1, 0.2, 0.2, 0.1],  # Row 1: Low intensity background\n    [0.2, 0.8, 0.8, 0.2],  # Row 2: High intensity for vertical stroke\n    [0.2, 0.8, 0.8, 0.2],  # Row 3: High intensity for vertical stroke\n    [0.1, 0.2, 0.2, 0.1]   # Row 4: Low intensity background\n])\n\n# Flatten the 4x4 image to a 16-dimensional vector\nimage_data = image_4x4.flatten()\ninputs = image_data\nprint(inputs)\nparams = torch.rand(12, requires_grad=True, dtype=torch.float32)\ninputs = torch.tensor(image_data, dtype=torch.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:07.788684Z","iopub.execute_input":"2026-05-11T16:17:07.788928Z","iopub.status.idle":"2026-05-11T16:17:14.784536Z","shell.execute_reply.started":"2026-05-11T16:17:07.788890Z","shell.execute_reply":"2026-05-11T16:17:14.783553Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define QNode for the full QSAN circuit\n@qml.qnode(dev, interface=\"torch\", diff_method=\"adjoint\")\ndef qsan_circuit(inputs, params):\n    #Image encoding\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m, 2 * n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(2 * n * m, 3 * n * m), normalize=True)\n    qml.Barrier(wires=range(total_wires))\n    #Query\n    apply_u(n, m, input_wires,params[0:4])\n    qml.Barrier(wires=range(total_wires))\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='q')\n    qml.Barrier(wires=range(total_wires))\n    #Key\n    apply_u(n, m, input_wires, params[4:8])\n    qml.Barrier(wires=range(total_wires))\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='k')\n    qml.Barrier(wires=range(total_wires))\n    #QLS\n    qls(n, m, query_wires, key_wires, qls_wires)\n    qml.Barrier(wires=range(total_wires))\n    #Value\n    apply_u(n, m, input_wires, params[8:12])\n    qml.Barrier(wires=range(total_wires))\n    compression_operation(n, m, input_wires)\n    qml.Barrier(wires=range(total_wires))\n    #Qbits selection\n    slicing_operation(n, m, query_wires, qls_wires)\n    qml.Barrier(wires=range(total_wires))\n    return [\n        qml.expval(qml.PauliZ(input_wires[n * m - 2])),\n        qml.expval(qml.PauliZ(input_wires[n * m - 1])),\n        qml.expval(qml.PauliZ(query_wires[0])),\n    ]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:14.785669Z","iopub.execute_input":"2026-05-11T16:17:14.785987Z","iopub.status.idle":"2026-05-11T16:17:14.793712Z","shell.execute_reply.started":"2026-05-11T16:17:14.785961Z","shell.execute_reply":"2026-05-11T16:17:14.793060Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"qml.draw_mpl(qsan_circuit)(inputs, params)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:14.794528Z","iopub.execute_input":"2026-05-11T16:17:14.794749Z","iopub.status.idle":"2026-05-11T16:17:15.901136Z","shell.execute_reply.started":"2026-05-11T16:17:14.794735Z","shell.execute_reply":"2026-05-11T16:17:15.900432Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# QBSASM","metadata":{}},{"cell_type":"code","source":"# Visualization-only QNode for QBSASM extraction.\n# Same circuit up through QLS, but returns the density matrix of the QLS register\n# instead of expectation values. No differentiation needed.\n@qml.qnode(dev, interface=\"torch\", diff_method=None)\ndef qsan_circuit_qbsasm(inputs, params):\n    # Image encoding\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m, 2 * n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(2 * n * m, 3 * n * m), normalize=True)\n    \n    # Query\n    apply_u(n, m, input_wires, params[0:4])\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='q')\n    \n    # Key\n    apply_u(n, m, input_wires, params[4:8])\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='k')\n    \n    # QLS — stop here, this is where QBSASM lives\n    qls(n, m, query_wires, key_wires, qls_wires)\n    \n    # Return the density matrix of the QLS register\n    return qml.density_matrix(wires=qls_wires)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:15.903160Z","iopub.execute_input":"2026-05-11T16:17:15.903508Z","iopub.status.idle":"2026-05-11T16:17:15.909643Z","shell.execute_reply.started":"2026-05-11T16:17:15.903465Z","shell.execute_reply":"2026-05-11T16:17:15.908925Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Heatmap function of the QBSASM","metadata":{}},{"cell_type":"code","source":"def plot_qbsasm_heatmap(model, sample_image, label=\"qbsasm\"):\n    \"\"\"\n    Generate the QBSASM heatmap for a given image sample.\n    Replicates Fig. 13 from the QSAN paper.\n    \n    sample_image: a single 8x8 image tensor (shape [1, 8, 8] or [8, 8])\n    \"\"\"\n    model.eval()\n    with torch.no_grad():\n        # Reproduce the model's preprocessing pipeline\n        img = sample_image\n        if img.dim() == 2:\n            img = img.unsqueeze(0).unsqueeze(0)  # [1, 1, 8, 8]\n        elif img.dim() == 3:\n            img = img.unsqueeze(0)  # [1, 1, 8, 8]\n        \n        # Extract patches the same way the model does\n        patches = img.unfold(2, 4, 4).unfold(3, 4, 4)\n        patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()\n        rnn_input = patches.reshape(1, 4, 16)  # [batch=1, seq=4, dim=16]\n        \n        # Pass through GRU to get context-aware patches\n        context_aware_patches, _ = model.rnn_context_encoder(rnn_input)\n        all_patches = context_aware_patches.reshape(4, 16)  # 4 patches of 16 features\n        \n        # Run the QBSASM circuit on each patch and average the resulting density matrices\n        # (one image = 4 patches, each produces a density matrix)\n        density_matrices = []\n        for patch in all_patches:\n            if torch.isnan(patch).any():\n                patch = torch.nan_to_num(patch, nan=1e-3)\n            dm = qsan_circuit_qbsasm(patch, model.qsan_classifier.q_params)\n            density_matrices.append(dm)\n        \n        # Average across patches to get one QBSASM per image\n        avg_density = torch.stack(density_matrices).mean(dim=0)\n        \n        # Diagonal of density matrix = measurement probabilities = self-attention scores\n        probs = torch.diag(avg_density).real.cpu().numpy()\n        \n        # Reshape into a square matrix\n        side = int(np.sqrt(len(probs)))\n        qbsasm = probs.reshape(side, side)\n    \n    plt.figure(figsize=(7, 6))\n    n_bits = int(np.log2(side))\n    labels = [format(i, f'0{n_bits}b') for i in range(side)]\n    \n    sns.heatmap(\n        qbsasm,\n        cmap='viridis',\n        xticklabels=labels,\n        yticklabels=labels,\n        cbar_kws={'label': 'Self-attention score'},\n        square=True,\n    )\n    plt.xlabel(\"|Q⟩\")\n    plt.ylabel(\"|K⟩\")\n    plt.title(f\"QBSASM — {label}\")\n    plt.tight_layout()\n    plt.savefig(f'qbsasm_{label}.png', dpi=150, bbox_inches='tight')\n    plt.close()\n    print(f\"  Saved QBSASM heatmap to 'qbsasm_{label}.png'\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:15.910320Z","iopub.execute_input":"2026-05-11T16:17:15.910539Z","iopub.status.idle":"2026-05-11T16:17:15.928045Z","shell.execute_reply.started":"2026-05-11T16:17:15.910525Z","shell.execute_reply":"2026-05-11T16:17:15.927323Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Visualization-only QNode for QBSASM extraction.\n# Same circuit up through QLS, but returns the density matrix of the QLS register\n# instead of expectation values. No differentiation needed.\n@qml.qnode(dev, interface=\"torch\", diff_method=None)\ndef qsan_circuit_qbsasm(inputs, params):\n    # Image encoding\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(n * m, 2 * n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs, wires=range(2 * n * m, 3 * n * m), normalize=True)\n    \n    # Query\n    apply_u(n, m, input_wires, params[0:4])\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='q')\n    \n    # Key\n    apply_u(n, m, input_wires, params[4:8])\n    barbell_operation(n, m, input_wires, query_wires, key_wires, step='k')\n    \n    # QLS — stop here, this is where QBSASM lives\n    qls(n, m, query_wires, key_wires, qls_wires)\n    \n    # Return the density matrix of the QLS register\n    return qml.density_matrix(wires=qls_wires)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:15.928848Z","iopub.execute_input":"2026-05-11T16:17:15.929128Z","iopub.status.idle":"2026-05-11T16:17:15.949486Z","shell.execute_reply.started":"2026-05-11T16:17:15.929102Z","shell.execute_reply":"2026-05-11T16:17:15.948642Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Results","metadata":{}},{"cell_type":"code","source":"# Simulate and get results\nqsan_circuit(inputs, params)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:15.950339Z","iopub.execute_input":"2026-05-11T16:17:15.950743Z","iopub.status.idle":"2026-05-11T16:17:16.234464Z","shell.execute_reply.started":"2026-05-11T16:17:15.950718Z","shell.execute_reply":"2026-05-11T16:17:16.233856Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Dataset Preparation","metadata":{}},{"cell_type":"markdown","source":"We are trying to recreate the exact experiment done by the QSAN paper using the same training parameters and dataset","metadata":{}},{"cell_type":"code","source":"import pennylane as qml\nfrom pennylane import numpy as np\nfrom pennylane.optimize import NesterovMomentumOptimizer\nimport numpy as np\nimport struct\nfrom array import array\nfrom os.path import join\nimport torch\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport torchvision # Import torchvision\nimport torchvision.transforms as transforms # Import transforms\nfrom torch.utils.data import TensorDataset, DataLoader, Subset\nimport os\nfrom scipy.ndimage import zoom\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:16.235076Z","iopub.execute_input":"2026-05-11T16:17:16.235267Z","iopub.status.idle":"2026-05-11T16:17:20.380368Z","shell.execute_reply.started":"2026-05-11T16:17:16.235252Z","shell.execute_reply":"2026-05-11T16:17:20.379473Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Loading the dataset","metadata":{}},{"cell_type":"code","source":"\n# # --- 1. AUTOMATED DATA LOADING AND PREPROCESSING ---\n# print(\"--- 1. Loading and Processing MNIST Data (Automated) ---\")\n\n# # --- Hyperparameters for Data ---\n# N_TRAIN_PER_CLASS = 100\n# N_TEST_PER_CLASS = 50\n# BATCH_SIZE = 15\n\n# # Use torchvision to download and load the dataset\n# # This handles all the file path and loading logic automatically.\n# full_train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)\n# full_test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True)\n\n# # --- Filter for digits 0 and 1 ---\n# def filter_and_select_subset(dataset, n_samples_per_class):\n#     # Get indices for classes 0 and 1\n#     idx_0 = (dataset.targets == 0).nonzero(as_tuple=True)[0]\n#     idx_1 = (dataset.targets == 1).nonzero(as_tuple=True)[0]\n\n#     # Select the required number of samples\n#     subset_idx_0 = idx_0[:n_samples_per_class]\n#     subset_idx_1 = idx_1[:n_samples_per_class]\n\n#     # Combine the indices\n#     final_indices = torch.cat([subset_idx_0, subset_idx_1])\n    \n#     # Get the data and labels for the final subset\n#     images = dataset.data[final_indices].float() / 255.0 # Normalize to [0, 1]\n#     labels = dataset.targets[final_indices].float().reshape(-1, 1)\n    \n#     return images, labels\n\n# # Create the small training and test sets\n# x_train, y_train = filter_and_select_subset(full_train_dataset, N_TRAIN_PER_CLASS)\n# x_test, y_test = filter_and_select_subset(full_test_dataset, N_TEST_PER_CLASS)\n\n# # --- Resize, Flatten, and Create DataLoaders ---\n# def process_data_for_qsan(images, labels, batch_size, shuffle=False):\n\n\n#     print(f\"Original images shape: {images.shape}\")\n\n#     resize_transform = transforms.Resize((8, 8), antialias=True)\n    \n#     # 2. Add a channel dimension. The input is [N, H, W], transforms need [N, C, H, W]\n#     images_with_channel = images.unsqueeze(1) # Shape becomes [N, 1, 28, 28]\n\n#     # 3. Apply the transformation\n#     resized_images = resize_transform(images_with_channel) # Shape becomes [N, 1, 8, 8]\n#     print(f\"Resized images shape: {resized_images.shape}\")\n    \n#     # Convert back to PyTorch Tensors\n#     X_tensor = torch.tensor(resized_images, dtype=torch.float32)\n#     print(X_tensor.shape)\n#     Y_tensor = labels \n    \n#     # Create the final TensorDataset and DataLoader\n#     dataset = TensorDataset(X_tensor, Y_tensor)\n#     loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)\n    \n#     return loader, X_tensor, Y_tensor\n\n# train_loader, _, _ = process_data_for_qsan(x_train, y_train, BATCH_SIZE, shuffle=True)\n# test_loader, X_test, Y_test = process_data_for_qsan(x_test, y_test, BATCH_SIZE, shuffle=False)\n\n# print(f\"Data preparation complete. Training batches: {len(train_loader)}, Test samples: {len(X_test)}\")\n\n\n# # (This section is unchanged and remains exactly as in the previous response)\n# print(\"\\n--- 2. Defining Full QSAN Quantum Circuit ---\")\n\n# N_QUBITS = 4\n# n, m = 2, 2\n# qls_reg_size = m * n * (n + 1) // 2\n# TOTAL_QUBITS = (N_QUBITS * 3) + qls_reg_size\n# PARAMS_PER_U = N_QUBITS\n# TOTAL_PARAMS = PARAMS_PER_U * 3\n\n# input_wires, query_wires, key_wires, qls_wires = (\n#     list(range(N_QUBITS)),\n#     list(range(N_QUBITS, 2 * N_QUBITS)),\n#     list(range(2 * N_QUBITS, 3 * N_QUBITS)),\n#     list(range(3 * N_QUBITS, TOTAL_QUBITS)),\n# )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:20.381313Z","iopub.execute_input":"2026-05-11T16:17:20.382290Z","iopub.status.idle":"2026-05-11T16:17:20.386847Z","shell.execute_reply.started":"2026-05-11T16:17:20.382259Z","shell.execute_reply":"2026-05-11T16:17:20.386052Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Reviewer's request","metadata":{}},{"cell_type":"code","source":"# --- 1. AUTOMATED DATA LOADING AND PREPROCESSING ---\nprint(\"--- 1. Loading and Processing MNIST Data (Automated) ---\")\n\n# --- Hyperparameters for Data ---\nDIGIT_PAIRS = [(0, 1), (3, 5), (4, 9), (7, 9)]  # Multiple pairs to test robustness\nN_TRAIN_PER_CLASS = 250  # Use much larger subsets (or set to None for all available)\nN_TEST_PER_CLASS = 100\nBATCH_SIZE = 15\nN_SEEDS = 5  # For reporting mean ± std\n\n# Use torchvision to download and load the dataset\nfull_train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True)\nfull_test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True)\n\n# --- Filter for any two digit classes ---\ndef filter_and_select_subset(dataset, class_a, class_b, n_samples_per_class=None, seed=42):\n    \"\"\"\n    Extracts samples for a binary classification task between class_a and class_b.\n    Labels are mapped to 0 (class_a) and 1 (class_b).\n    Samples are shuffled with a fixed seed for reproducibility.\n    \"\"\"\n    idx_a = (dataset.targets == class_a).nonzero(as_tuple=True)[0]\n    idx_b = (dataset.targets == class_b).nonzero(as_tuple=True)[0]\n\n    # Optionally limit the number of samples per class\n    if n_samples_per_class is not None:\n        # Shuffle before selecting to avoid ordering bias\n        gen = torch.Generator().manual_seed(seed)\n        idx_a = idx_a[torch.randperm(len(idx_a), generator=gen)[:n_samples_per_class]]\n        idx_b = idx_b[torch.randperm(len(idx_b), generator=gen)[:n_samples_per_class]]\n\n    final_indices = torch.cat([idx_a, idx_b])\n\n    # Shuffle combined indices so batches are mixed\n    gen = torch.Generator().manual_seed(seed)\n    final_indices = final_indices[torch.randperm(len(final_indices), generator=gen)]\n\n    images = dataset.data[final_indices].float() / 255.0  # Normalize to [0, 1]\n    # Binary labels: class_a -> 0, class_b -> 1\n    labels = (dataset.targets[final_indices] == class_b).float().reshape(-1, 1)\n\n    return images, labels\n\n# --- Resize, Flatten, and Create DataLoaders ---\ndef process_data_for_qsan(images, labels, batch_size, shuffle=False):\n    resize_transform = transforms.Resize((8, 8), antialias=True)\n\n    images_with_channel = images.unsqueeze(1)  # [N, 1, 28, 28]\n    resized_images = resize_transform(images_with_channel)  # [N, 1, 8, 8]\n\n    X_tensor = resized_images.clone().detach().float()\n    Y_tensor = labels\n\n    dataset = TensorDataset(X_tensor, Y_tensor)\n    loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)\n\n    return loader, X_tensor, Y_tensor\n\n# --- Prepare data for all digit pairs ---\n# This dict will hold data loaders for each pair, to be iterated over during training\ndigit_pair_data = {}\nfor class_a, class_b in DIGIT_PAIRS:\n    print(f\"\\nPreparing data for digit pair ({class_a} vs {class_b})...\")\n\n    x_train, y_train = filter_and_select_subset(\n        full_train_dataset, class_a, class_b, N_TRAIN_PER_CLASS, seed=42\n    )\n    x_test, y_test = filter_and_select_subset(\n        full_test_dataset, class_a, class_b, N_TEST_PER_CLASS, seed=42\n    )\n\n    train_loader, _, _ = process_data_for_qsan(x_train, y_train, BATCH_SIZE, shuffle=True)\n    test_loader, X_test, Y_test = process_data_for_qsan(x_test, y_test, BATCH_SIZE, shuffle=False)\n\n    digit_pair_data[(class_a, class_b)] = {\n        'train_loader': train_loader,\n        'test_loader': test_loader,\n        'X_test': X_test,\n        'Y_test': Y_test,\n    }\n    print(f\"  Train: {len(x_train)} samples, Test: {len(x_test)} samples\")\n\nprint(f\"\\nData preparation complete for {len(DIGIT_PAIRS)} digit pairs.\")\n\n\n# (Section 2 is unchanged)\nprint(\"\\n--- 2. Defining Full QSAN Quantum Circuit ---\")\nN_QUBITS = 4\nn, m = 2, 2\nqls_reg_size = m * n * (n + 1) // 2\nTOTAL_QUBITS = (N_QUBITS * 3) + qls_reg_size\nPARAMS_PER_U = N_QUBITS\nTOTAL_PARAMS = PARAMS_PER_U * 3\ninput_wires, query_wires, key_wires, qls_wires = (\n    list(range(N_QUBITS)),\n    list(range(N_QUBITS, 2 * N_QUBITS)),\n    list(range(2 * N_QUBITS, 3 * N_QUBITS)),\n    list(range(3 * N_QUBITS, TOTAL_QUBITS)),\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:20.388139Z","iopub.execute_input":"2026-05-11T16:17:20.388384Z","iopub.status.idle":"2026-05-11T16:17:21.734859Z","shell.execute_reply.started":"2026-05-11T16:17:20.388363Z","shell.execute_reply":"2026-05-11T16:17:21.733948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class QSAN_Classifier(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.q_params = nn.Parameter(torch.randn(TOTAL_PARAMS) * 0.1)\n        self.classical_layer = nn.Linear(3, 1) \n    def forward(self, x_batch):\n        batch_results = []\n        for x in x_batch:\n            q_output_list = qsan_circuit(inputs=x, params=self.q_params)\n            # Each output is now a scalar tensor, not a 2-element vector\n            q_output_tensor = torch.stack(q_output_list)\n            classical_output = self.classical_layer(q_output_tensor.float())\n            batch_results.append(classical_output)\n        return torch.stack(batch_results).view(-1, 1)\n\nclass HybridRNN_Q_Model(nn.Module):\n\n    def __init__(self, patch_size=4, patch_features=16):\n        super().__init__()\n        self.patch_size = patch_size\n        patch_dim = patch_size * patch_size # 16\n\n        self.rnn_context_encoder = nn.GRU(\n            input_size=patch_dim,\n            hidden_size=patch_features, # Outputting 16 features per patch\n            num_layers=1,\n            batch_first=True\n        )\n        self.qsan_classifier = QSAN_Classifier()\n\n    def forward(self, image_batch):\n        \n        if image_batch.dim() == 3:\n            image_batch = image_batch.unsqueeze(1) \n            \n        patches = image_batch.unfold(2, self.patch_size, self.patch_size).unfold(3, self.patch_size, self.patch_size)\n        patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()\n        \n        batch_size = image_batch.size(0)\n        sequence_len = 4 \n        patch_dim = self.patch_size * self.patch_size\n        \n        rnn_input = patches.reshape(batch_size, sequence_len, patch_dim)\n\n        context_aware_patches, _ = self.rnn_context_encoder(rnn_input)\n\n        all_patches_batch = context_aware_patches.reshape(batch_size * sequence_len, patch_dim)\n        if torch.isnan(all_patches_batch).any():\n            nan_count = torch.isnan(all_patches_batch).sum().item()\n            print(f\"WARNING: {nan_count} NaN values in patch embeddings — replacing with small noise\")\n            noise = torch.randn_like(all_patches_batch) * 1e-3\n            all_patches_batch = torch.where(torch.isnan(all_patches_batch), noise, all_patches_batch)\n        patch_predictions = self.qsan_classifier(all_patches_batch)\n\n\n        predictions_by_image = patch_predictions.reshape(batch_size, sequence_len, 1)\n\n        final_image_prediction = predictions_by_image.mean(dim=1)\n        \n        return final_image_prediction","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:21.735768Z","iopub.execute_input":"2026-05-11T16:17:21.736068Z","iopub.status.idle":"2026-05-11T16:17:21.745598Z","shell.execute_reply.started":"2026-05-11T16:17:21.736043Z","shell.execute_reply":"2026-05-11T16:17:21.744959Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Reviewer's edit","metadata":{}},{"cell_type":"code","source":"from tqdm import tqdm\nfrom sklearn.metrics import (\n    confusion_matrix, roc_auc_score, f1_score,\n    accuracy_score, precision_score, recall_score\n)\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n\ndef compute_metrics(y_true, y_pred_probs, threshold=0.5):\n    \"\"\"\n    Computes all clinically relevant metrics from true labels and predicted probabilities.\n    Returns a dict with accuracy, AUROC, F1, sensitivity (recall), specificity, and precision.\n    \"\"\"\n    y_true_np = y_true.cpu().numpy().flatten()\n    y_probs_np = y_pred_probs.cpu().detach().numpy().flatten()\n    y_pred_np = (y_probs_np >= threshold).astype(int)\n\n    tn, fp, fn, tp = confusion_matrix(y_true_np, y_pred_np, labels=[0, 1]).ravel()\n\n    metrics = {\n        'accuracy': accuracy_score(y_true_np, y_pred_np),\n        'auroc': roc_auc_score(y_true_np, y_probs_np),\n        'f1': f1_score(y_true_np, y_pred_np, zero_division=0),\n        'sensitivity': recall_score(y_true_np, y_pred_np, zero_division=0),  # TPR\n        'specificity': tn / (tn + fp) if (tn + fp) > 0 else 0.0,\n        'precision': precision_score(y_true_np, y_pred_np, zero_division=0),\n    }\n    return metrics, confusion_matrix(y_true_np, y_pred_np, labels=[0, 1])\n\n\ndef plot_confusion_matrix(y_true, y_pred_probs, label, save_prefix=\"cm\", threshold=0.5):\n    \"\"\"\n    Calculates and plots a confusion matrix.\n    label: a descriptive string, e.g. \"pair_3v5_seed_2\"\n    \"\"\"\n    y_true_np = y_true.cpu().numpy().flatten()\n    y_pred_np = (y_pred_probs.cpu().detach().numpy().flatten() >= threshold).astype(int)\n\n    cm = confusion_matrix(y_true_np, y_pred_np, labels=[0, 1])\n\n    plt.figure(figsize=(6, 5))\n    sns.heatmap(\n        cm,\n        annot=True,\n        fmt='d',\n        cmap='Blues',\n        xticklabels=['Predicted Negative', 'Predicted Positive'],\n        yticklabels=['Actual Negative', 'Actual Positive']\n    )\n    plt.xlabel(\"Predicted Label\")\n    plt.ylabel(\"True Label\")\n    plt.title(f\"Confusion Matrix — {label}\")\n\n    filename = f'{save_prefix}_{label}.png'\n    plt.savefig(filename, dpi=150, bbox_inches='tight')\n    print(f\"Saved confusion matrix plot to '{filename}'\")\n    plt.close()\n\n\ndef print_results_table(all_results):\n    \"\"\"\n    Prints a summary table: mean ± std for each metric across seeds, per digit pair.\n    all_results: dict of {(class_a, class_b): [list of metric dicts, one per seed]}\n    \"\"\"\n    metric_names = ['accuracy', 'auroc', 'f1', 'sensitivity', 'specificity']\n\n    print(\"\\n\" + \"=\" * 85)\n    print(f\"{'Digit Pair':<14}\", end=\"\")\n    for m in metric_names:\n        print(f\"  {m:>18}\", end=\"\")\n    print()\n    print(\"-\" * 85)\n\n    for pair, seed_metrics in all_results.items():\n        pair_label = f\"{pair[0]} vs {pair[1]}\"\n        print(f\"{pair_label:<14}\", end=\"\")\n        for m in metric_names:\n            values = [s[m] for s in seed_metrics]\n            mean = np.mean(values)\n            std = np.std(values)\n            print(f\"  {mean:.4f} ± {std:.4f}\", end=\"\")\n        print()\n\n    print(\"=\" * 85)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:21.746426Z","iopub.execute_input":"2026-05-11T16:17:21.746726Z","iopub.status.idle":"2026-05-11T16:17:22.718038Z","shell.execute_reply.started":"2026-05-11T16:17:21.746709Z","shell.execute_reply":"2026-05-11T16:17:22.717412Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Reviewer's edit","metadata":{}},{"cell_type":"code","source":"import copy\n\nNUM_EPOCHS = 3                # reduced from 5\nBATCH_SIZE = 15\nN_SEEDS = 5\nEARLY_STOP_PATIENCE = 1       # stop if no improvement for 1 epoch (was 2)\nNAN_PATIENCE = 1              # stop on FIRST NaN batch (was 3)\n\nall_results = {}\n\nfor class_a, class_b in DIGIT_PAIRS:\n    pair_label = f\"{class_a}v{class_b}\"\n    print(f\"\\n{'='*70}\")\n    print(f\"  DIGIT PAIR: {class_a} vs {class_b}\")\n    print(f\"{'='*70}\")\n\n    data = digit_pair_data[(class_a, class_b)]\n    train_loader = data['train_loader']\n    test_loader = data['test_loader']\n    X_test = data['X_test']\n    Y_test = data['Y_test']\n\n    seed_metrics = []\n    seed_confusion_matrices = []\n\n    for seed in range(N_SEEDS):\n        print(f\"\\n--- Pair {pair_label} | Seed {seed+1}/{N_SEEDS} ---\")\n\n        torch.manual_seed(seed)\n        np.random.seed(seed)\n\n        model = HybridRNN_Q_Model()\n        optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n        # CHANGED: ReduceLROnPlateau responds to actual loss, not fixed schedule\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, mode='min', factor=0.3, patience=1\n        )\n        criterion = nn.BCEWithLogitsLoss()\n\n        train_loss_history = []\n        train_acc_history = []\n        test_loss_history = []\n        test_acc_history = []\n        best_test_accuracy = 0.0\n        best_model_state = None\n        best_epoch = 0\n        epochs_without_improvement = 0\n        early_stopped = False\n        nan_triggered = False\n\n        try:\n            for epoch in range(NUM_EPOCHS):\n                if early_stopped or nan_triggered:\n                    break\n\n                model.train()\n                epoch_train_loss = 0.0\n                epoch_train_correct = 0\n                total_train_samples = 0\n\n                pbar = tqdm(train_loader,\n                            desc=f\"[{pair_label} s{seed+1}] Epoch {epoch+1}/{NUM_EPOCHS}\",\n                            unit=\"batch\")\n                for features, labels in pbar:\n                    optimizer.zero_grad()\n                    outputs = model(features)\n                    labels = labels.float().view_as(outputs)\n                    loss = criterion(outputs, labels)\n\n                    # IMMEDIATE NaN detection — stop training this seed right now\n                    if torch.isnan(loss).any() or torch.isinf(loss).any():\n                        print(f\"\\n  NaN/Inf detected — stopping seed immediately.\")\n                        if best_model_state is not None:\n                            print(f\"  Restoring best model from epoch {best_epoch} (acc={best_test_accuracy:.4f})\")\n                            model.load_state_dict(best_model_state)\n                        else:\n                            print(f\"  No prior best model available for this seed.\")\n                        nan_triggered = True\n                        break\n\n                    loss.backward()\n                    # Tighter clipping\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)\n                    optimizer.step()\n\n                    epoch_train_loss += loss.item() * features.size(0)\n                    preds = torch.round(torch.sigmoid(outputs))\n                    epoch_train_correct += (preds == labels).sum().item()\n                    total_train_samples += features.size(0)\n                    current_acc = epoch_train_correct / total_train_samples\n                    pbar.set_postfix({'loss': f\"{loss.item():.4f}\", 'acc': f\"{current_acc:.2f}\"})\n\n                if nan_triggered or total_train_samples == 0:\n                    break\n\n                avg_train_loss = epoch_train_loss / total_train_samples\n                avg_train_acc = epoch_train_correct / total_train_samples\n\n                # Evaluate\n                model.eval()\n                with torch.no_grad():\n                    test_outputs = model(X_test)\n                    test_labels = Y_test.float().view_as(test_outputs)\n                    test_loss = criterion(test_outputs, test_labels)\n                    test_preds = torch.round(torch.sigmoid(test_outputs))\n                    test_correct = (test_preds == test_labels).sum().item()\n                    test_accuracy = test_correct / len(test_labels)\n\n                train_loss_history.append(avg_train_loss)\n                train_acc_history.append(avg_train_acc)\n                test_loss_history.append(test_loss.item())\n                test_acc_history.append(test_accuracy)\n\n                print(f\"  Epoch {epoch+1} | Train Loss: {avg_train_loss:.4f} | Train Acc: {avg_train_acc:.2f} | \"\n                      f\"Test Loss: {test_loss.item():.4f} | Test Acc: {test_accuracy:.2f}\")\n\n                # Track best model\n                if test_accuracy > best_test_accuracy and not np.isnan(test_loss.item()):\n                    best_test_accuracy = test_accuracy\n                    best_model_state = copy.deepcopy(model.state_dict())\n                    best_epoch = epoch + 1\n                    epochs_without_improvement = 0\n                    print(f\"  ✓ New best at epoch {best_epoch}: {best_test_accuracy:.4f}\")\n                else:\n                    epochs_without_improvement += 1\n\n                # Early stop on plateau — stop when no further gain is expected\n                if epochs_without_improvement >= EARLY_STOP_PATIENCE:\n                    print(f\"  Plateau detected ({EARLY_STOP_PATIENCE} epoch(s) without improvement) — stopping.\")\n                    early_stopped = True\n\n                # Adaptive LR based on training loss\n                scheduler.step(avg_train_loss)\n\n        except Exception as e:\n            print(f\"\\n--- ERROR on pair {pair_label}, seed {seed+1}: {e} ---\")\n            if best_model_state is None:\n                print(f\"  No best model available — skipping this seed.\")\n                continue\n\n        # --- Final evaluation on best model ---\n        if best_model_state is None:\n            print(f\"  WARNING: No valid model trained for seed {seed+1}, skipping metrics.\")\n            continue\n\n        model.load_state_dict(best_model_state)\n        model.eval()\n        with torch.no_grad():\n            final_outputs = model(X_test)\n            final_probs = torch.sigmoid(final_outputs)\n\n        if torch.isnan(final_probs).any():\n            print(f\"  WARNING: NaN in final probabilities for seed {seed+1}, skipping.\")\n            continue\n\n        metrics, cm = compute_metrics(Y_test, final_probs)\n        seed_metrics.append(metrics)\n        seed_confusion_matrices.append(cm)\n\n        print(f\"  Seed {seed+1} best (epoch {best_epoch}) — \"\n              f\"Acc: {metrics['accuracy']:.4f} | AUROC: {metrics['auroc']:.4f} | \"\n              f\"F1: {metrics['f1']:.4f} | Sens: {metrics['sensitivity']:.4f} | \"\n              f\"Spec: {metrics['specificity']:.4f}\")\n\n        plot_confusion_matrix(Y_test, final_probs, label=f\"pair_{pair_label}_seed_{seed+1}\")\n        torch.save(best_model_state, f\"best_model_{pair_label}_seed{seed+1}.pth\")\n        # --- QBSASM heatmaps (only on the last successful seed per pair) ---\n        if seed == N_SEEDS - 1:\n            try:\n                sample_a_idx = (Y_test.flatten() == 0).nonzero(as_tuple=True)[0][0].item()\n                sample_b_idx = (Y_test.flatten() == 1).nonzero(as_tuple=True)[0][0].item()\n                \n                plot_qbsasm_heatmap(model, X_test[sample_a_idx], \n                                    label=f\"{pair_label}_class{class_a}\")\n                plot_qbsasm_heatmap(model, X_test[sample_b_idx], \n                                    label=f\"{pair_label}_class{class_b}\")\n            except Exception as e:\n                print(f\"  QBSASM plotting failed: {e}\")\n    # --- Aggregate per pair ---\n    if seed_confusion_matrices:\n        all_results[(class_a, class_b)] = seed_metrics\n        agg_cm = np.sum(seed_confusion_matrices, axis=0)\n        plt.figure(figsize=(6, 5))\n        sns.heatmap(agg_cm, annot=True, fmt='d', cmap='Blues',\n                    xticklabels=[f'Pred {class_a}', f'Pred {class_b}'],\n                    yticklabels=[f'Actual {class_a}', f'Actual {class_b}'])\n        plt.xlabel(\"Predicted Label\"); plt.ylabel(\"True Label\")\n        plt.title(f\"Aggregate CM — {class_a} vs {class_b} ({len(seed_confusion_matrices)} seeds)\")\n        plt.savefig(f'cm_aggregate_{pair_label}.png', dpi=150, bbox_inches='tight')\n        plt.close()\n        print(f\"  Aggregate CM saved with {len(seed_confusion_matrices)} valid seeds.\")\n    else:\n        print(f\"  No valid seeds for pair {pair_label}.\")\n\nprint_results_table(all_results)\nprint(\"\\n--- All Training Complete ---\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T16:17:22.718737Z","iopub.execute_input":"2026-05-11T16:17:22.719176Z"}},"outputs":[],"execution_count":null}]}