{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e7d97093-d746-4749-a754-cd090687a524",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/joaom/mambaforge/envs/qucint/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "from pyqtorch.core.circuit import QuantumCircuit\n",
    "from pyqtorch.core.operation import hamiltonian_evolution, hamiltonian_evolution_eig\n",
    "import copy\n",
    "\n",
    "from pyqtorch.matrices import generate_ising_from_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a069a95-3339-4296-aacf-ea5f8a49059e",
   "metadata": {},
   "source": [
    "# Comparing the single hamiltonian evolution\n",
    "Here we use an Ising hamiltonian over a graph, which is diagonal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7041a91f-4712-459c-8fd6-b147e12faf58",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 8\n",
    "p = 0.5\n",
    "\n",
    "graph = nx.fast_gnp_random_graph(N, p)\n",
    "\n",
    "batch_size = 1\n",
    "qc = QuantumCircuit(N)\n",
    "psi = qc.uniform_state(batch_size)\n",
    "psi_ini = copy.deepcopy(psi)\n",
    "\n",
    "H_diag = generate_ising_from_graph(graph)\n",
    "\n",
    "H = torch.diag(H_diag)\n",
    "\n",
    "qc = QuantumCircuit(N)\n",
    "psi = qc.uniform_state(batch_size)\n",
    "psi_ini = copy.deepcopy(psi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5a447259-e5fa-4154-aa1d-a9507592b021",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_range = torch.arange(0, 2*np.pi, 0.1)\n",
    "wf_save_rk  = torch.zeros((len(t_range),)+tuple(psi.shape)).cdouble()\n",
    "wf_save_eig = torch.zeros((len(t_range),)+tuple(psi.shape)).cdouble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d6103960-5f98-4369-b9b9-b576111f1508",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RK4: 3.183 secs.\n",
      "Eig: 0.038 secs.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "start = time.perf_counter()\n",
    "for i, t in enumerate(t_range):\n",
    "    t_evo = torch.tensor([t])\n",
    "    psi = copy.deepcopy(psi_ini)\n",
    "    hamiltonian_evolution(H, psi, t_evo, range(N), N)\n",
    "    wf_save_rk[i]  = psi\n",
    "end = time.perf_counter()\n",
    "\n",
    "secs = (end-start)\n",
    "print(f\"RK4: {secs:.03f} secs.\")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for i, t in enumerate(t_range):\n",
    "    t_evo = torch.tensor([t])\n",
    "    psi = copy.deepcopy(psi_ini)\n",
    "    hamiltonian_evolution_eig(H, psi, t_evo, range(N), N)\n",
    "    wf_save_eig[i]  = psi\n",
    "end = time.perf_counter()\n",
    "\n",
    "secs = (end-start)\n",
    "print(f\"Eig: {secs:.03f} secs.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7478d25-a006-495e-92fa-16658c8ddba8",
   "metadata": {},
   "source": [
    "The Hamiltonian is diagonal so the eig function skips the diagonalization and is very efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e47eac6f-ed01-43f6-b41c-36ee0a8b99a4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x117d82340>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24460 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "norm_rk = torch.tensor([(abs(wf_save_rk[i,...])**2).sum() for i in range(len(t_range))])\n",
    "norm_eig = torch.tensor([(abs(wf_save_eig[i,...])**2).sum() for i in range(len(t_range))])\n",
    "\n",
    "plt.plot(t_range, norm_rk, label = \"RK4\")\n",
    "plt.plot(t_range, norm_eig, label = \"Eig\")\n",
    "plt.ylim(0.95, 1.01)\n",
    "plt.ylabel(\"Normalization |<psi|psi>|^2\")\n",
    "plt.xlabel(\"Time t\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd9baa5-5584-44bc-a095-7f85ce9bff86",
   "metadata": {},
   "source": [
    "It is also exact, maintaining the normalization of the wavefunction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c66d1b1d-575f-4303-aa86-1bbd4197dcf0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Time t')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 29624 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "diff = torch.tensor([torch.max(abs(wf_save_rk[i,...] - wf_save_eig[i,...])) for i in range(len(t_range))])\n",
    "\n",
    "plt.plot(t_range, diff)\n",
    "plt.yscale(\"log\")\n",
    "plt.ylabel(\"Max(Abs(|psi_rk4> - |psi_eig>))\")\n",
    "plt.xlabel(\"Time t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "312f9b7f-60f7-4ae7-98de-55156a3067e9",
   "metadata": {},
   "source": [
    "As we can see both methods give a similar result for small t but the difference increases exponentially due to the error in the RK4 method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac2bf97-1fa2-4660-acd2-a048bd899d58",
   "metadata": {},
   "source": [
    "# Hamiltonian caching:\n",
    "The results of hamiltonian diagonalization are cached for repeated use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "eadbc7dc-7139-4f54-9651-d888f2867cab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_qubits: int = 10\n",
    "batch_size: int = 1\n",
    "\n",
    "qc = QuantumCircuit(n_qubits)\n",
    "psi = qc.uniform_state(batch_size)\n",
    "psi_ini = copy.deepcopy(psi)\n",
    "\n",
    "H_0 = torch.randn((2**n_qubits, 2**n_qubits), dtype = torch.double)\n",
    "H0 = (H_0 + torch.conj(H_0.transpose(0, 1))).cdouble()\n",
    "\n",
    "H_1 = torch.randn((2**n_qubits, 2**n_qubits), dtype = torch.double)\n",
    "H1 = (H_1 + torch.conj(H_1.transpose(0, 1))).cdouble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c7a135f0-fea0-4af6-80ad-fdf860a7b7af",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Hamiltonian\n",
      "Timing: 0.29912304878234863\n",
      "Timing: 0.07640600204467773\n",
      "Second Hamiltonian\n",
      "Timing: 0.27396464347839355\n",
      "Timing: 0.08124589920043945\n",
      "First Hamiltonian again\n",
      "Timing: 0.08870816230773926\n",
      "Timing: 0.07958817481994629\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "n_trials = 2\n",
    "wf_save_eig = torch.zeros((n_trials,)+tuple(psi.shape)).cdouble()\n",
    "\n",
    "print(\"First Hamiltonian\")\n",
    "for i in range(n_trials):\n",
    "    H = H0\n",
    "    start = time.time()\n",
    "    t_evo = torch.rand(batch_size)*0.5\n",
    "\n",
    "    psi = copy.deepcopy(psi_ini)\n",
    "    hamiltonian_evolution_eig(H, psi, t_evo, range(n_qubits), n_qubits)\n",
    "    wf_save_eig[i] = psi\n",
    "    end = time.time()\n",
    "    print(\"Timing:\", end-start)\n",
    "\n",
    "print(\"Second Hamiltonian\")\n",
    "for i in range(n_trials):\n",
    "    H = H1\n",
    "    start = time.time()\n",
    "    t_evo = torch.rand(batch_size)*0.5\n",
    "\n",
    "    psi = copy.deepcopy(psi_ini)\n",
    "    hamiltonian_evolution_eig(H, psi, t_evo, range(n_qubits), n_qubits)\n",
    "    wf_save_eig[i] = psi\n",
    "    end = time.time()\n",
    "    print(\"Timing:\", end-start)\n",
    "\n",
    "print(\"First Hamiltonian again\")\n",
    "for i in range(n_trials):\n",
    "    H = H0\n",
    "    start = time.time()\n",
    "    t_evo = torch.rand(batch_size)*0.5\n",
    "\n",
    "    psi = copy.deepcopy(psi_ini)\n",
    "    hamiltonian_evolution_eig(H, psi, t_evo, range(n_qubits), n_qubits)\n",
    "    wf_save_eig[i] = psi\n",
    "    end = time.time()\n",
    "    print(\"Timing:\", end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1aa927c-e366-4196-b837-b6c1682f6dab",
   "metadata": {},
   "source": [
    "# Comparing the batched hamiltonian evolution:\n",
    "Here we use general hermitian matrices as hamiltonians"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dc4d878c-2380-4279-aa1e-51a4425d1344",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "from pyqtorch.core.circuit import QuantumCircuit\n",
    "from pyqtorch.core.operation import hamiltonian_evolution, hamiltonian_evolution_eig\n",
    "from pyqtorch.core.batched_operation import batched_hamiltonian_evolution, batched_hamiltonian_evolution_eig\n",
    "import copy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783f5036-60e5-4b6a-87f6-71ca81197f1b",
   "metadata": {},
   "source": [
    "Let's create a batch of hamiltonians that are just repeating the same hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "740a1a03-d72b-4648-93fe-492a3ae60574",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 5\n",
    "batch_size = 10\n",
    "\n",
    "qc = QuantumCircuit(N)\n",
    "\n",
    "psi = qc.uniform_state(batch_size)\n",
    "psi_0 = copy.deepcopy(psi)\n",
    "\n",
    "H_batch = torch.zeros((2**N, 2**N, batch_size), dtype = torch.cdouble)\n",
    "for i in range(batch_size):\n",
    "    H_0 = torch.randn((2**N, 2**N), dtype = torch.cdouble)\n",
    "    H = (H_0 + torch.conj(H_0.transpose(0, 1))).cdouble()\n",
    "    H_batch[...,i] = H"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a4ce18d-fe39-4b83-9615-ae1c7639229f",
   "metadata": {},
   "source": [
    "Now we create a batch of linearly increasing times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bb18b45b-8c18-46c8-8a01-cfc6bd05446e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_size = batch_size\n",
    "t_start = 0.0\n",
    "t_end = 5.0\n",
    "t_evo = torch.arange(t_start, t_end, (t_end-t_start)/t_size).cdouble()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a255ba0-6060-46da-9294-444e9b6ccf1a",
   "metadata": {},
   "source": [
    "Now we evaluate the batched time-evolution, essentially evaluating the same generator for an increasing time. First with the RK4 method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "68c86af7-4539-40db-b169-1f23c7207b5a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.07417893409729004\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "psi = copy.deepcopy(psi_0)\n",
    "start = time.time()\n",
    "psi = batched_hamiltonian_evolution(H_batch, psi, t_evo, range(N), N)\n",
    "end = time.time()\n",
    "psi_save_rk = copy.deepcopy(psi)\n",
    "print(end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701aa46a-cad9-4ad7-b057-a788c27ed53e",
   "metadata": {},
   "source": [
    "And then with the eigenvalue method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "21cea2c7-c9d3-496d-afe0-eec03f7293c0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.005017757415771484\n"
     ]
    }
   ],
   "source": [
    "psi = copy.deepcopy(psi_0)\n",
    "start = time.time()\n",
    "psi = batched_hamiltonian_evolution_eig(H_batch, psi, t_evo, range(N), N)\n",
    "end = time.time()\n",
    "psi_save_eig = copy.deepcopy(psi)\n",
    "print(end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b154c4-7713-4b0e-9174-9e9a2a250e5f",
   "metadata": {},
   "source": [
    "Now we can compare the wavefunction normalization of both as t increases:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "88112d4d-23b6-4e11-b302-fe289cf2e907",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x137aaf7c0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24644 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "t_range = t_evo.real\n",
    "\n",
    "norm_rk = torch.tensor([(abs(psi_save_rk[...,b])**2).sum() for b in range(batch_size)])\n",
    "norm_eig = torch.tensor([(abs(psi_save_eig[...,b])**2).sum() for b in range(batch_size)])\n",
    "\n",
    "plt.plot(t_range, norm_rk, label = \"RK4\")\n",
    "plt.plot(t_range, norm_eig, label = \"Eig\")\n",
    "plt.ylim(0.95, 1.01)\n",
    "plt.ylabel(\"Normalization |<psi|psi>|^2\")\n",
    "plt.xlabel(\"Time t\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af5291ff-8c0c-4fe6-a7c4-5e53c448809b",
   "metadata": {},
   "source": [
    "And also the difference between the wavefunctions as t increases:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "14d90e9e-049c-4754-899d-cadbc4aa3a5b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Time t')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26820 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "diff = torch.tensor([torch.max(abs(psi_save_rk[..., b] - psi_save_eig[..., b])) for b in range(batch_size)])\n",
    "\n",
    "plt.plot(t_range, diff)\n",
    "plt.yscale(\"log\")\n",
    "plt.ylabel(\"Max(Abs(|psi_rk4> - |psi_eig>))\")\n",
    "plt.xlabel(\"Time t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8453624f-a170-4610-b908-1181abe038b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "50935ae23c6784d55197a912635e8caab0c12202c38f5b44c7779b3e5667ccc0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
