{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "schrodinger_folder = os.path.join('..', 'schrodinger') # path to the schrodinger package downloaded from https://github.com/agolovanov/schrodinger\n",
    "sys.path.append(schrodinger_folder)\n",
    "import decay # decay.py which comes with the supplementary materials\n",
    "import wavefunction\n",
    "import potential\n",
    "import solver\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "import numpy as np\n",
    "import eigen\n",
    "from scipy import optimize, integrate, interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure_width = 3.37\n",
    "figure_height = 0.7 * figure_width\n",
    "plt.rc('figure', figsize=(figure_width, figure_height), dpi=200)\n",
    "plt.rc('font', size=10)\n",
    "plt.rc('text', usetex=True)\n",
    "plt.rc('lines', linewidth=1.4)\n",
    "plt.rc('axes', grid=True, labelpad=1.0, axisbelow=True)\n",
    "plt.rc('grid', linewidth = 0.5, linestyle = '-', color='0.8')\n",
    "plt.rc('xtick', direction='in')\n",
    "plt.rc('ytick', direction='in')\n",
    "plt.rc('xtick.major', size = 0)\n",
    "plt.rc('ytick.major', size = 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D delta potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depth = 1.0\n",
    "fields = [0.3, 1, 5]\n",
    "\n",
    "v = potential.DeltaPotential1D(depth)\n",
    "solvers = [solver.CrankNicolsonSolver(100, 0.02, 0.01, potential=potential.UniformField1D(f, potential=v)) for f in fields]\n",
    "x = solvers[0].x\n",
    "psi0 = v.get_eigenfunction()(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 4.1\n",
    "iters = 50\n",
    "corrs_numeric = []\n",
    "corrs_analytic = []\n",
    "corrs_analytic_bf = []\n",
    "\n",
    "for f, s in zip(fields, solvers):\n",
    "    ts, psis = s.execute(tmax, output_dt=tmax/iters, psi0=psi0)\n",
    "    \n",
    "    corrs_numeric.append(np.array([np.abs(wavefunction.correlation(x, psi0, psi)) ** 2 for psi in psis]))\n",
    "    corrs_analytic.append(decay.calculate_state_decay_analytic(x, psi0, ts, f))\n",
    "    corrs_analytic_bf.append(decay.calculate_state_decay_analytic(x, psi0, ts, f, True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delta_corr(A):\n",
    "    return ((0.5 * A) ** 2 + 1) ** (-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=3, figsize=(figure_width, 1.2 * figure_width))\n",
    "\n",
    "for a, f, c1, c2, c3 in zip(ax, fields, corrs_numeric, corrs_analytic_bf, corrs_analytic):\n",
    "    a.plot(ts, c1, label='TDSE')\n",
    "    a.plot(ts, c2, '--', label='Motionless')\n",
    "    a.plot(ts, c3, ':', label='Free electron')\n",
    "    #a.plot(ts, delta_corr(f * ts))\n",
    "    a.set_ylim(0, 1.05)\n",
    "    a.set_ylabel('$C_0$')\n",
    "    #a.set_title(f'$E={f:g}$')\n",
    "    a.set_yticks([0,0.5, 1])\n",
    "    a.set_xlim(0, 4)\n",
    "    a.text(0.8, 0.85, f'$E={f:g}$', transform=a.transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "    \n",
    "ax[0].set_xticklabels([])\n",
    "ax[1].set_xticklabels([])\n",
    "ax[2].set_xlabel('$t$')\n",
    "\n",
    "ax[2].legend(fontsize=9, loc=4)\n",
    "\n",
    "for a in ax[:2]:\n",
    "    a.set_yticks([0.5, 1])\n",
    "\n",
    "plt.tight_layout(pad=0.4, h_pad=0.0)\n",
    "plt.savefig('fig/1d_delta_corr.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for wi(E)\n",
    "\n",
    "print(1 / (np.sqrt(np.sqrt(np.e) - 1) * 2 * np.sqrt(2)))\n",
    "print(1 / (np.sqrt(np.sqrt(np.e) - 1) * 2))\n",
    "\n",
    "alpha1 = 1 / optimize.brentq(lambda A: delta_corr(A) - np.exp(-1), 0, 10)\n",
    "print(alpha1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_ionization(alpha):\n",
    "    return lambda A: np.exp(- alpha * A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sq_diff = lambda alpha: integrate.quad(lambda A: (delta_corr(A) - model_ionization(alpha)(A)) ** 2, 0, np.infty)[0]\n",
    "\n",
    "alpha2 = optimize.minimize(sq_diff, alpha1).x[0]\n",
    "print(alpha2)\n",
    "\n",
    "abs_diff = lambda alpha: integrate.quad(lambda A: np.abs(delta_corr(A) - model_ionization(alpha)(A)), 0, np.infty)[0]\n",
    "\n",
    "alpha3 = optimize.minimize(abs_diff, alpha1).x[0]\n",
    "print(alpha3)\n",
    "\n",
    "A = np.linspace(0, 20, 400)\n",
    "min_diff = lambda alpha: np.max(np.abs(delta_corr(A) - model_ionization(alpha)(A)))\n",
    "\n",
    "alpha4 = optimize.minimize(min_diff, alpha1).x[0]\n",
    "print(alpha4)\n",
    "\n",
    "alpha = np.linspace(0.3, 2, 100)\n",
    "plt.plot(alpha, np.vectorize(sq_diff)(alpha))\n",
    "plt.plot(alpha, np.vectorize(abs_diff)(alpha), '--')\n",
    "plt.plot(alpha, np.vectorize(min_diff)(alpha), ':')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Soft-core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [0.05, 0.1, 0.2]\n",
    "\n",
    "v0 = lambda x: -1 / np.sqrt(x ** 2 + 2)\n",
    "solvers = {f : solver.SplitOperatorHalfSpectralSolver(100, 0.4, 0.1, potential=(lambda x: -1 / np.sqrt(x ** 2 + 2) - f * x), stationary=True) for f in fields}\n",
    "x = solvers[fields[0]].x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w, psis0 = eigen.calculate_eigenstates(x, v0, 10, -0.5)\n",
    "for i in range(len(psis0)):\n",
    "    psi = psis0[i]\n",
    "    if np.max(np.real(psi)) < np.max(-np.real(psi)):\n",
    "        psis0[i] = -psi\n",
    "psi0 = psis0[0]\n",
    "\n",
    "for psi in psis0:\n",
    "    plt.plot(x, np.real(psi))\n",
    "    \n",
    "plt.xlim(-30, 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p, psi_p = wavefunction.momentum_representation(x, psi0)\n",
    "p0_sq = wavefunction.correlation(p, np.conj(psi_p), p ** 2 * psi_p)\n",
    "print(p0_sq)\n",
    "p0 = np.sqrt(np.abs(p0_sq))\n",
    "print(f'p0 = {p0}')\n",
    "\n",
    "print(f'p0^3 = {p0 ** 3})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 4\n",
    "styles = [[1, 0], [3, 1], [1, 1], [3, 1, 1, 1]]\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for i in range(n):\n",
    "    ax.plot(x, np.real(psis0[i]), dashes=styles[i], label=f'$n = {i+1}$')\n",
    "    \n",
    "ax.set_xlim(-20, 20)\n",
    "ax.legend(fontsize='small')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel(r'$\\psi_n$')\n",
    "ax.set_ylim(-0.6, 0.6)\n",
    "\n",
    "plt.tight_layout(pad=0.2)\n",
    "plt.savefig('fig/softcore_wavefunctions.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_core_corr(A):\n",
    "    return np.sum([decay.correlation_p(x, psi, psi0, A) for psi in psis0], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 30\n",
    "iters = 50\n",
    "\n",
    "corrs_numeric = {}\n",
    "corrs_analytic = {}\n",
    "corrs_analytic_bf = {}\n",
    "corrs_numeric_tot = {}\n",
    "corrs_analytic_tot = {}\n",
    "corrs_analytic_bf_tot = {}\n",
    "\n",
    "for f, s in solvers.items():\n",
    "    ts, psis = s.execute(tmax, output_dt=tmax/iters, psi0=psi0)\n",
    "    corrs_numeric[f] = [np.array([np.abs(wavefunction.correlation(x, psi0, psi)) ** 2 for psi in psis]) for psi0 in psis0]\n",
    "    corrs_numeric_tot[f] = np.sum(corrs_numeric[f], axis=0)\n",
    "    corrs_analytic[f] = [decay.calculate_state_decay_analytic(x, psi0, ts, f, False, psi_comp=psi) for psi in psis0]\n",
    "    corrs_analytic_tot[f] = np.sum(corrs_analytic[f], axis=0)\n",
    "    corrs_analytic_bf[f] = [decay.calculate_state_decay_analytic(x, psi0, ts, f, True, psi_comp=psi) for psi in psis0]\n",
    "    corrs_analytic_bf_tot[f] = np.sum(corrs_analytic_bf[f], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotate_axes(ax):\n",
    "    letters = 'abcde'\n",
    "    for i, a in enumerate(ax):\n",
    "        a.text(0.95, 0.88, f'({letters[i]})', ha='center', va='center', fontsize=10, transform=a.transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "field0 = fields[2]\n",
    "print(field0)\n",
    "\n",
    "fig, ax = plt.subplots(nrows = 3, figsize=(figure_width, 1.2 * figure_width))\n",
    "\n",
    "ns = 3\n",
    "styles = [[3, 1], [1, 1], [3, 1, 1, 1]]\n",
    "\n",
    "ax[0].plot(ts, corrs_numeric_tot[field0])\n",
    "for corr, s in zip(corrs_numeric[field0][:ns], styles):\n",
    "    ax[0].plot(ts, corr, dashes=s)\n",
    "\n",
    "ax[1].plot(ts, corrs_analytic_tot[field0], label='$C$')\n",
    "for i, (corr, s) in enumerate(zip(corrs_analytic[field0][:ns], styles)):\n",
    "    ax[1].plot(ts, corr, dashes=s, label=r'$C_{%d}$' % i)\n",
    "    \n",
    "ax[1].legend(loc=4, fontsize='small')\n",
    "\n",
    "ax[2].plot(ts, corrs_analytic_bf_tot[field0])\n",
    "for corr, s in zip(corrs_analytic_bf[field0][:ns], styles):\n",
    "    ax[2].plot(ts, corr, dashes=s)\n",
    "\n",
    "for a in ax:\n",
    "    a.set_xlim(0, 12)\n",
    "    a.set_ylim(0, 1.05)\n",
    "    a.set_ylabel('$C$')\n",
    "    \n",
    "for a in ax[:2]:\n",
    "    a.set_yticks([0.5, 1])\n",
    "    \n",
    "ax[0].set_xticklabels([])\n",
    "ax[1].set_xticklabels([])\n",
    "\n",
    "ax[2].set_xlabel('$t$')\n",
    "#annotate_axes(ax)\n",
    "\n",
    "ax[0].text(0.98, 0.8, 'TDSE', ha='right', transform=ax[0].transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "ax[1].text(0.98, 0.8, 'Free electron', ha='right', transform=ax[1].transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "ax[2].text(0.98, 0.8, 'Motionles', ha='right', transform=ax[2].transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "\n",
    "plt.tight_layout(pad=0.05, h_pad=0)\n",
    "\n",
    "plt.savefig('fig/softcore_multiple.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows = 3, figsize=(figure_width, 1.2 * figure_width))\n",
    "\n",
    "for f, a in zip(fields, ax):\n",
    "    a.plot(ts, corrs_numeric_tot[f], label='TDSE')\n",
    "    a.plot(ts, corrs_analytic_bf_tot[f], '--', label='Motionless')\n",
    "    a.plot(ts, corrs_analytic_tot[f], ':', label='Free electron')\n",
    "    a.set_xlim(0, 20)\n",
    "    a.set_ylim(0, 1.05)\n",
    "    #a.set_title(f'$E={f:g}$')\n",
    "    a.text(0.04, 0.1, f'$E={f:g}$', transform=a.transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "    a.set_ylabel('$C$')\n",
    "    a.set_yticks([0, 0.5, 1])\n",
    "\n",
    "for a in ax[:2]:\n",
    "    a.set_yticks([0.5, 1])\n",
    "\n",
    "ax[2].legend(fontsize='small', loc=1)\n",
    "\n",
    "ax[0].set_xticklabels([])\n",
    "ax[1].set_xticklabels([])\n",
    "\n",
    "ax[2].set_xlabel('$t$')\n",
    "plt.tight_layout(pad=0.05, h_pad=0)\n",
    "\n",
    "plt.savefig('fig/softcore_corr.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.linspace(0, 5, 100)\n",
    "soft_core_corr_value = soft_core_corr(A)\n",
    "soft_core_interp = interpolate.interp1d(A, soft_core_corr_value, fill_value=(1, 0), bounds_error=False)\n",
    "\n",
    "alpha1 = 1 / optimize.brentq(lambda A: soft_core_interp(A) - np.exp(-1), 0, 10)\n",
    "print(alpha1)\n",
    "print(alpha1 / np.sqrt(2))\n",
    "\n",
    "plt.plot(A, soft_core_corr_value)\n",
    "plt.plot(A, model_ionization(alpha1)(A))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sq_diff = lambda alpha: integrate.quad(lambda A: (soft_core_interp(A) - model_ionization(alpha)(A)) ** 2, 0, 10)[0]\n",
    "\n",
    "alpha2 = optimize.minimize(sq_diff, alpha1).x[0]\n",
    "print(alpha2)\n",
    "\n",
    "abs_diff = lambda alpha: integrate.quad(lambda A: np.abs(soft_core_interp(A) - model_ionization(alpha)(A)), 0, 10)[0]\n",
    "\n",
    "alpha3 = optimize.minimize(abs_diff, alpha1).x[0]\n",
    "print(alpha3)\n",
    "\n",
    "A = np.linspace(0, 20, 400)\n",
    "min_diff = lambda alpha: np.max(np.abs(soft_core_interp(A) - model_ionization(alpha)(A)))\n",
    "\n",
    "alpha4 = optimize.minimize(min_diff, alpha1).x[0]\n",
    "print(alpha4)\n",
    "\n",
    "alpha = np.linspace(0.3, 2, 20)\n",
    "plt.plot(alpha, np.vectorize(sq_diff)(alpha))\n",
    "plt.plot(alpha, np.vectorize(abs_diff)(alpha), '--')\n",
    "plt.plot(alpha, np.vectorize(min_diff)(alpha), ':')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_func = lambda x, E: - 1 / np.sqrt(2 + x ** 2) - E * x\n",
    "x = np.linspace(0, 12, 5000)\n",
    "potential_max = lambda E: np.max(potential_func(x, E))\n",
    "\n",
    "Ecritical = optimize.brentq(lambda E: potential_max(E) + 0.5, 0.06, 0.07)\n",
    "\n",
    "print(Ecritical)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_closest(tarr, t0):\n",
    "    i0 = 1\n",
    "    for i, t in enumerate(tarr):\n",
    "        if t > t0:\n",
    "            i0 = i\n",
    "            break\n",
    "    return i0 if abs(tarr[i0] - t0) < abs(tarr[i0-1] - t0) else i0-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "E1 = fields[0]\n",
    "E2 = fields[2]\n",
    "print(E1, E2)\n",
    "\n",
    "ets = [0, 1, 2]\n",
    "\n",
    "fig = plt.figure(figsize=(figure_width, 1.5 * figure_width))\n",
    "\n",
    "grid_x = [ImageGrid(fig, 221, nrows_ncols=(3, 1), axes_pad=0, aspect=False), ImageGrid(fig, 223, nrows_ncols=(3, 1), axes_pad=0, aspect=False)]\n",
    "grid_p = [ImageGrid(fig, 222, nrows_ncols=(3, 1), axes_pad=0, aspect=False), ImageGrid(fig, 224, nrows_ncols=(3, 1), axes_pad=0, aspect=False)]\n",
    "\n",
    "for f, gx, gp in zip([E1, E2], grid_x, grid_p):\n",
    "    tmax = 3 / f\n",
    "    ts, psis = solvers[f].execute(tmax, output_dt=tmax/iters, psi0=psi0)\n",
    "    x = solvers[f].x\n",
    "    for i, et in enumerate(ets):\n",
    "        index = find_closest(ts, et / f)\n",
    "        psi = psis[index] \n",
    "        gx[i].plot(x, np.abs(psi) ** 2, label='TDSE')\n",
    "        psi_motionless = decay.calculate_psi(x, psi0, ts[index], f, big_field=True)\n",
    "        gx[i].plot(x, np.abs(psi_motionless) ** 2, '--', label='M')\n",
    "        psi_free = decay.calculate_psi(x, psi0, ts[index], f, big_field=False)\n",
    "        gx[i].plot(x, np.abs(psi_free) ** 2, ':', label='FE')\n",
    "        \n",
    "        p, psi_p = wavefunction.momentum_representation(x, psi)\n",
    "        gp[i].plot(p, np.abs(psi_p) ** 2, label='TDSE')\n",
    "        p, psi_motionless_p = wavefunction.momentum_representation(x, psi_motionless)\n",
    "        gp[i].plot(p, np.abs(psi_motionless_p) ** 2, '--', label='M')\n",
    "        p, psi_free_p = wavefunction.momentum_representation(x, psi_free)\n",
    "        gp[i].plot(p, np.abs(psi_free_p) ** 2, ':', label='FE')\n",
    "\n",
    "for g in grid_x:\n",
    "    for a in g:\n",
    "        a.set_xlim(-6, 6)\n",
    "        a.set_ylim(0, 0.4)\n",
    "    g[1].set_ylabel(r'$|\\psi|^2$')\n",
    "        \n",
    "\n",
    "for g in grid_p:\n",
    "    for a in g:\n",
    "        a.set_xlim(-3, 3)\n",
    "        a.set_ylim(0, 1.4)\n",
    "    g[1].set_ylabel(r'$|\\tilde\\psi|^2$')\n",
    "        \n",
    "        \n",
    "for g in grid_x + grid_p:\n",
    "    for i, a in enumerate(g):\n",
    "        a.set_yticklabels([])\n",
    "        a.text(0.05, 0.8, f'$t={ets[i]}/E$', transform=a.transAxes, bbox={'facecolor':'white', 'pad':1, 'lw':0})\n",
    "        \n",
    "grid_x[0][-1].set_xticklabels([])\n",
    "grid_p[0][-1].set_xticklabels([])\n",
    "\n",
    "grid_x[1][-1].set_xlabel('$x$')\n",
    "grid_p[1][-1].set_xlabel('$p$')\n",
    "\n",
    "grid_p[0][0].legend(fontsize='x-small', handlelength=1.5, handletextpad=0.3)\n",
    "\n",
    "plt.tight_layout(pad=0.1, h_pad=2, w_pad=0.4, rect=[0.03, 0, 1, 0.96])\n",
    "fig.text(0.54, 0.98, f'$E = {E1}$', ha='center', va='center')\n",
    "fig.text(0.54, 0.51, f'$E = {E2}$', ha='center', va='center')\n",
    "\n",
    "plt.savefig('fig/softcore_evolution.eps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coulomb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-20, 20, 101)\n",
    "r = np.meshgrid(x, x, x, indexing='ij')\n",
    "\n",
    "v = potential.CoulombPotential()\n",
    "\n",
    "levels = [(1,0), (2,0), (2,1), (3,0), (3,1), (3,2), (4,0), (4,1), (4,2), (4,3)]\n",
    "\n",
    "psis = [v.get_eigenfunction(*l, 0)(*r) for l in levels]\n",
    "psi0 = psis[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coulomb_corr(a, psi):\n",
    "    @np.vectorize\n",
    "    def corrf(a):\n",
    "        return wavefunction.correlation(r, psi, psi0 * np.exp(-1j * a * r[2]))\n",
    "    \n",
    "    return np.abs(corrf(a)) ** 2\n",
    "\n",
    "@np.vectorize\n",
    "def coulomb_corr_tot(a):\n",
    "    return np.sum([coulomb_corr(a, psi) for psi in psis], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.linspace(0, 5, 100)\n",
    "\n",
    "corrs = [coulomb_corr(a, psi) for psi in psis]\n",
    "corrs_tot = np.sum(corrs, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "E = 0.4\n",
    "t = a / E\n",
    "plt.plot(a, corrs_tot, label=r'$\\tilde{C}$')\n",
    "\n",
    "styles = [[3, 1], [1, 1], [3, 1, 1, 1]]\n",
    "for l, c, s in zip(levels[:3], corrs[:3], styles):\n",
    "    plt.plot(a, c, dashes=s, label=r'$\\tilde{C}_{%d,%d,0}$' % l)\n",
    "    \n",
    "plt.xlabel('$A$')\n",
    "plt.xlim(0, 4)\n",
    "plt.ylim(0, 1.05)\n",
    "plt.ylabel(r'$\\tilde{C}$')\n",
    "plt.legend(fontsize=10)\n",
    "plt.tight_layout(pad=0.2)\n",
    "\n",
    "plt.savefig('fig/coulomb_corr.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.linspace(0, 5, 100)\n",
    "coulomb_corr_value = coulomb_corr_tot(A)\n",
    "coulomb_core_interp = interpolate.interp1d(A, coulomb_corr_value, fill_value=(1, 0), bounds_error=False)\n",
    "\n",
    "alpha1 = 1 / optimize.brentq(lambda A: coulomb_corr_tot(A) - np.exp(-1), 0, 3)\n",
    "print(alpha1)\n",
    "print(alpha1 / np.sqrt(2))\n",
    "\n",
    "plt.plot(A, coulomb_corr_value)\n",
    "plt.plot(A, model_ionization(alpha1)(A))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sq_diff = lambda alpha: integrate.quad(lambda A: (coulomb_core_interp(A) - model_ionization(alpha)(A)) ** 2, 0, 10)[0]\n",
    "\n",
    "alpha2 = optimize.minimize(sq_diff, alpha1).x[0]\n",
    "print(alpha2)\n",
    "\n",
    "abs_diff = lambda alpha: integrate.quad(lambda A: np.abs(coulomb_core_interp(A) - model_ionization(alpha)(A)), 0, 10)[0]\n",
    "\n",
    "alpha3 = optimize.minimize(abs_diff, alpha1).x[0]\n",
    "print(alpha3)\n",
    "\n",
    "A = np.linspace(0, 20, 400)\n",
    "min_diff = lambda alpha: np.max(np.abs(coulomb_core_interp(A) - model_ionization(alpha)(A)))\n",
    "\n",
    "alpha4 = optimize.minimize(min_diff, alpha1).x[0]\n",
    "print(alpha4)\n",
    "\n",
    "alpha = np.linspace(0.3, 2, 20)\n",
    "plt.plot(alpha, np.vectorize(sq_diff)(alpha))\n",
    "plt.plot(alpha, np.vectorize(abs_diff)(alpha), '--')\n",
    "plt.plot(alpha, np.vectorize(min_diff)(alpha), ':')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tunnel formula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def w(E):\n",
    "    return 4 / np.abs(E) * np.exp(- 2 / 3 / np.abs(E))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "atomic_t = 2.418884326505e-17 # s\n",
    "\n",
    "pulse_length_s = 200e-18\n",
    "\n",
    "pulse_length = pulse_length_s / atomic_t\n",
    "print(pulse_length)\n",
    "\n",
    "amplitude = 10\n",
    "def E_attosecond(t):\n",
    "    return amplitude * np.exp(- (t - 3 * pulse_length) ** 2 / pulse_length ** 2)\n",
    "\n",
    "integral_func_attosecond = np.vectorize(lambda t: integrate.quad(lambda t1: w(E_attosecond(t1)), 0, t)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_l = 0.9 # um\n",
    "atomic_t = 2.418884326505e-17 # s\n",
    "c = 3e8 # m/s\n",
    "omega_l = 2 * np.pi * c / (lambda_l * 1e-6) * atomic_t\n",
    "print(omega_l)\n",
    "T_l = 2 * np.pi / omega_l\n",
    "print(T_l)\n",
    "\n",
    "a0 = 10\n",
    "E_l_SI = 3.21 * a0 * 1e12/ lambda_l\n",
    "print(f\"E = {E_l_SI / 1e12} TV/m\")\n",
    "E_atomic = 5.14220652e11\n",
    "E_l = E_l_SI / E_atomic\n",
    "print(E_l)\n",
    "\n",
    "N_pulse = 60e-15 / atomic_t / T_l\n",
    "phase = 0\n",
    "\n",
    "def E_gauss(t):\n",
    "    return E_l * np.exp(- (t - 1.6 * N_pulse * T_l) ** 2 / (0.5 * N_pulse * T_l) ** 2) * np.cos(omega_l * t + phase)\n",
    "\n",
    "integral_func_gauss = np.vectorize(lambda t: integrate.quad(lambda t1: w(E_gauss(t1)), 0, t)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_attosecond = np.linspace(0, 16, 100)\n",
    "probability_attosecond = 1 - np.exp(- integral_func_attosecond(t_attosecond))\n",
    "\n",
    "t_gauss = np.linspace(0, 1000, 400)\n",
    "probability_gauss = 1 - np.exp(- integral_func_gauss(t_gauss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(figure_width, figure_width))\n",
    "ax[0][1].plot(t_attosecond, E_attosecond(t_attosecond))\n",
    "ax[0][0].plot(t_gauss, E_gauss(t_gauss))\n",
    "\n",
    "ax[1][1].plot(t_attosecond, probability_attosecond)\n",
    "ax[1][0].plot(t_gauss, probability_gauss)\n",
    "\n",
    "ax[0][0].set_ylabel('$E$')\n",
    "ax[1][0].set_ylabel('$W$')\n",
    "\n",
    "ax[0][1].set_ylim(-2, 2)\n",
    "ax[0][0].set_ylim(-0.2, 0.2)\n",
    "\n",
    "ax[0][0].set_xlim(min(t_gauss), max(t_gauss))\n",
    "ax[1][0].set_xlim(min(t_gauss), max(t_gauss))\n",
    "\n",
    "ax[0][1].set_title('200 as pulse')\n",
    "\n",
    "ax[0][0].set_title('60 fs pulse')\n",
    "\n",
    "for a in (ax[0][1], ax[1][1]):\n",
    "    a.set_xlim(min(t_attosecond), max(t_attosecond))\n",
    "    a.fill_between(t_attosecond, 100 * (np.abs(E_attosecond(t_attosecond)) > 1 / 16), -100 * (np.abs(E_attosecond(t_attosecond)) > 1 / 16), color='#ddddf0', zorder=-20, lw=0)\n",
    "    \n",
    "for a in (ax[0][0], ax[1][0]):\n",
    "    a.set_xlim(min(t_gauss), max(t_gauss))\n",
    "    a.fill_between(t_gauss, 100 * (np.abs(E_gauss(t_gauss)) > 1 / 16), -100 * (np.abs(E_gauss(t_gauss)) > 1 / 16), color='#ddddf0', zorder=-20, lw=0)\n",
    "\n",
    "for a in ax[0]:\n",
    "    a.set_xticklabels([])\n",
    "\n",
    "for a in ax[1]:\n",
    "    a.set_xlabel('$t$')\n",
    "    a.set_ylim(0, 1.1)\n",
    "\n",
    "#ax[1][1].set_yticklabels([])\n",
    "\n",
    "plt.tight_layout(pad=0.05, w_pad=0.4, h_pad=0.2)\n",
    "\n",
    "plt.savefig('fig/tunnel_applicability.eps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Different formulas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tunnel(E):\n",
    "    return 4 / E * np.exp(- 2 / 3 / E)\n",
    "\n",
    "def tunnel_ADK(E):\n",
    "    return 2 * np.exp(2) / np.pi / E * np.exp(-2 / 3 / E)\n",
    "\n",
    "@np.vectorize\n",
    "def posthumus(E):\n",
    "    if E < 1 / 16:\n",
    "        return tunnel_ADK(E)\n",
    "    else:\n",
    "        return tunnel_ADK(1 / 16) + (1 - 1 / 16 / E) / 2 / np.pi\n",
    "\n",
    "def field_squared(E):\n",
    "    return 2.4 * E ** 2\n",
    "\n",
    "def field_linear(E):\n",
    "    return 0.8 * E\n",
    "\n",
    "def tunnel_tong(E):\n",
    "    return tunnel(E) * np.exp(-2 * 6 * E)\n",
    "\n",
    "def tunnel_zhang(E):\n",
    "    return tunnel(E) * np.exp((0.11714 * 16 ** 2 * E ** 2 - 0.90933 * 16 * E - 0.06034))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rc('xtick.major', size = 2)\n",
    "plt.rc('ytick.major', size = 2)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "\n",
    "power_min = -1.2\n",
    "power_max = 1\n",
    "\n",
    "E = np.logspace(power_min, power_max, 100)\n",
    "\n",
    "styles = [[3, 1], [1, 1], [3, 1, 1, 1], [3, 1, 1, 1, 1, 1], [2, 1, 1, 1, 1, 1, 1, 1]]\n",
    "\n",
    "plt.plot(E, tunnel(E), label='Tunnel')\n",
    "plt.plot(E, posthumus(E), label='Posthumus', dashes=styles[0])\n",
    "plt.plot(E, tunnel_tong(E), label='Tong', dashes=styles[1])\n",
    "plt.plot(E, tunnel_zhang(E), label='Zhang', dashes=styles[2])\n",
    "plt.plot(E, field_squared(E), label='Bauer', dashes=styles[3])\n",
    "plt.plot(E, field_linear(E), 'C6-o', label='Motionless', ms=2, markevery=5)\n",
    "plt.legend(fontsize=8)\n",
    "\n",
    "plt.xlim(xmin=10 ** power_min, xmax=10 ** power_max)\n",
    "plt.ylim(2e-3, 3)\n",
    "plt.xlabel(r'$E$')\n",
    "plt.ylabel(r'$w_\\mathrm{i}$')\n",
    "\n",
    "plt.tight_layout(pad=0.2)\n",
    "plt.savefig('fig/ionization_models.eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}