{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'fixed_counts_to_2d_array'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mqiskit_aer\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AerSimulator\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpyplot\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mplt\u001b[39;00m\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfixed_counts_to_2d_array\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m counts_to_2d_array\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minitialize_wavefunction\u001b[39m(qc, qubits, wave_function):\n\u001b[1;32m      9\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Initialize the quantum circuit with the given wave function.\"\"\"\u001b[39;00m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'fixed_counts_to_2d_array'"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, assemble\n",
    "from qiskit.visualization import plot_histogram\n",
    "from qiskit_aer import AerSimulator\n",
    "import matplotlib.pyplot as plt\n",
    "from fixed_counts_to_2d_array import counts_to_2d_array\n",
    "\n",
    "def initialize_wavefunction(qc, qubits, wave_function):\n",
    "    \"\"\"Initialize the quantum circuit with the given wave function.\"\"\"\n",
    "    for i, amplitude in enumerate(wave_function):\n",
    "        qc.ry(2 * np.arccos(amplitude), qubits[i])\n",
    "\n",
    "def apply_phase_shift(qc, qubits, phase_shifts):\n",
    "    \"\"\"Apply phase shift to simulate the weak phase object.\"\"\"\n",
    "    for i, phase in enumerate(phase_shifts):\n",
    "        qc.rz(phase, qubits[i])\n",
    "\n",
    "def simulate_electron_wave(phase_shifts, n_qubits=4):\n",
    "    # Initialize quantum circuit\n",
    "    qr = QuantumRegister(n_qubits, 'q')\n",
    "    cr = ClassicalRegister(n_qubits, 'c')\n",
    "    qc = QuantumCircuit(qr, cr)\n",
    "    \n",
    "    # Define initial wave function (Gaussian-like distribution)\n",
    "    initial_wave_function = np.array([np.exp(-0.5 * (i - n_qubits / 2) ** 2) for i in range(n_qubits)])\n",
    "    initial_wave_function /= np.linalg.norm(initial_wave_function)\n",
    "    \n",
    "    print(\"Initial wave function:\", initial_wave_function)  # Debug print\n",
    "    \n",
    "    # Initialize wave function in the quantum circuit\n",
    "    initialize_wavefunction(qc, qr, initial_wave_function)\n",
    "    \n",
    "    # Apply phase shift to simulate the weak phase object\n",
    "    apply_phase_shift(qc, qr, phase_shifts)\n",
    "    \n",
    "    # Measure the results\n",
    "    qc.measure(qr, cr)\n",
    "    \n",
    "    print(\"Quantum circuit:\\n\", qc)  # Debug print\n",
    "    \n",
    "    return qc\n",
    "\n",
    "# Define phase shifts to simulate the weak phase object\n",
    "phase_shifts = [0.1, 0.2, 0.3, 0.4]\n",
    "\n",
    "# Simulate electron wave\n",
    "qc = simulate_electron_wave(phase_shifts, n_qubits=4)\n",
    "\n",
    "# Execute the circuit using AerSimulator\n",
    "simulator = AerSimulator()\n",
    "compiled_circuit = transpile(qc, simulator)\n",
    "print(\"Compiled circuit:\\n\", compiled_circuit)  # Debug print\n",
    "sim_result = simulator.run(compiled_circuit).result()\n",
    "counts = sim_result.get_counts()\n",
    "\n",
    "print(\"Counts:\", counts)  # Debug print\n",
    "\n",
    "# Manually plot the histogram\n",
    "if counts:\n",
    "    outcomes = list(counts.keys())\n",
    "    values = list(counts.values())\n",
    "    \n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.bar(outcomes, values)\n",
    "    plt.title(\"Electron Wave Simulation Results\")\n",
    "    plt.xlabel(\"Measurement Outcome\")\n",
    "    plt.ylabel(\"Counts\")\n",
    "    plt.xticks(rotation=90)\n",
    "    plt.show()\n",
    "else:\n",
    "    print(\"No counts to display.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#expandig to 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial 2D wave function:\n",
      " [[0.14566959 0.22718971 0.22718971 0.14566959]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.14566959 0.22718971 0.22718971 0.14566959]]\n",
      "Quantum circuit:\n",
      "       ┌────────────┐┌─────────┐┌─┐                                             \n",
      " q_0: ┤ Ry(2.8492) ├┤ Rz(0.1) ├┤M├─────────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ Ry(2.4172) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ Ry(2.4172) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ Ry(2.8492) ├┤ Rz(0.7) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └────────────┘└─────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/══════════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                                0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Compiled circuit:\n",
      " global phase: 3.0832\n",
      "      ┌──────────────────┐┌─┐                                             \n",
      " q_0: ┤ U3(2.8492,0.1,0) ├┤M├─────────────────────────────────────────────\n",
      "      ├──────────────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ U3(2.6832,0.2,0) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├──────────────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ U3(2.6832,0.3,0) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ U3(2.6832,0.2,0) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ U3(2.4172,0.3,0) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ U3(2.6832,0.3,0) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ U3(2.4172,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ U3(2.8492,0.7,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └──────────────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/═════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                           0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Counts: {'1111111110001111': 1, '1110111111110111': 1, '1011111111110011': 1, '1111110111100111': 1, '1111111010011111': 1, '1111110010111011': 1, '1011101110111111': 1, '1111001101111111': 1, '1101101111111111': 1, '1101111111101111': 1, '1111111110110011': 1, '1011111010111111': 1, '1111110101111011': 1, '1111100111011111': 1, '1011101111111011': 2, '1010111111011111': 1, '1111101001101111': 1, '1111111000111111': 1, '1111110100111111': 1, '1011101111111111': 2, '1101011111011111': 1, '1111101011111111': 1, '1101111110111101': 1, '1111101111111110': 1, '1011111111110111': 1, '1111110010111111': 1, '0111111110111111': 1, '1110111100111111': 1, '1101110110111111': 1, '1111110111011101': 1, '1011101111011111': 2, '1111011101111111': 1, '1111101101111111': 1, '1111110111011011': 1, '1111101101111101': 1, '1101111111111110': 2, '1111110111101111': 3, '1011100100111111': 1, '1111001111011111': 1, '1111101101111011': 1, '1111101111110111': 2, '1011110110111111': 1, '1110111111111111': 12, '0110111111111111': 1, '1111111111011110': 1, '1011111111111101': 1, '1111001110111111': 1, '0111101110111111': 1, '1111111011111011': 1, '1111111110011111': 6, '1111111110010111': 1, '1111111101101101': 1, '1111011111111011': 1, '1111110111111110': 1, '1111111111011101': 5, '1111111101111011': 1, '1011110111111111': 5, '1101100111111111': 1, '1101111001111111': 1, '1111110110111011': 1, '1111100110111011': 1, '1101011111111111': 2, '1111111101111111': 8, '1011110111011111': 1, '1110110111011111': 1, '1111111010111111': 3, '1110101111011111': 1, '1111011111101111': 2, '1011111110111111': 1, '1011111111011111': 2, '1011011111111101': 1, '1111111101110111': 1, '1110111110111111': 1, '1111110110011111': 4, '0111111110011111': 1, '0111101111111111': 2, '1111101111111101': 4, '1110101110111111': 1, '1101111111111111': 20, '1110101111111111': 2, '1011111111010111': 1, '1111010111111111': 6, '1011111011111111': 1, '1111101110111111': 7, '1111111111001111': 2, '1111110111111111': 55, '1101011101111111': 1, '0111111011111101': 1, '1111011111111110': 1, '1111111011011111': 3, '1110110100111101': 1, '0111111111111101': 1, '0011111111111111': 2, '0111101111110111': 1, '1111111011111111': 25, '1111011111011111': 5, '1011001110111101': 1, '1101101110111111': 1, '1111111110110111': 2, '1011011111011111': 1, '1111110101111111': 3, '1111101111101011': 1, '1111110111110111': 2, '1111101010111111': 2, '1111111110011101': 1, '0111111111011111': 1, '1111101111101111': 1, '1111111111100111': 2, '1101110111111111': 2, '1110111111011111': 1, '1111111111111110': 9, '1111110111111101': 2, '1111011110111111': 6, '1111100110111111': 1, '1111111111011111': 50, '1111100011111111': 2, '1110110111101111': 1, '1111101101011111': 3, '1111101111111011': 1, '1111111110111101': 4, '0101111110101111': 1, '0111111111111111': 7, '1111101111111111': 52, '1111000111111111': 2, '0111101010111111': 1, '1111111110111111': 54, '1111010111011111': 1, '1111011110111101': 1, '1111111111011011': 2, '1111010101111111': 1, '1110101110111011': 1, '1011111101011111': 1, '1111111110111011': 2, '1111011111111111': 11, '1111111111101111': 16, '1111111111111000': 1, '1111111100101111': 1, '1111110111011111': 9, '0111011100111111': 1, '1101101111011111': 1, '1111111101011111': 1, '1101111111011111': 5, '1101111111111011': 2, '1111111111111101': 21, '1101111110111111': 6, '1111110011111111': 4, '1111111001111111': 1, '1011111101111111': 1, '1111111100111111': 3, '1111111101101111': 2, '0111111111111011': 2, '1111001100011111': 1, '1111111110101111': 2, '0111011011111111': 1, '1111111111110111': 5, '1111110111111011': 4, '1111110110111111': 11, '1011111111111111': 20, '1111101100101111': 1, '1111101111011111': 10, '1111100111111111': 6, '1111111101111101': 2, '1111111111111011': 30, '1111111111111111': 354}\n",
      "Expected number of qubits: 16\n",
      "Actual bitstring length: 16\n",
      "Bitstring: 1111111110001111, Index: 65423, Row: 16355, Col: 3, Count: 1\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "index 16355 is out of bounds for axis 0 with size 4",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 87\u001b[0m\n\u001b[1;32m     83\u001b[0m         result_array[row, col] \u001b[38;5;241m=\u001b[39m count\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n\u001b[0;32m---> 87\u001b[0m result_array \u001b[38;5;241m=\u001b[39m \u001b[43mcounts_to_2d_array\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcounts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mResult array:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, result_array)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n\u001b[1;32m     91\u001b[0m \u001b[38;5;66;03m# Plot the initial wave function and the result\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[2], line 83\u001b[0m, in \u001b[0;36mcounts_to_2d_array\u001b[0;34m(counts, grid_size)\u001b[0m\n\u001b[1;32m     81\u001b[0m     col \u001b[38;5;241m=\u001b[39m index \u001b[38;5;241m%\u001b[39m grid_size[\u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m     82\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBitstring: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtruncated_bitstring\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Index: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mindex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Row: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrow\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Col: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcol\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Count: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcount\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n\u001b[0;32m---> 83\u001b[0m     \u001b[43mresult_array\u001b[49m\u001b[43m[\u001b[49m\u001b[43mrow\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcol\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m count\n\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\u001b[0;31mIndexError\u001b[0m: index 16355 is out of bounds for axis 0 with size 4"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, assemble\n",
    "from qiskit.visualization import plot_histogram\n",
    "from qiskit_aer import AerSimulator\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def initialize_wavefunction_2d(qc, qubits, wave_function):\n",
    "    \"\"\"Initialize the quantum circuit with the given 2D wave function.\"\"\"\n",
    "    for i, amplitude in enumerate(wave_function.flatten()):\n",
    "        qc.ry(2 * np.arccos(amplitude), qubits[i])\n",
    "\n",
    "def apply_phase_shift_2d(qc, qubits, phase_shifts):\n",
    "    \"\"\"Apply phase shift to simulate the weak phase object in 2D.\"\"\"\n",
    "    for i, phase in enumerate(phase_shifts.flatten()):\n",
    "        qc.rz(phase, qubits[i])\n",
    "\n",
    "def simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4)):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    \n",
    "    # Initialize quantum circuit\n",
    "    qr = QuantumRegister(n_qubits, 'q')\n",
    "    cr = ClassicalRegister(n_qubits, 'c')\n",
    "    qc = QuantumCircuit(qr, cr)\n",
    "    \n",
    "    # Define initial 2D wave function (Gaussian-like distribution)\n",
    "    x = np.linspace(-1, 1, grid_size[0])\n",
    "    y = np.linspace(-1, 1, grid_size[1])\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    initial_wave_function = np.exp(-0.5 * (X**2 + Y**2))\n",
    "    initial_wave_function /= np.linalg.norm(initial_wave_function)\n",
    "    \n",
    "    print(\"Initial 2D wave function:\\n\", initial_wave_function)  # Debug print\n",
    "    \n",
    "    # Initialize wave function in the quantum circuit\n",
    "    initialize_wavefunction_2d(qc, qr, initial_wave_function)\n",
    "    \n",
    "    # Apply phase shift to simulate the weak phase object\n",
    "    apply_phase_shift_2d(qc, qr, phase_shifts)\n",
    "    \n",
    "    # Measure the results\n",
    "    qc.measure(qr, cr)\n",
    "    \n",
    "    print(\"Quantum circuit:\\n\", qc)  # Debug print\n",
    "    \n",
    "    return qc, initial_wave_function\n",
    "\n",
    "# Define phase shifts to simulate the weak phase object in 2D\n",
    "phase_shifts = np.array([[0.1, 0.2, 0.3, 0.4],\n",
    "                         [0.2, 0.3, 0.4, 0.5],\n",
    "                         [0.3, 0.4, 0.5, 0.6],\n",
    "                         [0.4, 0.5, 0.6, 0.7]])\n",
    "\n",
    "# Simulate electron wave in 2D\n",
    "qc, initial_wave_function = simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4))\n",
    "\n",
    "# Execute the circuit using AerSimulator\n",
    "simulator = AerSimulator()\n",
    "compiled_circuit = transpile(qc, simulator)\n",
    "print(\"Compiled circuit:\\n\", compiled_circuit)  # Debug print\n",
    "sim_result = simulator.run(compiled_circuit).result()\n",
    "counts = sim_result.get_counts()\n",
    "\n",
    "print(\"Counts:\", counts)  # Debug print\n",
    "\n",
    "# Verify the number of qubits\n",
    "n_qubits = 4 * 4\n",
    "bitstring_length = len(list(counts.keys())[0])\n",
    "print(f\"Expected number of qubits: {n_qubits}\")\n",
    "print(f\"Actual bitstring length: {bitstring_length}\")\n",
    "\n",
    "# Convert counts to 2D array for visualization\n",
    "def counts_to_2d_array(counts, grid_size):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    result_array = np.zeros((grid_size[0], grid_size[1]))\n",
    "    \n",
    "    for bitstring, count in counts.items():\n",
    "        # Truncate the bitstring to the expected length\n",
    "        truncated_bitstring = bitstring[-n_qubits:]\n",
    "        index = int(truncated_bitstring, 2)\n",
    "        row = index // grid_size[1]\n",
    "        col = index % grid_size[1]\n",
    "        print(f\"Bitstring: {truncated_bitstring}, Index: {index}, Row: {row}, Col: {col}, Count: {count}\")  # Debug print\n",
    "        result_array[row, col] = count\n",
    "    \n",
    "    return result_array\n",
    "\n",
    "result_array = counts_to_2d_array(counts, (4, 4))\n",
    "\n",
    "print(\"Result array:\\n\", result_array)  # Debug print\n",
    "\n",
    "# Plot the initial wave function and the result\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "ax[0].imshow(initial_wave_function, cmap='viridis', origin='lower')\n",
    "ax[0].set_title(\"Initial 2D Wave Function\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "ax[1].imshow(result_array, cmap='viridis', origin='lower')\n",
    "ax[1].set_title(\"Electron Wave Simulation Results\")\n",
    "ax[1].axis('off')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial 2D wave function:\n",
      " [[0.14566959 0.22718971 0.22718971 0.14566959]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.14566959 0.22718971 0.22718971 0.14566959]]\n",
      "Quantum circuit:\n",
      "       ┌────────────┐┌─────────┐┌─┐                                             \n",
      " q_0: ┤ Ry(2.8492) ├┤ Rz(0.1) ├┤M├─────────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ Ry(2.4172) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ Ry(2.4172) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ Ry(2.8492) ├┤ Rz(0.7) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └────────────┘└─────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/══════════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                                0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Compiled circuit:\n",
      " global phase: 3.0832\n",
      "      ┌──────────────────┐┌─┐                                             \n",
      " q_0: ┤ U3(2.8492,0.1,0) ├┤M├─────────────────────────────────────────────\n",
      "      ├──────────────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ U3(2.6832,0.2,0) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├──────────────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ U3(2.6832,0.3,0) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ U3(2.6832,0.2,0) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ U3(2.4172,0.3,0) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ U3(2.6832,0.3,0) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ U3(2.4172,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ U3(2.8492,0.7,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └──────────────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/═════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                           0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Counts: {'1111111110001111': 1, '1110111111110111': 1, '1011111111110011': 1, '1111110111100111': 1, '1111111010011111': 1, '1111110010111011': 1, '1011101110111111': 1, '1111001101111111': 1, '1101101111111111': 1, '1101111111101111': 1, '1111111110110011': 1, '1011111010111111': 1, '1111110101111011': 1, '1111100111011111': 1, '1011101111111011': 2, '1010111111011111': 1, '1111101001101111': 1, '1111111000111111': 1, '1111110100111111': 1, '1011101111111111': 2, '1101011111011111': 1, '1111101011111111': 1, '1101111110111101': 1, '1111101111111110': 1, '1011111111110111': 1, '1111110010111111': 1, '0111111110111111': 1, '1110111100111111': 1, '1101110110111111': 1, '1111110111011101': 1, '1011101111011111': 2, '1111011101111111': 1, '1111101101111111': 1, '1111110111011011': 1, '1111101101111101': 1, '1101111111111110': 2, '1111110111101111': 3, '1011100100111111': 1, '1111001111011111': 1, '1111101101111011': 1, '1111101111110111': 2, '1011110110111111': 1, '1110111111111111': 12, '0110111111111111': 1, '1111111111011110': 1, '1011111111111101': 1, '1111001110111111': 1, '0111101110111111': 1, '1111111011111011': 1, '1111111110011111': 6, '1111111110010111': 1, '1111111101101101': 1, '1111011111111011': 1, '1111110111111110': 1, '1111111111011101': 5, '1111111101111011': 1, '1011110111111111': 5, '1101100111111111': 1, '1101111001111111': 1, '1111110110111011': 1, '1111100110111011': 1, '1101011111111111': 2, '1111111101111111': 8, '1011110111011111': 1, '1110110111011111': 1, '1111111010111111': 3, '1110101111011111': 1, '1111011111101111': 2, '1011111110111111': 1, '1011111111011111': 2, '1011011111111101': 1, '1111111101110111': 1, '1110111110111111': 1, '1111110110011111': 4, '0111111110011111': 1, '0111101111111111': 2, '1111101111111101': 4, '1110101110111111': 1, '1101111111111111': 20, '1110101111111111': 2, '1011111111010111': 1, '1111010111111111': 6, '1011111011111111': 1, '1111101110111111': 7, '1111111111001111': 2, '1111110111111111': 55, '1101011101111111': 1, '0111111011111101': 1, '1111011111111110': 1, '1111111011011111': 3, '1110110100111101': 1, '0111111111111101': 1, '0011111111111111': 2, '0111101111110111': 1, '1111111011111111': 25, '1111011111011111': 5, '1011001110111101': 1, '1101101110111111': 1, '1111111110110111': 2, '1011011111011111': 1, '1111110101111111': 3, '1111101111101011': 1, '1111110111110111': 2, '1111101010111111': 2, '1111111110011101': 1, '0111111111011111': 1, '1111101111101111': 1, '1111111111100111': 2, '1101110111111111': 2, '1110111111011111': 1, '1111111111111110': 9, '1111110111111101': 2, '1111011110111111': 6, '1111100110111111': 1, '1111111111011111': 50, '1111100011111111': 2, '1110110111101111': 1, '1111101101011111': 3, '1111101111111011': 1, '1111111110111101': 4, '0101111110101111': 1, '0111111111111111': 7, '1111101111111111': 52, '1111000111111111': 2, '0111101010111111': 1, '1111111110111111': 54, '1111010111011111': 1, '1111011110111101': 1, '1111111111011011': 2, '1111010101111111': 1, '1110101110111011': 1, '1011111101011111': 1, '1111111110111011': 2, '1111011111111111': 11, '1111111111101111': 16, '1111111111111000': 1, '1111111100101111': 1, '1111110111011111': 9, '0111011100111111': 1, '1101101111011111': 1, '1111111101011111': 1, '1101111111011111': 5, '1101111111111011': 2, '1111111111111101': 21, '1101111110111111': 6, '1111110011111111': 4, '1111111001111111': 1, '1011111101111111': 1, '1111111100111111': 3, '1111111101101111': 2, '0111111111111011': 2, '1111001100011111': 1, '1111111110101111': 2, '0111011011111111': 1, '1111111111110111': 5, '1111110111111011': 4, '1111110110111111': 11, '1011111111111111': 20, '1111101100101111': 1, '1111101111011111': 10, '1111100111111111': 6, '1111111101111101': 2, '1111111111111011': 30, '1111111111111111': 354}\n",
      "Expected number of qubits: 16\n",
      "Actual bitstring length: 16\n",
      "Bitstring: 1111111110001111, Index: 65423, Row: 16355, Col: 3, Count: 1\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "index 16355 is out of bounds for axis 0 with size 4",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[2], line 87\u001b[0m\n",
      "\u001b[1;32m     83\u001b[0m         result_array[row, col] \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\u001b[0;32m---> 87\u001b[0m result_array \u001b[38;5;241m=\u001b[39m \u001b[43mcounts_to_2d_array\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcounts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;32m     89\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mResult array:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, result_array)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[1;32m     91\u001b[0m \u001b[38;5;66;03m# Plot the initial wave function and the result\u001b[39;00m\n",
      "\n",
      "Cell \u001b[0;32mIn[2], line 83\u001b[0m, in \u001b[0;36mcounts_to_2d_array\u001b[0;34m(counts, grid_size)\u001b[0m\n",
      "\u001b[1;32m     81\u001b[0m     col \u001b[38;5;241m=\u001b[39m index \u001b[38;5;241m%\u001b[39m grid_size[\u001b[38;5;241m1\u001b[39m]\n",
      "\u001b[1;32m     82\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBitstring: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtruncated_bitstring\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Index: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mindex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Row: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrow\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Col: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcol\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Count: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcount\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[0;32m---> 83\u001b[0m     \u001b[43mresult_array\u001b[49m\u001b[43m[\u001b[49m\u001b[43mrow\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcol\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\n",
      "\u001b[0;31mIndexError\u001b[0m: index 16355 is out of bounds for axis 0 with size 4"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, assemble\n",
    "from qiskit.visualization import plot_histogram\n",
    "from qiskit_aer import AerSimulator\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def initialize_wavefunction_2d(qc, qubits, wave_function):\n",
    "    \"\"\"Initialize the quantum circuit with the given 2D wave function.\"\"\"\n",
    "    for i, amplitude in enumerate(wave_function.flatten()):\n",
    "        qc.ry(2 * np.arccos(amplitude), qubits[i])\n",
    "\n",
    "def apply_phase_shift_2d(qc, qubits, phase_shifts):\n",
    "    \"\"\"Apply phase shift to simulate the weak phase object in 2D.\"\"\"\n",
    "    for i, phase in enumerate(phase_shifts.flatten()):\n",
    "        qc.rz(phase, qubits[i])\n",
    "\n",
    "def simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4)):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    \n",
    "    # Initialize quantum circuit\n",
    "    qr = QuantumRegister(n_qubits, 'q')\n",
    "    cr = ClassicalRegister(n_qubits, 'c')\n",
    "    qc = QuantumCircuit(qr, cr)\n",
    "    \n",
    "    # Define initial 2D wave function (Gaussian-like distribution)\n",
    "    x = np.linspace(-1, 1, grid_size[0])\n",
    "    y = np.linspace(-1, 1, grid_size[1])\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    initial_wave_function = np.exp(-0.5 * (X**2 + Y**2))\n",
    "    initial_wave_function /= np.linalg.norm(initial_wave_function)\n",
    "    \n",
    "    print(\"Initial 2D wave function:\\n\", initial_wave_function)  # Debug print\n",
    "    \n",
    "    # Initialize wave function in the quantum circuit\n",
    "    initialize_wavefunction_2d(qc, qr, initial_wave_function)\n",
    "    \n",
    "    # Apply phase shift to simulate the weak phase object\n",
    "    apply_phase_shift_2d(qc, qr, phase_shifts)\n",
    "    \n",
    "    # Measure the results\n",
    "    qc.measure(qr, cr)\n",
    "    \n",
    "    print(\"Quantum circuit:\\n\", qc)  # Debug print\n",
    "    \n",
    "    return qc, initial_wave_function\n",
    "\n",
    "# Define phase shifts to simulate the weak phase object in 2D\n",
    "phase_shifts = np.array([[0.1, 0.2, 0.3, 0.4],\n",
    "                         [0.2, 0.3, 0.4, 0.5],\n",
    "                         [0.3, 0.4, 0.5, 0.6],\n",
    "                         [0.4, 0.5, 0.6, 0.7]])\n",
    "\n",
    "# Simulate electron wave in 2D\n",
    "qc, initial_wave_function = simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4))\n",
    "\n",
    "# Execute the circuit using AerSimulator\n",
    "simulator = AerSimulator()\n",
    "compiled_circuit = transpile(qc, simulator)\n",
    "print(\"Compiled circuit:\\n\", compiled_circuit)  # Debug print\n",
    "sim_result = simulator.run(compiled_circuit).result()\n",
    "counts = sim_result.get_counts()\n",
    "\n",
    "print(\"Counts:\", counts)  # Debug print\n",
    "\n",
    "# Verify the number of qubits\n",
    "n_qubits = 4 * 4\n",
    "bitstring_length = len(list(counts.keys())[0])\n",
    "print(f\"Expected number of qubits: {n_qubits}\")\n",
    "print(f\"Actual bitstring length: {bitstring_length}\")\n",
    "\n",
    "# Convert counts to 2D array for visualization\n",
    "def counts_to_2d_array(counts, grid_size):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    result_array = np.zeros((grid_size[0], grid_size[1]))\n",
    "    \n",
    "    for bitstring, count in counts.items():\n",
    "        # Truncate the bitstring to the expected length\n",
    "        truncated_bitstring = bitstring[-n_qubits:]\n",
    "        index = int(truncated_bitstring, 2)\n",
    "        row = index // grid_size[1]\n",
    "        col = index % grid_size[1]\n",
    "        print(f\"Bitstring: {truncated_bitstring}, Index: {index}, Row: {row}, Col: {col}, Count: {count}\")  # Debug print\n",
    "        result_array[row, col] = count\n",
    "    \n",
    "    return result_array\n",
    "\n",
    "result_array = counts_to_2d_array(counts, (4, 4))\n",
    "\n",
    "print(\"Result array:\\n\", result_array)  # Debug print\n",
    "\n",
    "# Plot the initial wave function and the result\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "ax[0].imshow(initial_wave_function, cmap='viridis', origin='lower')\n",
    "ax[0].set_title(\"Initial 2D Wave Function\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "ax[1].imshow(result_array, cmap='viridis', origin='lower')\n",
    "ax[1].set_title(\"Electron Wave Simulation Results\")\n",
    "ax[1].axis('off')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial 2D wave function:\n",
      " [[0.14566959 0.22718971 0.22718971 0.14566959]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.14566959 0.22718971 0.22718971 0.14566959]]\n",
      "Quantum circuit:\n",
      "       ┌────────────┐┌─────────┐┌─┐                                             \n",
      " q_0: ┤ Ry(2.8492) ├┤ Rz(0.1) ├┤M├─────────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ Ry(2.4172) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ Ry(2.4172) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ Ry(2.8492) ├┤ Rz(0.7) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └────────────┘└─────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/══════════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                                0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Compiled circuit:\n",
      " global phase: 3.0832\n",
      "      ┌──────────────────┐┌─┐                                             \n",
      " q_0: ┤ U3(2.8492,0.1,0) ├┤M├─────────────────────────────────────────────\n",
      "      ├──────────────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ U3(2.6832,0.2,0) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├──────────────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ U3(2.6832,0.3,0) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ U3(2.6832,0.2,0) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ U3(2.4172,0.3,0) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ U3(2.6832,0.3,0) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ U3(2.4172,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ U3(2.8492,0.7,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └──────────────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/═════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                           0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Counts: {'1111111110001111': 1, '1110111111110111': 1, '1011111111110011': 1, '1111110111100111': 1, '1111111010011111': 1, '1111110010111011': 1, '1011101110111111': 1, '1111001101111111': 1, '1101101111111111': 1, '1101111111101111': 1, '1111111110110011': 1, '1011111010111111': 1, '1111110101111011': 1, '1111100111011111': 1, '1011101111111011': 2, '1010111111011111': 1, '1111101001101111': 1, '1111111000111111': 1, '1111110100111111': 1, '1011101111111111': 2, '1101011111011111': 1, '1111101011111111': 1, '1101111110111101': 1, '1111101111111110': 1, '1011111111110111': 1, '1111110010111111': 1, '0111111110111111': 1, '1110111100111111': 1, '1101110110111111': 1, '1111110111011101': 1, '1011101111011111': 2, '1111011101111111': 1, '1111101101111111': 1, '1111110111011011': 1, '1111101101111101': 1, '1101111111111110': 2, '1111110111101111': 3, '1011100100111111': 1, '1111001111011111': 1, '1111101101111011': 1, '1111101111110111': 2, '1011110110111111': 1, '1110111111111111': 12, '0110111111111111': 1, '1111111111011110': 1, '1011111111111101': 1, '1111001110111111': 1, '0111101110111111': 1, '1111111011111011': 1, '1111111110011111': 6, '1111111110010111': 1, '1111111101101101': 1, '1111011111111011': 1, '1111110111111110': 1, '1111111111011101': 5, '1111111101111011': 1, '1011110111111111': 5, '1101100111111111': 1, '1101111001111111': 1, '1111110110111011': 1, '1111100110111011': 1, '1101011111111111': 2, '1111111101111111': 8, '1011110111011111': 1, '1110110111011111': 1, '1111111010111111': 3, '1110101111011111': 1, '1111011111101111': 2, '1011111110111111': 1, '1011111111011111': 2, '1011011111111101': 1, '1111111101110111': 1, '1110111110111111': 1, '1111110110011111': 4, '0111111110011111': 1, '0111101111111111': 2, '1111101111111101': 4, '1110101110111111': 1, '1101111111111111': 20, '1110101111111111': 2, '1011111111010111': 1, '1111010111111111': 6, '1011111011111111': 1, '1111101110111111': 7, '1111111111001111': 2, '1111110111111111': 55, '1101011101111111': 1, '0111111011111101': 1, '1111011111111110': 1, '1111111011011111': 3, '1110110100111101': 1, '0111111111111101': 1, '0011111111111111': 2, '0111101111110111': 1, '1111111011111111': 25, '1111011111011111': 5, '1011001110111101': 1, '1101101110111111': 1, '1111111110110111': 2, '1011011111011111': 1, '1111110101111111': 3, '1111101111101011': 1, '1111110111110111': 2, '1111101010111111': 2, '1111111110011101': 1, '0111111111011111': 1, '1111101111101111': 1, '1111111111100111': 2, '1101110111111111': 2, '1110111111011111': 1, '1111111111111110': 9, '1111110111111101': 2, '1111011110111111': 6, '1111100110111111': 1, '1111111111011111': 50, '1111100011111111': 2, '1110110111101111': 1, '1111101101011111': 3, '1111101111111011': 1, '1111111110111101': 4, '0101111110101111': 1, '0111111111111111': 7, '1111101111111111': 52, '1111000111111111': 2, '0111101010111111': 1, '1111111110111111': 54, '1111010111011111': 1, '1111011110111101': 1, '1111111111011011': 2, '1111010101111111': 1, '1110101110111011': 1, '1011111101011111': 1, '1111111110111011': 2, '1111011111111111': 11, '1111111111101111': 16, '1111111111111000': 1, '1111111100101111': 1, '1111110111011111': 9, '0111011100111111': 1, '1101101111011111': 1, '1111111101011111': 1, '1101111111011111': 5, '1101111111111011': 2, '1111111111111101': 21, '1101111110111111': 6, '1111110011111111': 4, '1111111001111111': 1, '1011111101111111': 1, '1111111100111111': 3, '1111111101101111': 2, '0111111111111011': 2, '1111001100011111': 1, '1111111110101111': 2, '0111011011111111': 1, '1111111111110111': 5, '1111110111111011': 4, '1111110110111111': 11, '1011111111111111': 20, '1111101100101111': 1, '1111101111011111': 10, '1111100111111111': 6, '1111111101111101': 2, '1111111111111011': 30, '1111111111111111': 354}\n",
      "Expected number of qubits: 16\n",
      "Actual bitstring length: 16\n",
      "Bitstring: 1111111110001111, Index: 65423, Row: 16355, Col: 3, Count: 1\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "index 16355 is out of bounds for axis 0 with size 4",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[2], line 87\u001b[0m\n",
      "\u001b[1;32m     83\u001b[0m         result_array[row, col] \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\u001b[0;32m---> 87\u001b[0m result_array \u001b[38;5;241m=\u001b[39m \u001b[43mcounts_to_2d_array\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcounts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;32m     89\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mResult array:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, result_array)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[1;32m     91\u001b[0m \u001b[38;5;66;03m# Plot the initial wave function and the result\u001b[39;00m\n",
      "\n",
      "Cell \u001b[0;32mIn[2], line 83\u001b[0m, in \u001b[0;36mcounts_to_2d_array\u001b[0;34m(counts, grid_size)\u001b[0m\n",
      "\u001b[1;32m     81\u001b[0m     col \u001b[38;5;241m=\u001b[39m index \u001b[38;5;241m%\u001b[39m grid_size[\u001b[38;5;241m1\u001b[39m]\n",
      "\u001b[1;32m     82\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBitstring: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtruncated_bitstring\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Index: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mindex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Row: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrow\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Col: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcol\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Count: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcount\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[0;32m---> 83\u001b[0m     \u001b[43mresult_array\u001b[49m\u001b[43m[\u001b[49m\u001b[43mrow\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcol\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\n",
      "\u001b[0;31mIndexError\u001b[0m: index 16355 is out of bounds for axis 0 with size 4"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, assemble\n",
    "from qiskit.visualization import plot_histogram\n",
    "from qiskit_aer import AerSimulator\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def initialize_wavefunction_2d(qc, qubits, wave_function):\n",
    "    \"\"\"Initialize the quantum circuit with the given 2D wave function.\"\"\"\n",
    "    for i, amplitude in enumerate(wave_function.flatten()):\n",
    "        qc.ry(2 * np.arccos(amplitude), qubits[i])\n",
    "\n",
    "def apply_phase_shift_2d(qc, qubits, phase_shifts):\n",
    "    \"\"\"Apply phase shift to simulate the weak phase object in 2D.\"\"\"\n",
    "    for i, phase in enumerate(phase_shifts.flatten()):\n",
    "        qc.rz(phase, qubits[i])\n",
    "\n",
    "def simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4)):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    \n",
    "    # Initialize quantum circuit\n",
    "    qr = QuantumRegister(n_qubits, 'q')\n",
    "    cr = ClassicalRegister(n_qubits, 'c')\n",
    "    qc = QuantumCircuit(qr, cr)\n",
    "    \n",
    "    # Define initial 2D wave function (Gaussian-like distribution)\n",
    "    x = np.linspace(-1, 1, grid_size[0])\n",
    "    y = np.linspace(-1, 1, grid_size[1])\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    initial_wave_function = np.exp(-0.5 * (X**2 + Y**2))\n",
    "    initial_wave_function /= np.linalg.norm(initial_wave_function)\n",
    "    \n",
    "    print(\"Initial 2D wave function:\\n\", initial_wave_function)  # Debug print\n",
    "    \n",
    "    # Initialize wave function in the quantum circuit\n",
    "    initialize_wavefunction_2d(qc, qr, initial_wave_function)\n",
    "    \n",
    "    # Apply phase shift to simulate the weak phase object\n",
    "    apply_phase_shift_2d(qc, qr, phase_shifts)\n",
    "    \n",
    "    # Measure the results\n",
    "    qc.measure(qr, cr)\n",
    "    \n",
    "    print(\"Quantum circuit:\\n\", qc)  # Debug print\n",
    "    \n",
    "    return qc, initial_wave_function\n",
    "\n",
    "# Define phase shifts to simulate the weak phase object in 2D\n",
    "phase_shifts = np.array([[0.1, 0.2, 0.3, 0.4],\n",
    "                         [0.2, 0.3, 0.4, 0.5],\n",
    "                         [0.3, 0.4, 0.5, 0.6],\n",
    "                         [0.4, 0.5, 0.6, 0.7]])\n",
    "\n",
    "# Simulate electron wave in 2D\n",
    "qc, initial_wave_function = simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4))\n",
    "\n",
    "# Execute the circuit using AerSimulator\n",
    "simulator = AerSimulator()\n",
    "compiled_circuit = transpile(qc, simulator)\n",
    "print(\"Compiled circuit:\\n\", compiled_circuit)  # Debug print\n",
    "sim_result = simulator.run(compiled_circuit).result()\n",
    "counts = sim_result.get_counts()\n",
    "\n",
    "print(\"Counts:\", counts)  # Debug print\n",
    "\n",
    "# Verify the number of qubits\n",
    "n_qubits = 4 * 4\n",
    "bitstring_length = len(list(counts.keys())[0])\n",
    "print(f\"Expected number of qubits: {n_qubits}\")\n",
    "print(f\"Actual bitstring length: {bitstring_length}\")\n",
    "\n",
    "# Convert counts to 2D array for visualization\n",
    "def counts_to_2d_array(counts, grid_size):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    result_array = np.zeros((grid_size[0], grid_size[1]))\n",
    "    \n",
    "    for bitstring, count in counts.items():\n",
    "        # Truncate the bitstring to the expected length\n",
    "        truncated_bitstring = bitstring[-n_qubits:]\n",
    "        index = int(truncated_bitstring, 2)\n",
    "        row = index // grid_size[1]\n",
    "        col = index % grid_size[1]\n",
    "        print(f\"Bitstring: {truncated_bitstring}, Index: {index}, Row: {row}, Col: {col}, Count: {count}\")  # Debug print\n",
    "        result_array[row, col] = count\n",
    "    \n",
    "    return result_array\n",
    "\n",
    "result_array = counts_to_2d_array(counts, (4, 4))\n",
    "\n",
    "print(\"Result array:\\n\", result_array)  # Debug print\n",
    "\n",
    "# Plot the initial wave function and the result\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "ax[0].imshow(initial_wave_function, cmap='viridis', origin='lower')\n",
    "ax[0].set_title(\"Initial 2D Wave Function\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "ax[1].imshow(result_array, cmap='viridis', origin='lower')\n",
    "ax[1].set_title(\"Electron Wave Simulation Results\")\n",
    "ax[1].axis('off')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial 2D wave function:\n",
      " [[0.14566959 0.22718971 0.22718971 0.14566959]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.22718971 0.35433041 0.35433041 0.22718971]\n",
      " [0.14566959 0.22718971 0.22718971 0.14566959]]\n",
      "Quantum circuit:\n",
      "       ┌────────────┐┌─────────┐┌─┐                                             \n",
      " q_0: ┤ Ry(2.8492) ├┤ Rz(0.1) ├┤M├─────────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ Ry(2.6832) ├┤ Rz(0.2) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ Ry(2.4172) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ Ry(2.6832) ├┤ Rz(0.3) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ Ry(2.4172) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ Ry(2.4172) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ Ry(2.8492) ├┤ Rz(0.4) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ Ry(2.6832) ├┤ Rz(0.5) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ Ry(2.6832) ├┤ Rz(0.6) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├────────────┤├─────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ Ry(2.8492) ├┤ Rz(0.7) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └────────────┘└─────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/══════════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                                0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Compiled circuit:\n",
      " global phase: 3.0832\n",
      "      ┌──────────────────┐┌─┐                                             \n",
      " q_0: ┤ U3(2.8492,0.1,0) ├┤M├─────────────────────────────────────────────\n",
      "      ├──────────────────┤└╥┘┌─┐                                          \n",
      " q_1: ┤ U3(2.6832,0.2,0) ├─╫─┤M├──────────────────────────────────────────\n",
      "      ├──────────────────┤ ║ └╥┘┌─┐                                       \n",
      " q_2: ┤ U3(2.6832,0.3,0) ├─╫──╫─┤M├───────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║ └╥┘┌─┐                                    \n",
      " q_3: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫─┤M├────────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║ └╥┘┌─┐                                 \n",
      " q_4: ┤ U3(2.6832,0.2,0) ├─╫──╫──╫──╫─┤M├─────────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║ └╥┘┌─┐                              \n",
      " q_5: ┤ U3(2.4172,0.3,0) ├─╫──╫──╫──╫──╫─┤M├──────────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║ └╥┘┌─┐                           \n",
      " q_6: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫─┤M├───────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║ └╥┘┌─┐                        \n",
      " q_7: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫─┤M├────────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                     \n",
      " q_8: ┤ U3(2.6832,0.3,0) ├─╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐                  \n",
      " q_9: ┤ U3(2.4172,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐               \n",
      "q_10: ┤ U3(2.4172,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐            \n",
      "q_11: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├────────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐         \n",
      "q_12: ┤ U3(2.8492,0.4,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├─────────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐      \n",
      "q_13: ┤ U3(2.6832,0.5,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├──────\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐   \n",
      "q_14: ┤ U3(2.6832,0.6,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├───\n",
      "      ├──────────────────┤ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘┌─┐\n",
      "q_15: ┤ U3(2.8492,0.7,0) ├─╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫──╫─┤M├\n",
      "      └──────────────────┘ ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║  ║ └╥┘\n",
      "c: 16/═════════════════════╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩══╩═\n",
      "                           0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15\n",
      "Counts: {'1111111110001111': 1, '1110111111110111': 1, '1011111111110011': 1, '1111110111100111': 1, '1111111010011111': 1, '1111110010111011': 1, '1011101110111111': 1, '1111001101111111': 1, '1101101111111111': 1, '1101111111101111': 1, '1111111110110011': 1, '1011111010111111': 1, '1111110101111011': 1, '1111100111011111': 1, '1011101111111011': 2, '1010111111011111': 1, '1111101001101111': 1, '1111111000111111': 1, '1111110100111111': 1, '1011101111111111': 2, '1101011111011111': 1, '1111101011111111': 1, '1101111110111101': 1, '1111101111111110': 1, '1011111111110111': 1, '1111110010111111': 1, '0111111110111111': 1, '1110111100111111': 1, '1101110110111111': 1, '1111110111011101': 1, '1011101111011111': 2, '1111011101111111': 1, '1111101101111111': 1, '1111110111011011': 1, '1111101101111101': 1, '1101111111111110': 2, '1111110111101111': 3, '1011100100111111': 1, '1111001111011111': 1, '1111101101111011': 1, '1111101111110111': 2, '1011110110111111': 1, '1110111111111111': 12, '0110111111111111': 1, '1111111111011110': 1, '1011111111111101': 1, '1111001110111111': 1, '0111101110111111': 1, '1111111011111011': 1, '1111111110011111': 6, '1111111110010111': 1, '1111111101101101': 1, '1111011111111011': 1, '1111110111111110': 1, '1111111111011101': 5, '1111111101111011': 1, '1011110111111111': 5, '1101100111111111': 1, '1101111001111111': 1, '1111110110111011': 1, '1111100110111011': 1, '1101011111111111': 2, '1111111101111111': 8, '1011110111011111': 1, '1110110111011111': 1, '1111111010111111': 3, '1110101111011111': 1, '1111011111101111': 2, '1011111110111111': 1, '1011111111011111': 2, '1011011111111101': 1, '1111111101110111': 1, '1110111110111111': 1, '1111110110011111': 4, '0111111110011111': 1, '0111101111111111': 2, '1111101111111101': 4, '1110101110111111': 1, '1101111111111111': 20, '1110101111111111': 2, '1011111111010111': 1, '1111010111111111': 6, '1011111011111111': 1, '1111101110111111': 7, '1111111111001111': 2, '1111110111111111': 55, '1101011101111111': 1, '0111111011111101': 1, '1111011111111110': 1, '1111111011011111': 3, '1110110100111101': 1, '0111111111111101': 1, '0011111111111111': 2, '0111101111110111': 1, '1111111011111111': 25, '1111011111011111': 5, '1011001110111101': 1, '1101101110111111': 1, '1111111110110111': 2, '1011011111011111': 1, '1111110101111111': 3, '1111101111101011': 1, '1111110111110111': 2, '1111101010111111': 2, '1111111110011101': 1, '0111111111011111': 1, '1111101111101111': 1, '1111111111100111': 2, '1101110111111111': 2, '1110111111011111': 1, '1111111111111110': 9, '1111110111111101': 2, '1111011110111111': 6, '1111100110111111': 1, '1111111111011111': 50, '1111100011111111': 2, '1110110111101111': 1, '1111101101011111': 3, '1111101111111011': 1, '1111111110111101': 4, '0101111110101111': 1, '0111111111111111': 7, '1111101111111111': 52, '1111000111111111': 2, '0111101010111111': 1, '1111111110111111': 54, '1111010111011111': 1, '1111011110111101': 1, '1111111111011011': 2, '1111010101111111': 1, '1110101110111011': 1, '1011111101011111': 1, '1111111110111011': 2, '1111011111111111': 11, '1111111111101111': 16, '1111111111111000': 1, '1111111100101111': 1, '1111110111011111': 9, '0111011100111111': 1, '1101101111011111': 1, '1111111101011111': 1, '1101111111011111': 5, '1101111111111011': 2, '1111111111111101': 21, '1101111110111111': 6, '1111110011111111': 4, '1111111001111111': 1, '1011111101111111': 1, '1111111100111111': 3, '1111111101101111': 2, '0111111111111011': 2, '1111001100011111': 1, '1111111110101111': 2, '0111011011111111': 1, '1111111111110111': 5, '1111110111111011': 4, '1111110110111111': 11, '1011111111111111': 20, '1111101100101111': 1, '1111101111011111': 10, '1111100111111111': 6, '1111111101111101': 2, '1111111111111011': 30, '1111111111111111': 354}\n",
      "Expected number of qubits: 16\n",
      "Actual bitstring length: 16\n",
      "Bitstring: 1111111110001111, Index: 65423, Row: 16355, Col: 3, Count: 1\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "index 16355 is out of bounds for axis 0 with size 4",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)\n",
      "Cell \u001b[0;32mIn[2], line 87\u001b[0m\n",
      "\u001b[1;32m     83\u001b[0m         result_array[row, col] \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\u001b[0;32m---> 87\u001b[0m result_array \u001b[38;5;241m=\u001b[39m \u001b[43mcounts_to_2d_array\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcounts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;32m     89\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mResult array:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, result_array)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[1;32m     91\u001b[0m \u001b[38;5;66;03m# Plot the initial wave function and the result\u001b[39;00m\n",
      "\n",
      "Cell \u001b[0;32mIn[2], line 83\u001b[0m, in \u001b[0;36mcounts_to_2d_array\u001b[0;34m(counts, grid_size)\u001b[0m\n",
      "\u001b[1;32m     81\u001b[0m     col \u001b[38;5;241m=\u001b[39m index \u001b[38;5;241m%\u001b[39m grid_size[\u001b[38;5;241m1\u001b[39m]\n",
      "\u001b[1;32m     82\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBitstring: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtruncated_bitstring\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Index: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mindex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Row: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrow\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Col: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcol\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Count: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcount\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)  \u001b[38;5;66;03m# Debug print\u001b[39;00m\n",
      "\u001b[0;32m---> 83\u001b[0m     \u001b[43mresult_array\u001b[49m\u001b[43m[\u001b[49m\u001b[43mrow\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcol\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m count\n",
      "\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result_array\n",
      "\n",
      "\u001b[0;31mIndexError\u001b[0m: index 16355 is out of bounds for axis 0 with size 4"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile, assemble\n",
    "from qiskit.visualization import plot_histogram\n",
    "from qiskit_aer import AerSimulator\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def initialize_wavefunction_2d(qc, qubits, wave_function):\n",
    "    \"\"\"Initialize the quantum circuit with the given 2D wave function.\"\"\"\n",
    "    for i, amplitude in enumerate(wave_function.flatten()):\n",
    "        qc.ry(2 * np.arccos(amplitude), qubits[i])\n",
    "\n",
    "def apply_phase_shift_2d(qc, qubits, phase_shifts):\n",
    "    \"\"\"Apply phase shift to simulate the weak phase object in 2D.\"\"\"\n",
    "    for i, phase in enumerate(phase_shifts.flatten()):\n",
    "        qc.rz(phase, qubits[i])\n",
    "\n",
    "def simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4)):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    \n",
    "    # Initialize quantum circuit\n",
    "    qr = QuantumRegister(n_qubits, 'q')\n",
    "    cr = ClassicalRegister(n_qubits, 'c')\n",
    "    qc = QuantumCircuit(qr, cr)\n",
    "    \n",
    "    # Define initial 2D wave function (Gaussian-like distribution)\n",
    "    x = np.linspace(-1, 1, grid_size[0])\n",
    "    y = np.linspace(-1, 1, grid_size[1])\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    initial_wave_function = np.exp(-0.5 * (X**2 + Y**2))\n",
    "    initial_wave_function /= np.linalg.norm(initial_wave_function)\n",
    "    \n",
    "    print(\"Initial 2D wave function:\\n\", initial_wave_function)  # Debug print\n",
    "    \n",
    "    # Initialize wave function in the quantum circuit\n",
    "    initialize_wavefunction_2d(qc, qr, initial_wave_function)\n",
    "    \n",
    "    # Apply phase shift to simulate the weak phase object\n",
    "    apply_phase_shift_2d(qc, qr, phase_shifts)\n",
    "    \n",
    "    # Measure the results\n",
    "    qc.measure(qr, cr)\n",
    "    \n",
    "    print(\"Quantum circuit:\\n\", qc)  # Debug print\n",
    "    \n",
    "    return qc, initial_wave_function\n",
    "\n",
    "# Define phase shifts to simulate the weak phase object in 2D\n",
    "phase_shifts = np.array([[0.1, 0.2, 0.3, 0.4],\n",
    "                         [0.2, 0.3, 0.4, 0.5],\n",
    "                         [0.3, 0.4, 0.5, 0.6],\n",
    "                         [0.4, 0.5, 0.6, 0.7]])\n",
    "\n",
    "# Simulate electron wave in 2D\n",
    "qc, initial_wave_function = simulate_electron_wave_2d(phase_shifts, grid_size=(4, 4))\n",
    "\n",
    "# Execute the circuit using AerSimulator\n",
    "simulator = AerSimulator()\n",
    "compiled_circuit = transpile(qc, simulator)\n",
    "print(\"Compiled circuit:\\n\", compiled_circuit)  # Debug print\n",
    "sim_result = simulator.run(compiled_circuit).result()\n",
    "counts = sim_result.get_counts()\n",
    "\n",
    "print(\"Counts:\", counts)  # Debug print\n",
    "\n",
    "# Verify the number of qubits\n",
    "n_qubits = 4 * 4\n",
    "bitstring_length = len(list(counts.keys())[0])\n",
    "print(f\"Expected number of qubits: {n_qubits}\")\n",
    "print(f\"Actual bitstring length: {bitstring_length}\")\n",
    "\n",
    "# Convert counts to 2D array for visualization\n",
    "def counts_to_2d_array(counts, grid_size):\n",
    "    n_qubits = grid_size[0] * grid_size[1]\n",
    "    result_array = np.zeros((grid_size[0], grid_size[1]))\n",
    "    \n",
    "    for bitstring, count in counts.items():\n",
    "        # Truncate the bitstring to the expected length\n",
    "        truncated_bitstring = bitstring[-n_qubits:]\n",
    "        index = int(truncated_bitstring, 2)\n",
    "        row = index // grid_size[1]\n",
    "        col = index % grid_size[1]\n",
    "        print(f\"Bitstring: {truncated_bitstring}, Index: {index}, Row: {row}, Col: {col}, Count: {count}\")  # Debug print\n",
    "        result_array[row, col] = count\n",
    "    \n",
    "    return result_array\n",
    "\n",
    "result_array = counts_to_2d_array(counts, (4, 4))\n",
    "\n",
    "print(\"Result array:\\n\", result_array)  # Debug print\n",
    "\n",
    "# Plot the initial wave function and the result\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "ax[0].imshow(initial_wave_function, cmap='viridis', origin='lower')\n",
    "ax[0].set_title(\"Initial 2D Wave Function\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "ax[1].imshow(result_array, cmap='viridis', origin='lower')\n",
    "ax[1].set_title(\"Electron Wave Simulation Results\")\n",
    "ax[1].axis('off')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cwq",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}