{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test time complexity *vs.* linear chain size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from fuNEGF.models import LinearChain\n",
    "import timeit\n",
    "from scipy.optimize import curve_fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define static parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# static parameters\n",
    "eps_0 = 0\n",
    "t = 1\n",
    "a = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time complexity of model construction and $T(E)$ calculation for a clean system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47040 characters]",
      "text/plain": [
       "<Figure size 700x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "N_all = np.arange(1, 20, 2)\n",
    "times_constr_only_all = []\n",
    "times_constr_and_transmission_all = []\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))\n",
    "\n",
    "for N in N_all:\n",
    "\n",
    "    def time_constr():\n",
    "        chain = LinearChain(N, eps_0, t, a, plot_dispersion=False)\n",
    "\n",
    "    chain = LinearChain(N, eps_0, t, a, plot_dispersion=False)\n",
    "\n",
    "    def time_transmission_calculation():\n",
    "        chain.plot_transmission(ax=ax)\n",
    "\n",
    "    t_constr_only = timeit.timeit(lambda: time_constr(), number=100)\n",
    "    t_constr_and_transmission = timeit.timeit(\n",
    "        lambda: time_transmission_calculation(), number=30\n",
    "    )\n",
    "\n",
    "    times_constr_only_all.append(t_constr_only)\n",
    "    times_constr_and_transmission_all.append(t_constr_and_transmission)\n",
    "plt.close()\n",
    "\n",
    "\n",
    "def fit_power_law_and_plot(N_all, times_all, ax, title, color=None, marker=None):\n",
    "    # fit data with a power law\n",
    "    def power_law(x, a, b, c):\n",
    "        return a + b * np.power(x, c)\n",
    "\n",
    "    N_all = np.array(N_all)\n",
    "    times_all = np.array(times_all)\n",
    "    p0 = [0, 0.01, 1]\n",
    "    p = curve_fit(power_law, N_all, times_all, p0=p0)[0]\n",
    "\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(1, 1, figsize=(3.5, 3))\n",
    "    if color is None:\n",
    "        color = \"lightblue\"\n",
    "    if marker is None:\n",
    "        marker = \"o\"\n",
    "    ax.plot(N_all, times_all, \"-\", color=color, marker=marker)\n",
    "    ax.plot(\n",
    "        N_all,\n",
    "        power_law(N_all, *p),\n",
    "        \"k--\",\n",
    "        label=f\"$t(N) = {p[0]:.2g} + {p[1]:.2g}\"\n",
    "        + r\" \\cdot N^{\"\n",
    "        + f\"{p[2]:.2f}\"\n",
    "        + r\"} $\",\n",
    "    )\n",
    "    ax.set_xlabel(r\"$N$\", fontsize=11)\n",
    "    ax.set_ylabel(\"Time (s)\", fontsize=11)\n",
    "    ax.legend(\n",
    "        title=r\"scaling $\\mathcal{O}(\" + f\"{p[2]:.3g}\" + r\")$\",\n",
    "        loc=\"upper left\",\n",
    "        fontsize=9,\n",
    "    )\n",
    "    # legend title bold\n",
    "    ax.get_legend().get_title().set_fontweight(\"bold\")\n",
    "    ax.axhline(0, color=\"black\", lw=0.5)\n",
    "    ax.axvline(0, color=\"black\", lw=0.5)\n",
    "    if title is None:\n",
    "        title = (\n",
    "            r\"Time complexity of size-$N$ \" + \"\\n\" + r\"linear chain $T(E)$ calculation\"\n",
    "        )\n",
    "    ax.set_title(title, fontsize=11)\n",
    "    plt.tight_layout()\n",
    "\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(7.0, 3))\n",
    "plt.suptitle(\"Time complexity of size-$N$ linear chain calculation\", fontsize=12)\n",
    "fit_power_law_and_plot(\n",
    "    N_all,\n",
    "    times_constr_only_all,\n",
    "    ax=axes[0],\n",
    "    color=\"C0\",\n",
    "    marker=\"o\",\n",
    "    title=\"Hamiltonian construction\",\n",
    ")\n",
    "fit_power_law_and_plot(\n",
    "    N_all,\n",
    "    times_constr_and_transmission_all,\n",
    "    ax=axes[1],\n",
    "    color=\"C1\",\n",
    "    marker=\"o\",\n",
    "    title=r\"$T(E)$ calculation\",\n",
    ")\n",
    "axes[1].set_ylabel(\"\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Hamiltonian construction $\\mathcal{O}(N)$\n",
    "- $T(E)$ calculation $\\mathcal{O}(N^2)$\n",
    "  - $\\mathcal{O}(N^{\\approx 2})$ correponds to matrix multiplication\n",
    "    - \"As of January 2024, the best bound on the asymptotic complexity of a matrix multiplication algorithm is O(n2.371552).\" (wiki)\n",
    "    - \"numpy is \"incredibly fast\" in matrix multiplication, using a highly optimized BLAS (Basic Linear Algebra Subprograms) implementation\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
