{"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":"none","dataSources":[{"sourceType":"competition","sourceId":3004,"databundleVersionId":861823},{"sourceType":"datasetVersion","sourceId":12601980,"datasetId":7959846,"databundleVersionId":13199131},{"sourceType":"kernelVersion","sourceId":320601042}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"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 uninstall -y jax jaxlib jax-cuda12-pjrt jax-cuda12-plugin\n\n!pip install --quiet \\\n    \"numpy==1.26.4\" \\\n    \"scipy==1.13.1\" \\\n    \"autoray==0.6.12\" \\\n    \"qiskit==1.4.2\" \\\n    \"qiskit-aer==0.15.1\" \\\n    \"qiskit-ibm-runtime==0.30.0\" \\\n    \"pennylane==0.39.0\" \\\n    \"pennylane-lightning==0.39.0\" \\\n    pylatexenc\n\n!pip install --quiet --no-deps \"numpy==1.26.4\"\n\nprint(\"Done — NOW RESTART THE KERNEL before running anything else.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:01:54.418052Z","iopub.execute_input":"2026-05-25T12:01:54.421505Z","iopub.status.idle":"2026-05-25T12:02:03.370041Z","shell.execute_reply.started":"2026-05-25T12:01:54.421468Z","shell.execute_reply":"2026-05-25T12:02:03.368504Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy, scipy, qiskit, qiskit_aer, pennylane\nprint(\"numpy     \", numpy.__version__)\nprint(\"scipy     \", scipy.__version__)\nprint(\"qiskit    \", qiskit.__version__)\nprint(\"qiskit_aer\", qiskit_aer.__version__)\nprint(\"pennylane \", pennylane.__version__)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:02:03.372881Z","iopub.execute_input":"2026-05-25T12:02:03.373238Z","iopub.status.idle":"2026-05-25T12:02:03.380695Z","shell.execute_reply.started":"2026-05-25T12:02:03.37321Z","shell.execute_reply":"2026-05-25T12:02:03.379772Z"}},"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-25T12:02:03.381774Z","iopub.execute_input":"2026-05-25T12:02:03.382154Z","iopub.status.idle":"2026-05-25T12:02:03.400348Z","shell.execute_reply.started":"2026-05-25T12:02:03.382126Z","shell.execute_reply":"2026-05-25T12:02:03.39929Z"}},"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-25T12:02:03.401413Z","iopub.execute_input":"2026-05-25T12:02:03.401665Z","iopub.status.idle":"2026-05-25T12:02:03.421111Z","shell.execute_reply.started":"2026-05-25T12:02:03.401647Z","shell.execute_reply":"2026-05-25T12:02:03.420283Z"}},"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-25T12:02:03.424412Z","iopub.execute_input":"2026-05-25T12:02:03.42472Z","iopub.status.idle":"2026-05-25T12:02:03.443024Z","shell.execute_reply.started":"2026-05-25T12:02:03.4247Z","shell.execute_reply":"2026-05-25T12:02:03.442072Z"}},"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-25T12:02:03.444036Z","iopub.execute_input":"2026-05-25T12:02:03.444689Z","iopub.status.idle":"2026-05-25T12:02:03.459714Z","shell.execute_reply.started":"2026-05-25T12:02:03.44466Z","shell.execute_reply":"2026-05-25T12:02:03.458663Z"}},"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-25T12:02:03.460829Z","iopub.execute_input":"2026-05-25T12:02:03.461132Z","iopub.status.idle":"2026-05-25T12:02:03.475575Z","shell.execute_reply.started":"2026-05-25T12:02:03.461107Z","shell.execute_reply":"2026-05-25T12:02:03.47467Z"}},"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-25T12:02:03.476486Z","iopub.execute_input":"2026-05-25T12:02:03.476741Z","iopub.status.idle":"2026-05-25T12:02:03.499381Z","shell.execute_reply.started":"2026-05-25T12:02:03.476721Z","shell.execute_reply":"2026-05-25T12:02:03.498365Z"}},"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-25T12:02:03.500341Z","iopub.execute_input":"2026-05-25T12:02:03.500731Z","iopub.status.idle":"2026-05-25T12:02:03.515759Z","shell.execute_reply.started":"2026-05-25T12:02:03.500703Z","shell.execute_reply":"2026-05-25T12:02:03.514815Z"}},"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-25T12:02:03.517031Z","iopub.execute_input":"2026-05-25T12:02:03.51736Z","iopub.status.idle":"2026-05-25T12:02:03.535243Z","shell.execute_reply.started":"2026-05-25T12:02:03.517339Z","shell.execute_reply":"2026-05-25T12:02:03.534309Z"}},"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-25T12:02:03.535987Z","iopub.execute_input":"2026-05-25T12:02:03.536277Z","iopub.status.idle":"2026-05-25T12:02:03.553106Z","shell.execute_reply.started":"2026-05-25T12:02:03.536256Z","shell.execute_reply":"2026-05-25T12:02:03.552064Z"}},"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-25T12:02:03.554128Z","iopub.execute_input":"2026-05-25T12:02:03.554459Z","iopub.status.idle":"2026-05-25T12:02:03.573402Z","shell.execute_reply.started":"2026-05-25T12:02:03.554437Z","shell.execute_reply":"2026-05-25T12:02:03.572516Z"}},"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\n    \n    # Image encoding — three DIFFERENT views to give Q, K, V genuine asymmetry\n    inputs_q = inputs\n    inputs_k = torch.roll(inputs, shifts=1, dims=-1)\n    sign_pattern = torch.tensor([1.0, -1.0] * (inputs.shape[-1] // 2),\n                                dtype=inputs.dtype, device=inputs.device)\n    inputs_v = torch.roll(inputs * sign_pattern, shifts=2, dims=-1)\n\n    qml.AmplitudeEmbedding(features=inputs_q, wires=range(n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs_k, wires=range(n * m, 2 * n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs_v, 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        qml.expval(qml.PauliZ(query_wires[2])),\n    ]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:59:52.432949Z","iopub.execute_input":"2026-05-25T12:59:52.433298Z","iopub.status.idle":"2026-05-25T12:59:52.44513Z","shell.execute_reply.started":"2026-05-25T12:59:52.433277Z","shell.execute_reply":"2026-05-25T12:59:52.444042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"qml.draw_mpl(qsan_circuit)(inputs, params)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:59:53.933932Z","iopub.execute_input":"2026-05-25T12:59:53.934635Z","iopub.status.idle":"2026-05-25T12:59:55.025977Z","shell.execute_reply.started":"2026-05-25T12:59:53.934609Z","shell.execute_reply":"2026-05-25T12:59:55.024886Z"}},"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    # Image encoding — three DIFFERENT views to give Q, K, V genuine asymmetry\n    inputs_q = inputs\n    inputs_k = torch.roll(inputs, shifts=1, dims=-1)\n    sign_pattern = torch.tensor([1.0, -1.0] * (inputs.shape[-1] // 2),\n                                dtype=inputs.dtype, device=inputs.device)\n    inputs_v = torch.roll(inputs * sign_pattern, shifts=2, dims=-1)\n\n    qml.AmplitudeEmbedding(features=inputs_q, wires=range(n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs_k, wires=range(n * m, 2 * n * m), normalize=True)\n    qml.AmplitudeEmbedding(features=inputs_v, wires=range(2 * n * m, 3 * n * m), normalize=True)\n    qml.Barrier(wires=range(total_wires))\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-25T13:01:04.412834Z","iopub.execute_input":"2026-05-25T13:01:04.413458Z","iopub.status.idle":"2026-05-25T13:01:04.422454Z","shell.execute_reply.started":"2026-05-25T13:01:04.413433Z","shell.execute_reply":"2026-05-25T13:01:04.421495Z"}},"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-25T12:02:04.750567Z","iopub.execute_input":"2026-05-25T12:02:04.751584Z","iopub.status.idle":"2026-05-25T12:02:04.781879Z","shell.execute_reply.started":"2026-05-25T12:02:04.75156Z","shell.execute_reply":"2026-05-25T12:02:04.781096Z"}},"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-25T12:02:04.782791Z","iopub.execute_input":"2026-05-25T12:02:04.783098Z","iopub.status.idle":"2026-05-25T12:02:04.806928Z","shell.execute_reply.started":"2026-05-25T12:02:04.783053Z","shell.execute_reply":"2026-05-25T12:02:04.805952Z"}},"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-25T12:02:04.807957Z","iopub.execute_input":"2026-05-25T12:02:04.808261Z","iopub.status.idle":"2026-05-25T12:02:04.972819Z","shell.execute_reply.started":"2026-05-25T12:02:04.808241Z","shell.execute_reply":"2026-05-25T12:02:04.972141Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Dataset Preparation","metadata":{}},{"cell_type":"code","source":"import os\nimport copy\nimport json\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms as transforms\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport pennylane as qml\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom sklearn.metrics import (\n    confusion_matrix, roc_auc_score, f1_score,\n    accuracy_score, precision_score, recall_score\n)\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:02:04.973678Z","iopub.execute_input":"2026-05-25T12:02:04.973907Z","iopub.status.idle":"2026-05-25T12:02:04.980406Z","shell.execute_reply.started":"2026-05-25T12:02:04.973889Z","shell.execute_reply":"2026-05-25T12:02:04.97938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. CONFIGURATION ---\nprint(\"--- 1. Configuring ---\")\n\n# Paths\nBASE_DATA_DIR = '/kaggle/input/datasets/bissenabbassi/224-tumors/processed_data'\nSPLITS_DIR = '/kaggle/input/notebooks/bissenabbassi/train-dataset-lidc-idri'\n\n# Hyperparameters\nIMG_SIZE = 8\nBATCH_SIZE = 16\nN_SEEDS = 5\nNUM_EPOCHS = 3\nEARLY_STOP_PATIENCE = 1\nRANDOM_SEED = 42\n\ntorch.manual_seed(RANDOM_SEED)\nnp.random.seed(RANDOM_SEED)\n\n# --- 2. LOAD PRE-SPLIT DATA ---\nprint(\"\\n--- 2. Loading pre-split data ---\")\ntrain_df = pd.read_csv(os.path.join(SPLITS_DIR, 'train_slices.csv'))\ntest_df = pd.read_csv(os.path.join(SPLITS_DIR, 'test_slices.csv'))\n\n# Verify leakage-free\nassert len(set(train_df['Patient_ID']) & set(test_df['Patient_ID'])) == 0, \"PATIENT LEAKAGE!\"\nprint(f\"  ✓ No patient overlap (verified)\")\nprint(f\"  Train: {len(train_df)} slices ({dict(train_df['Cancer_Label'].value_counts())})\")\nprint(f\"  Test:  {len(test_df)} slices ({dict(test_df['Cancer_Label'].value_counts())})\")\nprint(f\"  Train patients: {train_df['Patient_ID'].nunique()} | Test patients: {test_df['Patient_ID'].nunique()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:02:04.981487Z","iopub.execute_input":"2026-05-25T12:02:04.981883Z","iopub.status.idle":"2026-05-25T12:02:05.037115Z","shell.execute_reply.started":"2026-05-25T12:02:04.981854Z","shell.execute_reply":"2026-05-25T12:02:05.036185Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pytorch Datasets and Loaders","metadata":{}},{"cell_type":"code","source":"# --- 3. DATASET CLASS AND TRANSFORMS ---\nprint(\"\\n--- 3. Creating PyTorch datasets and loaders ---\")\n\nclass LIDC_Dataset(Dataset):\n    def __init__(self, dataframe, transform=None):\n        self.df = dataframe.reset_index(drop=True)\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img = Image.open(row['image_path']).convert('L')  # grayscale\n        if self.transform:\n            img = self.transform(img)\n        label = torch.tensor([row['Cancer_Label']], dtype=torch.float32)\n        return img, label\n\n# Resize 224×224 → 8×8 to match QSAN input\ntransform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE), antialias=True),\n    transforms.ToTensor(),  # → [1, 8, 8] in [0,1]\n])\n\n# train_dataset = LIDC_Dataset(train_df, transform=transform)\ntrain_dataset = LIDC_Dataset(train_df.iloc[:480], transform=transform)  \ntest_dataset = LIDC_Dataset(test_df, transform=transform)\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\ntest_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# Build full test tensor for end-of-epoch evaluation (faster than iterating)\nX_test_list, Y_test_list = [], []\nfor img, lbl in test_dataset:\n    X_test_list.append(img)\n    Y_test_list.append(lbl)\nX_test = torch.stack(X_test_list)\nY_test = torch.stack(Y_test_list)\nprint(f\"  Train batches: {len(train_loader)}\")\nprint(f\"  X_test shape: {X_test.shape}, Y_test shape: {Y_test.shape}\")\n\n# Visualize a sample batch\nfig, axes = plt.subplots(2, 4, figsize=(12, 6))\nfor i, (img, lbl) in enumerate(train_dataset):\n    if i >= 8: break\n    axes[i//4, i%4].imshow(img.squeeze(), cmap='gray')\n    axes[i//4, i%4].set_title(f\"Label: {int(lbl.item())}\")\n    axes[i//4, i%4].axis('off')\nplt.tight_layout()\nplt.savefig('sample_batch.png', dpi=120)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T13:48:03.094107Z","iopub.execute_input":"2026-05-25T13:48:03.094805Z","iopub.status.idle":"2026-05-25T13:48:14.8862Z","shell.execute_reply.started":"2026-05-25T13:48:03.094779Z","shell.execute_reply":"2026-05-25T13:48:14.885343Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Cell 1 : Model Architecture","metadata":{}},{"cell_type":"code","source":"# --- 5. MODEL ARCHITECTURE ---\nN_QUBITS = 4\nTOTAL_PARAMS = 12\nN_QUANTUM_OBS = 4   # matches what qsan_circuit currently returns\n\n\nclass QSAN_TransformerBlock(nn.Module):\n    \"\"\"\n    Full transformer encoder block with QSAN as the attention mechanism.\n    Architecture: x → LN1 → QSAN → linear_out → +residual → LN2 → FFN → +residual\n    \"\"\"\n    def __init__(self, patch_dim=16, d_model=32):\n        super().__init__()\n        self.q_params = nn.Parameter(torch.randn(TOTAL_PARAMS) * 0.1)\n        self.patch_dim = patch_dim\n\n        # Project raw patch to d_model for residual stream\n        self.input_to_dmodel = nn.Linear(patch_dim, d_model)\n\n        # Attention sub-block\n        self.norm1 = nn.LayerNorm(d_model)\n        self.qsan_to_dmodel = nn.Linear(N_QUANTUM_OBS, d_model)\n\n        # FFN sub-block\n        self.norm2 = nn.LayerNorm(d_model)\n        self.ffn = nn.Sequential(\n            nn.Linear(d_model, d_model * 2),\n            nn.GELU(),\n            nn.Dropout(0.1),\n            nn.Linear(d_model * 2, d_model),\n            nn.Dropout(0.1),\n        )\n\n    def forward(self, x_patches):\n        # x_patches: [N_patches, patch_dim] — already unit-norm\n        residual = self.input_to_dmodel(x_patches)\n\n        # QSAN attention\n        q_outputs = []\n        for x in x_patches:\n            q_out = qsan_circuit(inputs=x, params=self.q_params)\n            q_outputs.append(torch.stack(q_out).float())\n        q_feats = torch.stack(q_outputs)                  # [N, 4]\n        attn_out = self.qsan_to_dmodel(q_feats)           # [N, d_model]\n\n        # Residual #1 (around attention)\n        x = residual + attn_out\n\n        # Residual #2 (around FFN, pre-norm)\n        x = x + self.ffn(self.norm2(x))\n        return x\n\n\nclass AttentionPool(nn.Module):\n    def __init__(self, d_model):\n        super().__init__()\n        self.score = nn.Linear(d_model, 1)\n\n    def forward(self, x):                                  # [B, S, d_model]\n        scores = self.score(x)\n        weights = torch.softmax(scores, dim=1)\n        return (weights * x).sum(dim=1)\n\n\nclass HybridRNN_Q_Model(nn.Module):\n    def __init__(self, patch_size=4, patch_features=16, d_model=32):\n        super().__init__()\n        self.patch_size = patch_size\n        patch_dim = patch_size * patch_size\n        self.rnn_context_encoder = nn.GRU(\n            input_size=patch_dim, hidden_size=patch_features,\n            num_layers=1, batch_first=True,\n        )\n        self.transformer = QSAN_TransformerBlock(patch_dim=patch_dim, d_model=d_model)\n        self.attn_pool = AttentionPool(d_model)\n        self.head = nn.Sequential(\n            nn.LayerNorm(d_model),\n            nn.Linear(d_model, 1),\n        )\n\n    def forward(self, image_batch):\n        if image_batch.dim() == 3:\n            image_batch = image_batch.unsqueeze(1)\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        batch_size = image_batch.size(0)\n        sequence_len = 4\n        patch_dim = self.patch_size * self.patch_size\n        rnn_input = patches.reshape(batch_size, sequence_len, patch_dim)\n\n        context_aware_patches, _ = self.rnn_context_encoder(rnn_input)\n        context_aware_patches = torch.tanh(context_aware_patches)\n\n        all_patches_batch = context_aware_patches.reshape(batch_size * sequence_len, patch_dim)\n\n        # NaN defenses + unit-norm\n        all_patches_batch = torch.nan_to_num(all_patches_batch, nan=0.0, posinf=0.0, neginf=0.0)\n        row_norms = all_patches_batch.norm(dim=1, keepdim=True)\n        zero_mask = (row_norms < 1e-8).expand_as(all_patches_batch)\n        fallback = torch.full_like(all_patches_batch, 1.0 / (patch_dim ** 0.5))\n        all_patches_batch = torch.where(zero_mask, fallback, all_patches_batch)\n        row_norms = all_patches_batch.norm(dim=1, keepdim=True).clamp(min=1e-6)\n        all_patches_batch = all_patches_batch / row_norms\n\n        if not torch.isfinite(all_patches_batch).all():\n            raise RuntimeError(\"non-finite values reached the quantum boundary\")\n\n        # Full transformer block\n        patch_feats = self.transformer(all_patches_batch)\n        patch_feats = patch_feats.reshape(batch_size, sequence_len, -1)\n\n        img_feat = self.attn_pool(patch_feats)\n        return self.head(img_feat)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T13:48:45.27894Z","iopub.execute_input":"2026-05-25T13:48:45.279697Z","iopub.status.idle":"2026-05-25T13:48:45.296971Z","shell.execute_reply.started":"2026-05-25T13:48:45.279664Z","shell.execute_reply":"2026-05-25T13:48:45.29614Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluation Functions","metadata":{}},{"cell_type":"code","source":"# --- 6. EVALUATION UTILITIES ---\n\ndef compute_metrics(y_true, y_pred_probs, threshold=0.5):\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    tn, fp, fn, tp = confusion_matrix(y_true_np, y_pred_np, labels=[0, 1]).ravel()\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),\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, threshold=0.5):\n    y_true_np = y_true.cpu().numpy().flatten()\n    y_pred_np = (y_pred_probs.cpu().detach().numpy().flatten() >= threshold).astype(int)\n    cm = confusion_matrix(y_true_np, y_pred_np, labels=[0, 1])\n    plt.figure(figsize=(6, 5))\n    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',\n                xticklabels=['Pred No-Cancer', 'Pred Cancer'],\n                yticklabels=['Actual No-Cancer', 'Actual Cancer'])\n    plt.xlabel(\"Predicted\"); plt.ylabel(\"True\")\n    plt.title(f\"Confusion Matrix — {label}\")\n    plt.savefig(f'cm_{label}.png', dpi=150, bbox_inches='tight')\n    plt.close()\n\n\ndef plot_qbsasm_aggregate(model, X_data, Y_data, label, max_samples_per_class=100):\n    \"\"\"\n    Aggregate QBSASM across many samples per class to show class-level attention patterns.\n    Plots two side-by-side heatmaps: avg attention for no-cancer vs cancer samples.\n    \"\"\"\n    model.eval()\n    densities_per_class = {0: [], 1: []}\n    \n    with torch.no_grad():\n        for cls in [0, 1]:\n            idxs = (Y_data.flatten() == cls).nonzero(as_tuple=True)[0]\n            idxs = idxs[:max_samples_per_class]\n            for idx in idxs:\n                img = X_data[idx]\n                if img.dim() == 2: img = img.unsqueeze(0).unsqueeze(0)\n                elif img.dim() == 3: img = img.unsqueeze(0)\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)\n                context_aware, _ = model.rnn_context_encoder(rnn_input)\n                all_patches = context_aware.reshape(4, 16)\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                    densities_per_class[cls].append(dm)\n    \n    # Average densities per class\n    avg_dm = {}\n    for cls in [0, 1]:\n        if densities_per_class[cls]:\n            avg_dm[cls] = torch.stack(densities_per_class[cls]).mean(dim=0)\n        else:\n            print(f\"  WARNING: no samples for class {cls}\")\n            return\n    \n    # Extract diagonal probabilities → attention scores\n    fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n    class_names = ['No-Cancer', 'Cancer']\n    \n    for cls in [0, 1]:\n        probs = torch.diag(avg_dm[cls]).real.cpu().numpy()\n        side = int(np.sqrt(len(probs)))\n        qbsasm = probs.reshape(side, side)\n        n_bits = int(np.log2(side))\n        labels_bits = [format(i, f'0{n_bits}b') for i in range(side)]\n        \n        sns.heatmap(qbsasm, ax=axes[cls], cmap='viridis',\n                    xticklabels=labels_bits, yticklabels=labels_bits,\n                    cbar_kws={'label': 'Attention score'}, square=True)\n        axes[cls].set_xlabel(\"|Q⟩\"); axes[cls].set_ylabel(\"|K⟩\")\n        axes[cls].set_title(f\"QBSASM — {class_names[cls]} (avg over {len(densities_per_class[cls])} samples)\")\n    \n    plt.suptitle(f\"Class-Level Attention Patterns — {label}\", y=1.02)\n    plt.tight_layout()\n    plt.savefig(f'qbsasm_{label}.png', dpi=150, bbox_inches='tight')\n    plt.close()\n    print(f\"  Saved QBSASM aggregate to qbsasm_{label}.png\")\n\n\ndef print_results_table(seed_metrics_list):\n    \"\"\"Prints mean ± std across seeds for a single task.\"\"\"\n    metric_names = ['accuracy', 'auroc', 'f1', 'sensitivity', 'specificity']\n    print(\"\\n\" + \"=\" * 80)\n    print(f\"  LIDC-IDRI RESULTS (mean ± std across {len(seed_metrics_list)} seeds)\")\n    print(\"=\" * 80)\n    for m in metric_names:\n        vals = [s[m] for s in seed_metrics_list]\n        print(f\"  {m:<14}: {np.mean(vals):.4f} ± {np.std(vals):.4f}\")\n    print(\"=\" * 80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T13:48:50.485524Z","iopub.execute_input":"2026-05-25T13:48:50.486144Z","iopub.status.idle":"2026-05-25T13:48:50.506616Z","shell.execute_reply.started":"2026-05-25T13:48:50.486111Z","shell.execute_reply":"2026-05-25T13:48:50.505593Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Loop","metadata":{}},{"cell_type":"code","source":"# --- 7. TRAIN RQSAN ON LIDC-IDRI ---\nprint(\"--- Training RQSAN on LIDC-IDRI ---\")\n\nN_SEEDS = 1\nNUM_EPOCHS = 1\n\nseed_metrics = []\nseed_confusion_matrices = []\nlast_successful_model = None\n\nfor seed in range(N_SEEDS):\n    print(f\"\\n{'='*70}\\n  SEED {seed+1}/{N_SEEDS}\\n{'='*70}\")\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.0005)\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, train_acc_history = [], []\n    test_loss_history, 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, desc=f\"[Seed {seed+1}] Epoch {epoch+1}/{NUM_EPOCHS}\", 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                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}\")\n                        model.load_state_dict(best_model_state)\n                    nan_triggered = True\n                    break\n                \n                loss.backward()\n                bad_grads = False\n                for p in model.parameters():\n                    if p.grad is not None and not torch.isfinite(p.grad).all():\n                        bad_grads = True\n                        break\n                \n                if bad_grads:\n                    print(f\"\\n  Non-finite gradient on batch — zeroing and skipping optimizer step.\")\n                    optimizer.zero_grad()\n                    continue\n                \n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.25)                \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                pbar.set_postfix({'loss': f\"{loss.item():.4f}\", 'acc': f\"{epoch_train_correct/total_train_samples:.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 on test set\n            model.eval()\n            with torch.no_grad():\n                test_outputs = model(X_test)\n                test_labels_t = Y_test.float().view_as(test_outputs)\n                test_loss = criterion(test_outputs, test_labels_t)\n                test_probs = torch.sigmoid(test_outputs)\n                test_preds = torch.round(test_probs)\n                test_accuracy = (test_preds == test_labels_t).sum().item() / len(test_labels_t)\n            \n            # Full test metrics (uses compute_metrics we already have)\n            try:\n                epoch_metrics, _ = compute_metrics(Y_test, test_probs)\n            except Exception as e:\n                print(f\"  (could not compute full metrics: {e})\")\n                epoch_metrics = {'auroc': float('nan'), 'f1': float('nan'),\n                                 'sensitivity': float('nan'), 'specificity': float('nan'),\n                                 'precision': float('nan')}\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} Acc: {avg_train_acc:.4f} | \"\n                  f\"Test Loss: {test_loss.item():.4f} Acc: {test_accuracy:.4f}\")\n            print(f\"             | AUROC: {epoch_metrics['auroc']:.4f}  F1: {epoch_metrics['f1']:.4f}  \"\n                  f\"Sens: {epoch_metrics['sensitivity']:.4f}  Spec: {epoch_metrics['specificity']:.4f}  \"\n                  f\"Prec: {epoch_metrics['precision']:.4f}\")\n            \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            current_auroc = epoch_metrics['auroc']\n            if not np.isnan(current_auroc) and current_auroc > best_test_accuracy:\n                best_test_accuracy = current_auroc  # variable kept named the same for code simplicity\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}: AUROC={best_test_accuracy:.4f}\")\n            else:\n                epochs_without_improvement += 1\n            \n            if epochs_without_improvement >= EARLY_STOP_PATIENCE:\n                print(f\"  Plateau — stopping.\")\n                early_stopped = True\n            \n            scheduler.step(avg_train_loss)\n    \n    except Exception as e:\n        print(f\"\\n  ERROR on seed {seed+1}: {e}\")\n        if best_model_state is None:\n            print(f\"  No best model — skipping seed.\")\n            continue\n    \n    # Final evaluation\n    if best_model_state is None:\n        print(f\"  WARNING: No valid model for seed {seed+1}, skipping.\")\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 probs, skipping seed {seed+1}.\")\n        continue\n    \n    metrics, cm = compute_metrics(Y_test, final_probs)\n    seed_metrics.append(metrics)\n    seed_confusion_matrices.append(cm)\n    last_successful_model = model\n    \n    print(f\"\\n  ★ 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} | Spec={metrics['specificity']:.4f}\")\n    \n    plot_confusion_matrix(Y_test, final_probs, label=f\"seed_{seed+1}\")\n    torch.save(best_model_state, f\"best_model_seed{seed+1}.pth\")\n\n# --- Aggregate results ---\nif seed_confusion_matrices:\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=['Pred No-Cancer', 'Pred Cancer'],\n                yticklabels=['Actual No-Cancer', 'Actual Cancer'])\n    plt.xlabel(\"Predicted\"); plt.ylabel(\"True\")\n    plt.title(f\"Aggregate Confusion Matrix ({len(seed_confusion_matrices)} seeds)\")\n    plt.savefig('cm_aggregate.png', dpi=150, bbox_inches='tight')\n    plt.close()\n    print(f\"\\n  Aggregate CM saved ({len(seed_confusion_matrices)} valid seeds).\")\n\nprint_results_table(seed_metrics)\n\n# Save results JSON for reproducibility\nresults_json = {\n    'n_seeds_completed': len(seed_metrics),\n    'seed_metrics': [{k: float(v) for k, v in m.items()} for m in seed_metrics],\n    'mean_metrics': {m: float(np.mean([s[m] for s in seed_metrics])) for m in ['accuracy','auroc','f1','sensitivity','specificity']},\n    'std_metrics': {m: float(np.std([s[m] for s in seed_metrics])) for m in ['accuracy','auroc','f1','sensitivity','specificity']},\n}\nwith open('lidc_rqsan_results.json', 'w') as f:\n    json.dump(results_json, f, indent=2)\nprint(\"Results saved to lidc_rqsan_results.json\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T13:49:06.178941Z","iopub.execute_input":"2026-05-25T13:49:06.179286Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 8. QBSASM CLASS-LEVEL ATTENTION VISUALIZATION ---\nprint(\"\\n--- Generating QBSASM heatmaps (class-aggregate) ---\")\n\nif last_successful_model is not None:\n    plot_qbsasm_aggregate(last_successful_model, X_test, Y_test, \n                          label=\"lidc_idri\", max_samples_per_class=100)\nelse:\n    print(\"  No successful model — skipping QBSASM.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-25T12:03:26.992495Z","iopub.status.idle":"2026-05-25T12:03:26.992931Z","shell.execute_reply.started":"2026-05-25T12:03:26.992714Z","shell.execute_reply":"2026-05-25T12:03:26.992735Z"}},"outputs":[],"execution_count":null}]}