Plotting_all.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b913731",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import cmath\n",
    "import math\n",
    "import os\n",
    "import h5py\n",
    "import matplotlib.pyplot as plt   # plots\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.colors as clrs\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import Normalize, TwoSlopeNorm\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "import numpy as np\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "from liblibra_core import *\n",
    "import util.libutil as comn\n",
    "from libra_py import units\n",
    "import libra_py.models.Holstein as Holstein\n",
    "import libra_py.models.Phenol as Phenol\n",
    "\n",
    "from libra_py import dynamics_plotting\n",
    "import libra_py.dynamics.tsh.compute as tsh_dynamics\n",
    "import libra_py.dynamics.tsh.plot as tsh_dynamics_plot\n",
    "import libra_py.dynamics.exact.plot as exact_plot\n",
    "import libra_py.data_savers as data_savers\n",
    "import libra_py.data_read as data_read\n",
    "import libra_py.data_visualize as data_visualize\n",
    "\n",
    "import libra_py.dynamics.exact.compute as dvr\n",
    "import libra_py.dynamics.exact.save as dvr_save\n",
    "\n",
    "import re\n",
    "\n",
    "%matplotlib inline\n",
    "#from matplotlib.mlab import griddata\n",
    "#%matplotlib inline \n",
    "#warnings.filterwarnings('ignore')\n",
    "\n",
    "colors = {}\n",
    "colors.update({\"11\": \"#CB3310\"}) #red\n",
    "colors.update({\"21\": \"#0077BB\"}) #blue\n",
    "colors.update({\"31\": \"#009988\"}) #teal\n",
    "colors.update({\"41\": \"#EE7733\"}) #orange\n",
    "colors.update({\"51\": \"#00FFFF\"}) #cyan\n",
    "colors.update({\"61\": \"#EE3377\"}) #magenta\n",
    "colors.update({\"71\": \"#AA3377\"}) #purple\n",
    "colors.update({\"81\": \"#BBBBBB\" }) #grey\n",
    "\n",
    "clrs_index = [\"11\", \"21\", \"31\", \"41\", \"51\", \"61\", \"71\", \"81\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "477858a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rc('axes', titlesize=30)      # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=30)      # fontsize of the x and y labels\n",
    "plt.rc('legend', fontsize=30)     # legend fontsize\n",
    "plt.rc('xtick', labelsize=30)     # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=30)     # fontsize of the tick labels\n",
    "\n",
    "plt.rc('figure', facecolor='white')\n",
    "\n",
    "#plt.rc('figure.subplot', left=0.2)\n",
    "#plt.rc('figure.subplot', right=0.95)\n",
    "#plt.rc('figure.subplot', bottom=0.2)\n",
    "plt.rc('figure.subplot', top=0.88)\n",
    "\n",
    "\n",
    "font = {'family': 'serif',\n",
    "        'color':  'blue',\n",
    "        'weight': 'bold',\n",
    "        'size': 30,\n",
    "        }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd232049",
   "metadata": {},
   "source": [
    "# PES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a22f30f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class tmp:\n",
    "    pass\n",
    "\n",
    "def superexchange(q, params, full_id):\n",
    "\n",
    "    critical_params = []\n",
    "    default_params = { \"V_11\":0.01, \"V_22\":0.005, \"A_01\":0.001, \"A_12\":0.01, \"A_02\":0.0  }\n",
    "    comn.check_input(params, default_params, critical_params)\n",
    "\n",
    "    V_11 = params[\"V_11\"]\n",
    "    V_22 = params[\"V_22\"]\n",
    "    A_01 = params[\"A_01\"]\n",
    "    A_12 = params[\"A_12\"]\n",
    "    A_02 = params[\"A_02\"]\n",
    "\n",
    "    n = 3\n",
    "\n",
    "    Hdia = CMATRIX(n,n)\n",
    "    Sdia = CMATRIX(n,n)\n",
    "    d1ham_dia = CMATRIXList();  d1ham_dia.append( CMATRIX(n,n) )\n",
    "    dc1_dia = CMATRIXList();  dc1_dia.append( CMATRIX(n,n) )\n",
    "\n",
    "    Id = Cpp2Py(full_id)\n",
    "    indx = Id[-1]\n",
    "\n",
    "    x = q.col(indx).get(0)\n",
    "\n",
    "    Sdia.identity()\n",
    "\n",
    "    Hdia.set(0,0,  0 )\n",
    "    Hdia.set(1,1,  V_11)\n",
    "    Hdia.set(2,2,  V_22)\n",
    "\n",
    "    exp_factor = math.exp(-0.5 *x * x)\n",
    "    dexp_factor = -x * math.exp(-0.5 *x * x)\n",
    "\n",
    "    Hdia.set(0,1,  A_01*exp_factor); Hdia.set(1,0,  A_01*exp_factor)\n",
    "    Hdia.set(1,2,  A_12*exp_factor); Hdia.set(2,1,  A_12*exp_factor)\n",
    "\n",
    "    d1ham_dia[0].set(0, 1, A_01*dexp_factor); d1ham_dia[0].set(1, 0, A_01*dexp_factor)\n",
    "    d1ham_dia[0].set(1, 2, A_12*dexp_factor); d1ham_dia[0].set(2, 1, A_12*dexp_factor)\n",
    "\n",
    "    obj = tmp()\n",
    "    obj.ham_dia = Hdia\n",
    "    obj.ovlp_dia = Sdia\n",
    "    obj.d1ham_dia = d1ham_dia\n",
    "    obj.dc1_dia = dc1_dia\n",
    "\n",
    "    return obj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb187372",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_model(q, params, full_id):\n",
    "\n",
    "    model = params[\"model\"]\n",
    "    res = None\n",
    "\n",
    "    if model==1:\n",
    "        res = Holstein.Holstein2(q, params, full_id)\n",
    "    elif model==2:\n",
    "        res = superexchange(q, params, full_id)\n",
    "    elif model==3:\n",
    "        res = Phenol.Pollien_Arribas_Agostini(q, params, full_id)\n",
    "    else:\n",
    "        pass\n",
    "\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a3e08e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_params1 = {\"model\":1, \"model0\":1, \"nstates\":2, \"E_n\":[0.0, -0.01], \"x_n\":[0.0,  0.5],\"k_n\":[0.002, 0.008],\"V\":0.001} # holstein\n",
    "model_params2 = {\"model\":2, \"model0\":2, \"nstates\":3} # superexchange\n",
    "model_params3 = {\"model\":3, \"model0\":3, \"nstates\":3} # phenol\n",
    "\n",
    "all_model_params = [model_params1, model_params2, model_params3]\n",
    "\n",
    "#################################\n",
    "# Give the model used an index\n",
    "model_indx = 0\n",
    "################################\n",
    "\n",
    "model_params = all_model_params[model_indx]\n",
    "\n",
    "list_states = [x for x in range(model_params[\"nstates\"])]\n",
    "NSTATES = model_params[\"nstates\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2c3b8b0",
   "metadata": {},
   "source": [
    "## 2D PES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8cd133d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def potential(q, params):\n",
    "    full_id = Py2Cpp_int([0,0]) \n",
    "    \n",
    "    return compute_model(q, params, full_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc3d68b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_E(nst, q, _model_params):\n",
    "    res = potential(q, _model_params)\n",
    "    Hdia = np.zeros((nst, nst), dtype=complex)\n",
    "    for i in range(nst):\n",
    "        for j in range(nst):\n",
    "            Hdia[i,j] = res.ham_dia.get(i,j)\n",
    "    H = np.diag(Hdia)\n",
    "    E, U = np.linalg.eig(Hdia) \n",
    "    idx= np.argsort(E) \n",
    "    E = E[idx] \n",
    "    U = U[:,idx]\n",
    "\n",
    "    return H, E"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac00799",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_PES(Rs, Ts, _model_params):\n",
    "    nR = len(Rs)\n",
    "    nT = len(Ts)\n",
    "\n",
    "    dR = Rs[1] - Rs[0]\n",
    "    dT = Ts[1] - Ts[0]\n",
    "\n",
    "    PES_dia = np.zeros((nR, nT, 3)) # nR, nT, nst\n",
    "    PES_adi = np.zeros((nR, nT, 3))\n",
    "\n",
    "    for ix in range(nR):\n",
    "        for iy in range(nT):\n",
    "            q = MATRIX(2,1)\n",
    "    \n",
    "            q.set(0,0, Rs[0] + ix*dR)\n",
    "            q.set(1,0, Ts[0] + iy*dT)\n",
    "    \n",
    "            H, E = get_E(3, q, _model_params)\n",
    "            PES_dia[ix, iy] = H\n",
    "            PES_adi[ix, iy] = E\n",
    "\n",
    "    return PES_dia, PES_adi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efb0f2b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "Angst = 1.88973\n",
    "\n",
    "nR = 100\n",
    "nT = 100\n",
    "\n",
    "dR = 3.0 / nR\n",
    "dT = 2.0 / nT\n",
    "\n",
    "Rs = [(0.5 + dR*ix)*Angst for ix in range(nR + 1)]\n",
    "Ts = [0 + dT*iy for iy in range(nT + 1)]\n",
    "\n",
    "PES_dia, PES_adi = get_PES(Rs, Ts, all_model_params[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650ba3df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_PES_3D(PES, Rs, Ts, title=\"Diabatic PES\"):\n",
    "    # Unit convert\n",
    "    Rs_t = list(Rs)\n",
    "    Rs_t = [x/1.88973 for x in Rs_t]\n",
    "    PES_t = np.copy(PES)\n",
    "    \n",
    "    PES_t *= 27.21\n",
    "\n",
    "    z_lim = 10\n",
    "    PES_t = np.where(PES_t > z_lim, np.nan, PES_t)\n",
    "    \n",
    "    # Generate a meshgrid for plotting\n",
    "    R_mesh, T_mesh = np.meshgrid(Rs_t, Ts)\n",
    "    num_states = PES.shape[2]\n",
    "\n",
    "    #vmin = np.min(PES)\n",
    "    #vmax = np.max(PES)\n",
    "    \n",
    "    # Set up a single 3D plot for all states\n",
    "    fig = plt.figure(figsize=(15, 13))\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "    # Plot each state in the same 3D space\n",
    "    for state in range(num_states):\n",
    "        #surf = ax.plot_surface(\n",
    "        #    R_mesh, T_mesh, PES[:, :, state].T,\n",
    "        #    cmap='viridis', edgecolor='none', alpha=0.6, # Set alpha for transparency\n",
    "        #    vmin=vmin, vmax=vmax\n",
    "        #)\n",
    "        surf = ax.plot_surface(\n",
    "            R_mesh, T_mesh, PES_t[:, :, state].T,\n",
    "            color=colors[clrs_index[state]], edgecolor='none', shade=False\n",
    "        )\n",
    "    \n",
    "    # Color bar for energy and axis labels\n",
    "    #fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5)\n",
    "\n",
    "    ax.margins(x=0, y=0, z=0)\n",
    "\n",
    "    ax.set_xticks([1.0, 2.0, 3.0])\n",
    "    ax.set_yticks([0, math.pi/3])\n",
    "    ax.set_yticklabels([\"0\", r\"    $\\pi/4$\"])\n",
    "    ax.set_zticks([2.0, 4.0, 6.0, 8.0, 10.0])\n",
    "    ax.set_zlim([0, 10])\n",
    "\n",
    "    #ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    #ax.zaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    \n",
    "    ax.set_xlabel(r\"OH length $r$ (Å)\")\n",
    "    ax.set_ylabel(r\"Dih. $\\theta$ (rad)\")\n",
    "    ax.set_zlabel(\"energy (eV)\")\n",
    "    if title != \"\":\n",
    "        ax.set_title(title)\n",
    "\n",
    "    ax.xaxis.labelpad = 30\n",
    "    ax.yaxis.labelpad = 40\n",
    "    ax.zaxis.labelpad = 30\n",
    "\n",
    "    ax.view_init(elev=5, azim=275)\n",
    "\n",
    "    ax.xaxis.pane.fill = False\n",
    "    ax.yaxis.pane.fill = False\n",
    "    ax.zaxis.pane.fill = False\n",
    "\n",
    "    fig.tight_layout()\n",
    "    #plt.savefig(\"fig1_pes_phenol.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acadb1b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_PES_3D(PES_adi, Rs, Ts, title=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "209aac85",
   "metadata": {},
   "source": [
    "## PES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f445af9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_pec(_compute_model, model_params, xmin, xmax, dx, states_of_interest=[0,1], _ndof=1, _active_dof=0, _all_coordinates=[0.0], \n",
    "             show_nac_abs=0, frac_nac=1, _legend_ncol=1, _legend_font=30,  _loc=\"best\", _bbox_to_anchor=\"\"):\n",
    "    X = []\n",
    "    nsteps = int((xmax - xmin) / dx) + 1\n",
    "\n",
    "    for i in range(nsteps):\n",
    "        X.append(xmin + i * dx)\n",
    "    comn.check_input(model_params, {}, [\"nstates\"])\n",
    "    n = model_params[\"nstates\"]\n",
    "    nstates = n\n",
    "    \n",
    "    ham = nHamiltonian(nstates, nstates, _ndof) # ndia, nadi, nnucl\n",
    "    ham.init_all(2)\n",
    "    \n",
    "    hdia, hadi, nac, nac_abs  = [], [], [], []\n",
    "    uij = []              # projecitions of the MOs onto elementary basis\n",
    "    \n",
    "    for k1 in range(nstates):\n",
    "        hadi.append([])\n",
    "        hdia.append([])\n",
    "        nac_k1 = []\n",
    "        nac_abs_k1 = []\n",
    "        uij_k1 = []\n",
    "        for k2 in range(nstates):\n",
    "            uij_k1.append([])\n",
    "            nac_k1.append([])\n",
    "            nac_abs_k1.append([])\n",
    "        uij.append(uij_k1)\n",
    "        nac.append(nac_k1)\n",
    "        nac_abs.append(nac_abs_k1)\n",
    "    \n",
    "    for i in range(nsteps):\n",
    "\n",
    "        scan_coord = MATRIX(_ndof, 1);\n",
    "        for j in range(_ndof):\n",
    "            scan_coord.set(j, 0, _all_coordinates[j])\n",
    "        scan_coord.set(_active_dof, 0, X[i])\n",
    "\n",
    "        # Diabatic properties\n",
    "        ham.compute_diabatic(_compute_model, scan_coord, model_params)\n",
    "\n",
    "        # Adiabatic properties\n",
    "        ham.compute_adiabatic(1);\n",
    "\n",
    "        U = ham.get_basis_transform()\n",
    "        #P = U * U.H()  # population matrix\n",
    "\n",
    "        for k1 in range(nstates):\n",
    "            hadi[k1].append(ham.get_ham_adi().get(k1, k1).real)\n",
    "            hdia[k1].append(ham.get_ham_dia().get(k1, k1).real)\n",
    "\n",
    "            for k2 in range(nstates):\n",
    "                uij[k1][k2].append(U.get(k1,k2).real**2 + U.get(k1,k2).imag**2)\n",
    "                nac_k1_k2 = ham.get_dc1_adi(0).get(k1, k2).real\n",
    "                nac[k1][k2].append(nac_k1_k2)\n",
    "                nac_abs[k1][k2].append( abs(nac_k1_k2) )\n",
    "    \n",
    "    fig, ax1 = plt.subplots(figsize=(10,6), dpi=300)\n",
    "    ax1.tick_params(axis='both', labelsize=30)\n",
    "    \n",
    "    ax1.set_xlabel('$R$ (Bohr)', fontsize=30)\n",
    "    #ax1.set_ylabel('energy (Ha)', labelpad=15, fontsize=30)\n",
    "    ax1.set_ylabel('energy (Ha)', fontsize=30)\n",
    "    ax1.margins(x=0)\n",
    "\n",
    "    for k1 in states_of_interest:\n",
    "        ax1.plot(X, hadi[k1], label='$E_{%i}$' % (k1), linewidth=5, color = colors[clrs_index[k1]], zorder=5)\n",
    "            \n",
    "    for k1 in states_of_interest:\n",
    "        if k1 == 0:\n",
    "            plt.plot(X, hdia[k1], label='$H_{%i%i}$' % (k1,k1), ls=\"dotted\", lw=4, color = colors[clrs_index[k1]], zorder=10)\n",
    "        else:\n",
    "            plt.plot(X, hdia[k1], label='$H_{%i%i}$' % (k1,k1), ls=\"dotted\", lw=4, color = colors[clrs_index[k1]])\n",
    "    \n",
    "    ax1.legend(frameon=False, ncol=2)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    #plt.savefig(\"fig1_pes_holstein.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8b8790",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_pec(compute_model, model_params1, -4.0, 5.0, 0.05, states_of_interest=[0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de83d487",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_pec(compute_model, model_params2, -10.0, 10.0, 0.05, states_of_interest=[0,1,2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b14fbe7",
   "metadata": {},
   "source": [
    "## Holstein"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd560ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def holstein_all(dirs, labels, title=\"Populations\", onlegend=True, name_out=\"res.png\"):\n",
    "\n",
    "    with h5py.File(\"holstein/DVR/exact-model0-icond0-p0_0/data.hdf\", 'r') as f:\n",
    "        time_DVR = f[\"time/data\"][:]/41.0\n",
    "        pop_adi_DVR = np.array(f[\"pop_adi/data\"])\n",
    "        coh_adi_DVR = np.array(f[\"coherence_adi/data\"])\n",
    "    \n",
    "    MQC_se_pop = []\n",
    "    MQC_sh_pop = []\n",
    "    MQC_coh = []\n",
    "    MQC_Etot = []\n",
    "    \n",
    "    for idir in dirs:\n",
    "        with h5py.File(idir + \"/mem_data.hdf\") as f:\n",
    "            time = np.array(f[\"time/data\"])/41.0\n",
    "            se_pop = np.array(f[\"se_pop_adi/data\"])\n",
    "            sh_pop = np.array(f[\"sh_pop_adi/data\"])\n",
    "            coh_adi = np.array(f[\"coherence_adi/data\"])\n",
    "            #C_adi = np.array(f[\"Cadi/data\"][:,:,:])\n",
    "            #hvib_adi = np.array(f[\"hvib_adi/data\"][:,:,:,:])\n",
    "            Etot = np.array(f[\"Etot_ave/data\"])\n",
    "    \n",
    "        Etot0 = Etot[0]\n",
    "        Etot -= Etot0\n",
    "        Etot = np.abs(Etot)/Etot0\n",
    "    \n",
    "        MQC_se_pop.append(se_pop)\n",
    "        MQC_sh_pop.append(sh_pop)\n",
    "        MQC_coh.append(coh_adi)\n",
    "        MQC_Etot.append(Etot)\n",
    "    \n",
    "    nsteps = se_pop.shape[0]\n",
    "    nst = se_pop.shape[1]\n",
    "    \n",
    "    fig, axs = plt.subplots(nrows=3, figsize=(10,16) )\n",
    "    \n",
    "    ## Energy conservation\n",
    "    axs[0].margins(x=0)\n",
    "    axs[0].set_ylabel(r\"$|\\Delta E_{tot}|/E_{tot}(0)$\")\n",
    "    axs[0].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "    \n",
    "    for i, idir in enumerate(dirs):  \n",
    "        axs[0].plot(time, MQC_Etot[i], linewidth=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    # For label\n",
    "    for i, idir in enumerate(dirs):\n",
    "        axs[0].plot([],[], linewidth=5, color=colors[clrs_index[i]], label=labels[i])\n",
    "    axs[0].plot([],[], linewidth=5, color='k', label=\"DVR\")\n",
    "    \n",
    "    axs[0].legend(frameon=False, loc=\"upper left\", bbox_to_anchor=(0, 1.03))\n",
    "    \n",
    "    ## POP\n",
    "    axs[1].margins(x=0)\n",
    "    axs[1].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "    # DVR\n",
    "    axs[1].plot(time_DVR, pop_adi_DVR[:,0], linewidth=7, color=\"k\")\n",
    "    \n",
    "    # MQC\n",
    "    for i, idir in enumerate(dirs):\n",
    "        axs[1].plot(time, MQC_sh_pop[i][:,0], linewidth=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    for i, idir in enumerate(dirs):\n",
    "        axs[1].plot(time, MQC_se_pop[i][:,0], linewidth=5, \n",
    "                           linestyle=\"dotted\", color=colors[clrs_index[i]], zorder = -10)\n",
    "                \n",
    "    axs[1].set_ylabel(rF\"population\")\n",
    "        \n",
    "    axs[-1].set_xlabel(\"time (fs)\")\n",
    "\n",
    "    ## COH\n",
    "    axs[2].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "    axs[2].set_ylabel(rF\"coherence\")\n",
    "    axs[2].margins(x=0)\n",
    "    axs[2].plot(time_DVR, coh_adi_DVR[:,0, 1], linewidth=7, color=\"k\")\n",
    "    for i, idir in enumerate(dirs):\n",
    "        axs[2].plot(time, MQC_coh[i][:, 0, 1], linewidth=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    axs[-1].set_xlabel(\"time (fs)\")\n",
    "        \n",
    "    #fig.legend(loc='upper center', ncol=3, fontsize=30, frameon=False, bbox_to_anchor=(0.565,1.01))\n",
    "    fig.tight_layout()\n",
    "    fig.subplots_adjust(top=0.9)  # Move subplots down to make space for the legend\n",
    "    \n",
    "    #plt.savefig(name_out, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8892c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "dirs = [\"model0-method0-icond0-p0_0/\", \"model0-method1-icond0-p0_0/\", \n",
    "        \"model0-method2-icond0-p0_0-eps0_0/\", \"model0-method3-icond0-p0_0-w0_1\"]\n",
    "\n",
    "dirs = [\"holstein/TSH/\" + idir for idir in dirs]\n",
    "\n",
    "labels=[\"FSSH\", \"QTSH\", \"QTSH-SDM\", \"QTSH-XF\"]\n",
    "\n",
    "holstein_all(dirs, labels, title=\"\", onlegend=True, name_out=\"fig2_holstein.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "767893a5",
   "metadata": {},
   "source": [
    "## Superexchange: scattering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c556debd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_RT_DVR3(_model_indx, _icond_indx, _ks):\n",
    "    R0_QD =np.zeros(len(_ks))\n",
    "    R1_QD =np.zeros(len(_ks))\n",
    "    R2_QD =np.zeros(len(_ks))\n",
    "    T0_QD =np.zeros(len(_ks))\n",
    "    T1_QD =np.zeros(len(_ks))\n",
    "    T2_QD =np.zeros(len(_ks))\n",
    "\n",
    "\n",
    "    for i, k in enumerate(_ks):\n",
    "        if k < 11:\n",
    "            fn = F\"superexchange/DVR/wfc{_model_indx}-icond0-p{k}/wfcr_snap_40_dens_rep_1\".replace(\".\", \"_\")\n",
    "        else:\n",
    "            fn = F\"superexchange/DVR/wfc{_model_indx}-icond0-p{k}/wfcr_snap_20_dens_rep_1\".replace(\".\", \"_\")\n",
    "        which_cols = [0]\n",
    "        nstates = 3\n",
    "    \n",
    "        for state in [0, 1, 2]:\n",
    "            which_cols.append(1+state)\n",
    "        data_exact = data_read.get_data_from_file2(fn, which_cols)\n",
    "    \n",
    "        ngrids = len(data_exact[0])\n",
    "    \n",
    "        #find necessary region\n",
    "        grids = np.array(data_exact[0])\n",
    "    \n",
    "        # Get wfc array\n",
    "        wfcs = []\n",
    "        for ist in range(nstates):\n",
    "            wfcs.append(np.array(data_exact[1+ist]))\n",
    "    \n",
    "        wfcs = np.array(wfcs)\n",
    "        \n",
    "        # find the zero\n",
    "        neg = list(grids).index(max(grids[grids<0]))\n",
    "        \n",
    "        dr = grids[1] - grids[0]\n",
    "        \n",
    "        P0 = wfcs[0]*dr\n",
    "        P1 = wfcs[1]*dr\n",
    "        P2 = wfcs[2]*dr\n",
    "        \n",
    "        R0_QD[i]=np.sum(P0[:neg+1])\n",
    "        T0_QD[i]=np.sum(P0[neg+1:])\n",
    "        R1_QD[i]=np.sum(P1[:neg+1])\n",
    "        T1_QD[i]=np.sum(P1[neg+1:])\n",
    "        R2_QD[i]=np.sum(P2[:neg+1])\n",
    "        T2_QD[i]=np.sum(P2[neg+1:])\n",
    "        \n",
    "    return R0_QD, R1_QD, R2_QD, T0_QD, T1_QD, T2_QD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2780e3e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_RT_MQC3(_model_indx, _method_indx, _icond_indx, _ks):\n",
    "    R0 = np.zeros(len(_ks))\n",
    "    R1 = np.zeros(len(_ks))\n",
    "    R2 = np.zeros(len(_ks))\n",
    "    T0 = np.zeros(len(_ks))\n",
    "    T1 = np.zeros(len(_ks))\n",
    "    T2 = np.zeros(len(_ks))\n",
    "    \n",
    "    for i, k in enumerate(_ks):\n",
    "        if _method_indx == 5:\n",
    "            dir_dyn = \"superexchange/TSH/\" + F\"model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{k:.1f}-eps0_1\".replace(\".\", \"_\")\n",
    "        elif _method_indx == 6:\n",
    "            dir_dyn = \"superexchange/TSH/\" + F\"model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{k:.1f}-w0_1\".replace(\".\", \"_\")\n",
    "        else:\n",
    "            dir_dyn = \"superexchange/TSH/\" + F\"model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{k:.1f}\".replace(\".\", \"_\")\n",
    "        \n",
    "        R0_tmp, R1_tmp, R2_tmp, T0_tmp, T1_tmp, T2_tmp = 0, 0, 0, 0, 0, 0\n",
    "        with h5py.File(dir_dyn + \"/mem_data.hdf\", 'r') as f:\n",
    "            qs = np.array(f[\"q/data\"][:,:,0])\n",
    "            ntraj = qs.shape[1]\n",
    "            \n",
    "            for itraj in range(ntraj):\n",
    "                act_st = f[\"states/data\"][-1, itraj]\n",
    "                q_end = f[\"q/data\"][-1,itraj,0]\n",
    "                p_end = f[\"p/data\"][-1,itraj,0]\n",
    "                #print(act_st, q_end, p_end)\n",
    "                if act_st==0: # at 0\n",
    "                    if q_end > 0.0:\n",
    "                        T0_tmp += 1\n",
    "                    elif q_end < 0.0:\n",
    "                        R0_tmp += 1\n",
    "                elif act_st==1: # at 1\n",
    "                    if q_end > 0.0:\n",
    "                        T1_tmp += 1\n",
    "                    elif q_end < 0.0:\n",
    "                        R1_tmp += 1\n",
    "                else: # at 2\n",
    "                    if q_end > 0.0:\n",
    "                        T2_tmp += 1\n",
    "                    elif q_end < 0.0:\n",
    "                        R2_tmp += 1\n",
    "            R0[i] = R0_tmp/ntraj\n",
    "            T0[i] = T0_tmp/ntraj\n",
    "            R1[i] = R1_tmp/ntraj\n",
    "            T1[i] = T1_tmp/ntraj\n",
    "            R2[i] = R2_tmp/ntraj\n",
    "            T2[i] = T2_tmp/ntraj\n",
    "    \n",
    "    return R0, R1, R2, T0, T1, T2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd7c7b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "incr_k = 1.0\n",
    "ks = [5.0 + incr_k*ik for ik in range(16)]\n",
    "\n",
    "R0_FSSH, R1_FSSH, R2_FSSH, T0_FSSH, T1_FSSH, T2_FSSH = get_RT_MQC3(1, 0, 0, ks)\n",
    "R0_QTSH, R1_QTSH, R2_QTSH, T0_QTSH, T1_QTSH, T2_QTSH = get_RT_MQC3(1, 1, 0, ks)\n",
    "R0_QTSH_SDM, R1_QTSH_SDM, R2_QTSH_SDM, T0_QTSH_SDM, T1_QTSH_SDM, T2_QTSH_SDM = get_RT_MQC3(1, 2, 0, ks)\n",
    "R0_QTSH_XF, R1_QTSH_XF, R2_QTSH_XF, T0_QTSH_XF, T1_QTSH_XF, T2_QTSH_XF = get_RT_MQC3(1, 3, 0, ks)\n",
    "R0_QD, R1_QD, R2_QD, T0_QD, T1_QD, T2_QD = get_RT_DVR3(1, 0, ks_DVR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea194c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "incr_k = 0.2\n",
    "ks_low = [3.0 + incr_k*ik for ik in range(11)]\n",
    "\n",
    "R0_FSSH_low, R1_FSSH_low, R2_FSSH_low, T0_FSSH_low, T1_FSSH_low, T2_FSSH_low = get_RT_MQC3_low(0, 0, 0, ks_low)\n",
    "R0_QTSH_low, R1_QTSH_low, R2_QTSH_low, T0_QTSH_low, T1_QTSH_low, T2_QTSH_low = get_RT_MQC3_low(0, 1, 0, ks_low)\n",
    "R0_QD_low, R1_QD_low, R2_QD_low, T0_QD_low, T1_QD_low, T2_QD_low = get_RT_DVR3_low(0, 0, ks_DVR_low)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df29348c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(ncols=3, figsize=(27,9))\n",
    "\n",
    "## For high k, transmission\n",
    "ax[0].margins(x=0)\n",
    "ax[0].xaxis.set_tick_params(labelsize=50)\n",
    "ax[0].yaxis.set_tick_params(labelsize=50)\n",
    "ax[0].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "ax[0].set_xticks([5, 10, 20])\n",
    "ax[0].set_yticks([1.0, 0.8, 0.2, 0.0])\n",
    "ax[0].set_ylabel(r\"transmission $T_{n}$\", fontsize=50)\n",
    "ax[0].set_xlabel(r\"$\\hbar k_{0}$ (a.u.)\", fontsize=50)\n",
    "\n",
    "# T0\n",
    "ax[0].plot(ks_DVR, T0_QD, linewidth=10, color=\"k\", label=\"DVR\", zorder=-10)\n",
    "ax[0].scatter(ks, T0_FSSH, marker=\"v\", color=colors[clrs_index[0]], s=150, clip_on=False, label=\"FSSH\")\n",
    "ax[0].scatter(ks, T0_QTSH, marker=\"v\", color=colors[clrs_index[1]], s=150, clip_on=False, label=\"QTSH\")\n",
    "ax[0].scatter(ks, T0_QTSH_SDM, marker=\"o\", color=colors[clrs_index[2]], s=150, clip_on=False, label=\"QTSH-SDM\")\n",
    "ax[0].scatter(ks, T0_QTSH_XF, marker=\"o\", color=colors[clrs_index[3]], s=150, clip_on=False, label=\"QTSH-XF\")\n",
    "\n",
    "ax[0].legend(handlelength=0.7, frameon=False, fontsize=50, labelspacing=0.3, loc=\"center\")\n",
    "\n",
    "ax[0].annotate(r\"$n=0$\", (15, 0.85), fontsize=50 )\n",
    "\n",
    "# T1\n",
    "ax[0].plot(ks_DVR, T1_QD, linewidth=10, color=\"k\", label=\"DVR\", zorder=-10)\n",
    "ax[0].scatter(ks, T1_FSSH, marker=\"v\", color=colors[clrs_index[0]], s=150, clip_on=False, label=\"FSSH\")\n",
    "ax[0].scatter(ks, T1_QTSH, marker=\"v\", color=colors[clrs_index[1]], s=150, clip_on=False, label=\"QTSH\")\n",
    "ax[0].scatter(ks, T1_QTSH_SDM, marker=\"o\", color=colors[clrs_index[2]], s=150, clip_on=False, label=\"QTSH-SDM\")\n",
    "ax[0].scatter(ks, T1_QTSH_XF, marker=\"o\", color=colors[clrs_index[3]], s=150, clip_on=False, label=\"QTSH-XF\")\n",
    "\n",
    "ax[0].annotate(r\"$n=1$\", (15, 0.08), fontsize=50 )\n",
    "\n",
    "## For lower k, transmission\n",
    "ax[1].margins(x=0)\n",
    "ax[1].xaxis.set_tick_params(labelsize=50)\n",
    "ax[1].yaxis.set_tick_params(labelsize=50)\n",
    "ax[1].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "ax[1].set_yticks([0.8, 0.2, 0.0])\n",
    "ax[1].set_ylabel(r\"transmission $T_{n}$\", fontsize=50)\n",
    "ax[1].set_xlabel(r\"$\\hbar k_{0}$ (a.u.)\", fontsize=50)\n",
    "\n",
    "# T0\n",
    "ax[1].plot(ks_DVR_low, T0_QD_low, linewidth=10, color=\"k\", label=\"DVR\", zorder=-10)\n",
    "ax[1].scatter(ks_low, T0_FSSH_low, marker=\"v\", color=colors[clrs_index[0]], s=150, clip_on=False, label=\"FSSH\")\n",
    "ax[1].scatter(ks_low, T0_QTSH_low, marker=\"v\", color=colors[clrs_index[1]], s=150, clip_on=False, label=\"QTSH\")\n",
    "\n",
    "ax[1].legend(handlelength=0.7, frameon=False, fontsize=50, labelspacing=0.3, loc=\"center\")\n",
    "\n",
    "ax[1].annotate(r\"$n=0$\", (4.35, 0.67), fontsize=50 )\n",
    "\n",
    "# T1\n",
    "ax[1].plot(ks_DVR_low, T1_QD_low, linewidth=10, color=\"k\", label=\"DVR\", zorder=-10)\n",
    "ax[1].scatter(ks_low, T1_FSSH_low, marker=\"v\", color=colors[clrs_index[0]], s=150, clip_on=False, label=\"FSSH\")\n",
    "ax[1].scatter(ks_low, T1_QTSH_low, marker=\"v\", color=colors[clrs_index[1]], s=150, clip_on=False, label=\"QTSH\")\n",
    "\n",
    "ax[1].annotate(r\"$n=1$\", (4.35, 0.03), fontsize=50 )\n",
    "\n",
    "## For lower k, reflection\n",
    "ax[2].margins(x=0)\n",
    "ax[2].xaxis.set_tick_params(labelsize=50)\n",
    "ax[2].yaxis.set_tick_params(labelsize=50)\n",
    "ax[2].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "ax[2].set_yticks([0.2, 0.0])\n",
    "ax[2].set_ylabel(r\"reflection $R_{n}$\", fontsize=50)\n",
    "ax[2].set_xlabel(r\"$\\hbar k_{0}$ (a.u.)\", fontsize=50)\n",
    "\n",
    "# R0\n",
    "ax[2].plot(ks_DVR_low, R0_QD_low, linewidth=10, color=\"k\", label=\"DVR\", zorder=-10)\n",
    "ax[2].scatter(ks_low, R0_FSSH_low, marker=\"v\", color=colors[clrs_index[0]], s=150, clip_on=False, label=\"FSSH\")\n",
    "ax[2].scatter(ks_low, R0_QTSH_low, marker=\"v\", color=colors[clrs_index[1]], s=150, clip_on=False, label=\"QTSH\")\n",
    "\n",
    "ax[2].annotate(r\"$n=0$\", (3.6, 0.1), fontsize=50 )\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "#plt.savefig(\"fig2_scattering.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a407a445",
   "metadata": {},
   "source": [
    "## Superexchange: dynamics evolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6f5f8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the canonical momentum\n",
    "def get_p_cano(file):\n",
    "\n",
    "    with h5py.File(file, 'r') as f:\n",
    "        p = np.array(f[\"p/data\"][:,:,0])\n",
    "        C_adi = np.array(f[\"Cadi/data\"][:,:,:])\n",
    "        dc1_adi = np.array(f[\"dc1_adi/data\"][:,:,:,:,:])\n",
    "    \n",
    "    p_cano = np.zeros(p.shape)\n",
    "    p_cano = np.copy(p)\n",
    "\n",
    "    ntraj = p.shape[1]\n",
    "    nst = C_adi.shape[2]\n",
    "\n",
    "    for itraj in range(ntraj):\n",
    "        for ist in range(nst):\n",
    "            for jst in range(nst):\n",
    "                if(ist>=jst): continue\n",
    "                p_cano[:,itraj] += 2* (C_adi[:, itraj, ist] * C_adi[:, itraj, jst].conj()).imag * dc1_adi.real[:,itraj,0, ist, jst]\n",
    "    \n",
    "    return p_cano"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc4c1bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the diagonal and coherence energy\n",
    "def get_E_cano(file):\n",
    "\n",
    "    with h5py.File(file, 'r') as f:\n",
    "        time = np.array(f[\"time/data\"])\n",
    "        Epot = np.array(f[\"Epot_ave/data\"])\n",
    "        Ekin = np.array(f[\"Ekin_ave/data\"])\n",
    "        Etot = np.array(f[\"Etot_ave/data\"])\n",
    "        states = np.array(f[\"states/data\"])\n",
    "        C_adi = np.array(f[\"Cadi/data\"][:,:,:])\n",
    "        hvib_adi = np.array(f[\"hvib_adi/data\"][:,:,:,:])\n",
    "        dc1_adi = np.array(f[\"dc1_adi/data\"][:,:,:,:,:])\n",
    "    \n",
    "    p_cano = get_p_cano(file)\n",
    "\n",
    "    Etot_cano = np.zeros(Etot.shape)\n",
    "    Etot_cano = np.copy(Etot_cano)\n",
    "\n",
    "    Ediag = np.zeros(Etot_cano.shape)\n",
    "    Ecoh = np.zeros(Etot_cano.shape)\n",
    "\n",
    "    ntraj = C_adi.shape[1]\n",
    "    nst = C_adi.shape[2]\n",
    "\n",
    "    M = 2000.\n",
    "\n",
    "    for itraj in range(ntraj):\n",
    "        # Diagonal part\n",
    "        Ediag[:] += 0.5 * p_cano[:,itraj] * p_cano[:,itraj] / M \n",
    "        for ist in range(nst):\n",
    "            delta = (ist == states[:,itraj]).astype(float)\n",
    "            Ediag[:] += delta* hvib_adi.real[:, itraj, ist, ist]\n",
    "        \n",
    "        # Coherence part\n",
    "        for ist in range(nst):\n",
    "            for jst in range(nst):\n",
    "                if(ist>=jst): continue\n",
    "                Ecoh[:] -= 2* (C_adi[:, itraj, ist] * C_adi[:, itraj, jst].conj()).imag * (-hvib_adi.imag[:, itraj, ist, jst])\n",
    "\n",
    "    Ediag /= ntraj\n",
    "    Ecoh /= ntraj\n",
    "\n",
    "    Etot_cano = Ediag + Ecoh\n",
    "    \n",
    "    return Ediag, Ecoh, Etot_cano"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbb78c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def comp_epc(file, file_DVR, pos_label=0, broken = (-0.1, 0.3, 3.8, 4.2), title=\"QTSH\", onlegend=True, name_out=\"res.png\"):\n",
    "\n",
    "    with h5py.File(file_DVR, 'r') as f:\n",
    "        pop_adi_DVR = np.array(f[\"pop_adi/data\"])\n",
    "        coh_adi_DVR = np.array(f[\"coherence_adi/data\"])\n",
    "    \n",
    "    with h5py.File(file) as f:\n",
    "        time = np.array(f[\"time/data\"])/41.0\n",
    "        se_pop = np.array(f[\"se_pop_adi/data\"])\n",
    "        sh_pop = np.array(f[\"sh_pop_adi/data\"])\n",
    "        coh_adi = np.array(f[\"coherence_adi/data\"])\n",
    "        Etot = np.array(f[\"Etot_ave/data\"])\n",
    "    \n",
    "    nst = coh_adi.shape[1]\n",
    "    \n",
    "    Ediag = np.loadtxt(name_out+\"_edig.txt\")\n",
    "    Ecoh = np.loadtxt(name_out+\"_ecoh.txt\")\n",
    "    Etot_cano = np.loadtxt(name_out+\"_etot_cano.txt\")\n",
    "    \n",
    "    fig, axs = plt.subplots(nrows=4, figsize=(10,16), gridspec_kw={'height_ratios': [0.8, 0.8, 2, 2]}) # axs[0], axs[1] is for broken axes\n",
    "    \n",
    "    plt.suptitle(title, fontsize=30)\n",
    "    axs[0].margins(x=0)\n",
    "    axs[0].set_ylabel(r\"energy (mHa)\")\n",
    "    axs[0].yaxis.set_label_coords(-0.1, pos_label)\n",
    "    axs[0].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    \n",
    "    axs[0].plot(time, Ediag*1000, color=colors[clrs_index[0]], linewidth=5, label=r\"$E_{diag}$\")\n",
    "    axs[0].plot(time, Ecoh*1000, color=colors[clrs_index[1]], linewidth=5, label=r\"$E_{coh}$\")\n",
    "    axs[0].plot(time, Etot_cano*1000, color=colors[clrs_index[2]], linewidth=5, label=r\"$E_{tot}$\")\n",
    "    \n",
    "    if(onlegend):\n",
    "        axs[0].legend(frameon=False, ncol=3, fontsize=30, loc='upper right', bbox_to_anchor=(1.0, 0.80), \n",
    "                      handlelength=1.5, columnspacing=1.0)\n",
    "\n",
    "    axs[1].margins(x=0)\n",
    "    axs[1].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    \n",
    "    axs[1].plot(time, Ediag*1000, color=colors[clrs_index[0]], linewidth=5, label=r\"$E_{diag}$\")\n",
    "    axs[1].plot(time, Ecoh*1000, color=colors[clrs_index[1]], linewidth=5, label=r\"$E_{coh}$\")\n",
    "    axs[1].plot(time, Etot_cano*1000, color=colors[clrs_index[2]], linewidth=5, label=r\"$E_{tot}$\")\n",
    "        \n",
    "    # breaking\n",
    "    diff = (broken[3]- broken[2])*0.1\n",
    "    axs[0].set_ylim(broken[2]-diff,broken[3]+diff)\n",
    "    \n",
    "    diff = (broken[1]- broken[0])*0.1\n",
    "    axs[1].set_ylim(broken[0]-diff, broken[1]+diff)\n",
    "    \n",
    "    # hide the spines between axs[0] and axs[1]\n",
    "    axs[0].spines.bottom.set_visible(False)\n",
    "    axs[1].spines.top.set_visible(False)\n",
    "    axs[0].xaxis.tick_top()\n",
    "    axs[0].tick_params(labeltop=False)  # don't put tick labels at the top\n",
    "    axs[1].xaxis.tick_bottom()\n",
    "    \n",
    "    d = .5  # proportion of vertical to horizontal extent of the slanted line\n",
    "    kwargs = dict(marker=[(-1, -d), (1, d)], markersize=12,\n",
    "                  linestyle=\"none\", color='k', mec='k', mew=1, clip_on=False)\n",
    "    axs[0].plot([0, 1], [0, 0], transform=axs[0].transAxes, **kwargs)\n",
    "    axs[1].plot([0, 1], [1, 1], transform=axs[1].transAxes, **kwargs)\n",
    "    \n",
    "    # POP\n",
    "    axs[2].margins(x=0)\n",
    "    axs[2].set_ylabel(\"population\")\n",
    "    axs[2].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    axs[2].set_yticks([0.0, 0.5, 1.0])\n",
    "    axs[2].plot([],[], color=\"k\", linewidth=5, label=\"DVR\")\n",
    "    \n",
    "    for ist in range(nst):\n",
    "        axs[2].plot(time, pop_adi_DVR[:,ist], linewidth=7, color=\"k\")\n",
    "        axs[2].plot(time, sh_pop[:,ist], linewidth=5, color=colors[clrs_index[ist]], label=rF\"$P^{{SH}}_{{{ist}}}$\")\n",
    "    \n",
    "    for ist in range(nst):\n",
    "        axs[2].plot(time, se_pop[:,ist], linestyle=\"dotted\", linewidth=5, color=colors[clrs_index[ist]], \n",
    "                    label=rF\"$P^{{SE}}_{{{ist}}}$\")\n",
    "    \n",
    "    if(onlegend):\n",
    "        axs[2].legend(ncol=2, frameon=False, fontsize=30, loc='upper right', labelspacing=0.6 , bbox_to_anchor=(1,1.01))\n",
    "    \n",
    "    # COH\n",
    "    axs[3].plot([],[], color=\"k\", label=\"DVR\", linewidth=5)\n",
    "    \n",
    "    cnt = -1\n",
    "    for ist in range(nst):\n",
    "        for jst in range(nst):\n",
    "            if ist < jst:\n",
    "                cnt += 1\n",
    "                axs[3].plot(time, coh_adi_DVR[:,ist,jst], color=\"k\", linewidth=7)\n",
    "                axs[3].plot(time, coh_adi[:,ist,jst], color=colors[clrs_index[cnt]], linewidth=5, \n",
    "                            label=rF\"$\\langle |\\rho_{{{ist}{jst}}}|^2 \\rangle$\")\n",
    "    \n",
    "    axs[3].margins(x=0)\n",
    "    axs[3].yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "    axs[3].set_yticks([0.0, 0.1, 0.2])\n",
    "    axs[3].set_ylabel(\"coherence\")\n",
    "    axs[3].set_xlabel(\"time (fs)\")\n",
    "    \n",
    "    if(onlegend):\n",
    "        axs[3].legend(ncol=2, frameon=False, loc=\"center right\", fontsize=30)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    #plt.savefig(name_out + \".png\", dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c96309",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 3, 0, 4.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, pos_label=-0.2, title=rF\"$\\hbar k_{0}=4.0$ a.u.\", onlegend=False, name_out=\"fig4_qtsh_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb51a92a",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 3, 0, 5.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, pos_label=-0.25, broken = (-0.1, 0.5, 5.8, 6.5), title=rF\"$\\hbar k_{0}=5.0$ a.u.\", onlegend=True, name_out=\"fig4_qtsh_5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b91b0002",
   "metadata": {},
   "source": [
    "## Superexchange: linear correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952f690e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx = 0\n",
    "_method_indx = 1\n",
    "_icond_indx = 0\n",
    "\n",
    "ks = [3.0 + 0.2*ik for ik in range(11)]\n",
    "\n",
    "tavg_de = [] # RMSE\n",
    "tavg_dP = []\n",
    "\n",
    "for i, k in enumerate(ks):\n",
    "    dir_dyn = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{k:.1f}\"\n",
    "    dir_dyn = dir_dyn.replace(\".\", \"_\")\n",
    "        \n",
    "    with h5py.File(dir_dyn + \"/mem_data.hdf\", \"r\") as f:\n",
    "        etot = np.array(f[\"Etot_ave/data\"])\n",
    "        se_pop = np.array(f[\"se_pop_adi/data\"])\n",
    "        sh_pop = np.array(f[\"sh_pop_adi/data\"])\n",
    "\n",
    "    etot0 = etot[0]\n",
    "    etot -= etot0 # E - E(0)\n",
    "    etot /= etot0\n",
    "    etot = etot * etot\n",
    "    tavg_de.append( np.sqrt(np.average(etot)) )\n",
    "\n",
    "    dP = np.zeros(se_pop.shape[0])\n",
    "    for ist in range(3):\n",
    "        dP += (se_pop[:,ist] - sh_pop[:,ist]) **2\n",
    "    tavg_dP.append( np.sqrt(np.average(dP)/3) )\n",
    "\n",
    "coefficients = np.polyfit(tavg_de, tavg_dP, 1)  # Linear fit (degree=1)\n",
    "slope, intercept = coefficients\n",
    "\n",
    "fit_line = slope * np.array(tavg_de) + intercept\n",
    "\n",
    "# Calculate R^2\n",
    "y_mean = np.mean(tavg_dP)\n",
    "ss_total = np.sum((tavg_dP - y_mean) ** 2)\n",
    "ss_residual = np.sum((tavg_dP - fit_line) ** 2)\n",
    "r2 = 1 - (ss_residual / ss_total)\n",
    "\n",
    "x = np.linspace(0.0, 0.2, 100)\n",
    "fit_line_for_plot = slope * np.array(x) + intercept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2890f6b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (10, 9))\n",
    "\n",
    "ax.xaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "\n",
    "ax.set_xlabel(r\"RMSE of $[E_{tot}-E_{tot}(0)]/E_{tot}(0)$\")\n",
    "ax.set_ylabel(r\"RMSE of $P^{SE}-P^{SH}$\")\n",
    "\n",
    "ax.set_xlim([0.0, tavg_de[0]+0.01])\n",
    "ax.set_ylim([tavg_dP[-2]-0.015, tavg_dP[0]+0.015])\n",
    "\n",
    "#ax.margins(x=0, y=0)\n",
    "ax.plot(x, fit_line_for_plot, color='k', lw=5, linestyle=\"dashed\", label=r'Linear Fit, $R^{2}=$' + f'{r2:.2f}')\n",
    "ax.scatter(tavg_de, tavg_dP, color=colors[clrs_index[0]], label=\"QTSH\")\n",
    "\n",
    "ax.legend(frameon=False, handlelength=1.5, fontsize=30)\n",
    "\n",
    "for i, k in enumerate(ks):\n",
    "    if k == 4.8:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]-0.002, tavg_dP[i]-0.012), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k == 4.6:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]-0.01, tavg_dP[i]+0.003), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k == 5.0:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]+0.002, tavg_dP[i]-0.004), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k == 3.0:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]-0.002, tavg_dP[i]+0.002), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k in [4.2, 4.0]:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]-0.005, tavg_dP[i]+0.004), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k in [3.6, 3.4]:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]-0.0025, tavg_dP[i]+0.004), fontsize=30, color=colors[clrs_index[0]])\n",
    "    elif k == 3.2:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]+0.001, tavg_dP[i]-0.006), fontsize=30, color=colors[clrs_index[0]])\n",
    "    else:\n",
    "        ax.annotate(F\"{k:.1f}\", (tavg_de[i]+0.001, tavg_dP[i]+0.001), fontsize=30, color=colors[clrs_index[0]])\n",
    "\n",
    "fig.tight_layout()\n",
    "        \n",
    "#plt.savefig(\"fig5_pE.png\", dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf04b8ed",
   "metadata": {},
   "source": [
    "## Superexchange: QTSH-SDM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b31ca47b",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 2, 0, 4.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, broken = (-0.1, 0.3, 3.8, 4.1), title=r\"$C=0.01$ Ha\", onlegend=False, name_out=\"fig6_qtsh_sdm_0_01_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "639a9c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 2, 0, 4.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, broken = (-0.1, 0.3, 3.8, 4.1), title=r\"$C=0.1$ Ha\", onlegend=False, name_out=\"fig6_qtsh_sdm_0_1_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c1a0692",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 2, 0, 4.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps2_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, broken = (-0.1, 0.3, 3.8, 4.1), title=r\"$C=2.0$ Ha\", onlegend=False, name_out=\"fig6_qtsh_sdm_2_0_4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0de378cf",
   "metadata": {},
   "source": [
    "## Superexchange: QTSH-IDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c315c27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 5, 0, 4.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, broken = (-0.2, 0.1, 4.0, 4.6), pos_label=-0.25, title=rF\"$\\hbar k_{0}=4.0$ a.u.\", onlegend=True, name_out=\"s_qtsh_ida_4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e221e220",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 5, 0, 5.0\n",
    "\n",
    "if _method_indx == 2:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 3:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, broken = (-0.2, 0.1, 4.0, 4.6), pos_label=-0.25, title=rF\"$\\hbar k_{0}=5.0$ a.u.\", onlegend=True, name_out=\"s_qtsh_ida_5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b64b5c",
   "metadata": {},
   "source": [
    "## Superexchange: QTSH-XF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a384f7dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 3, 0, 5.0\n",
    "\n",
    "if _method_indx == 5:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 6:\n",
    "    dir_name_suff = F\"-w0_1\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, pos_label=-0.3, broken = (-0.1, 0.5, 6.0, 6.5), title=r\"$\\sigma=0.1$ Bohr\", onlegend=True, name_out=\"fig7_qtsh_xf_0_1_5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df4fbc12",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 3, 0, 5.0\n",
    "\n",
    "if _method_indx == 5:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 6:\n",
    "    dir_name_suff = F\"-w1_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, pos_label=-0.3, broken = (-0.1, 0.5, 6.0, 6.5), title=r\"$\\sigma=1.0$ Bohr\", onlegend=True, name_out=\"fig7_qtsh_xf_1_0_5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ccc6ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "_model_indx, _method_indx, _icond_indx, _k = 0, 3, 0, 5.0\n",
    "\n",
    "if _method_indx == 5:\n",
    "    dir_name_suff = F\"-eps0_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "elif _method_indx == 6:\n",
    "    dir_name_suff = F\"-w5_0\"\n",
    "    dir_name_suff = dir_name_suff.replace(\".\", \"_\")\n",
    "else:\n",
    "    dir_name_suff = \"\"\n",
    "    \n",
    "file = F\"superexchange/TSH/model{_model_indx}-method{_method_indx}-icond{_icond_indx}-p{_k:.1f}\" + dir_name_suff\n",
    "file = file.replace(\".\", \"_\")\n",
    "file += \"/mem_data.hdf\"\n",
    "\n",
    "file_DVR = F\"superexchange/DVR/exact-model{_model_indx}-icond{_icond_indx}-p{_k:.1f}\"\n",
    "file_DVR = file_DVR.replace(\".\", \"_\")\n",
    "file_DVR += \"/data.hdf\"\n",
    "\n",
    "comp_epc(file, file_DVR, pos_label=-0.3, broken = (-0.1, 0.5, 6.0, 6.5), title=r\"$\\sigma=5.0$ Bohr\", onlegend=True, name_out=\"fig7_qtsh_xf_5_0_5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0102619d",
   "metadata": {},
   "source": [
    "## Phenol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e652abec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the dissociation probability from DVR\n",
    "def plot_disso_DVR(pref):\n",
    "    nstep = 4001\n",
    "    freq = 100\n",
    "    \n",
    "    nst = 3\n",
    "    r_cut = 2.6*units.Angst\n",
    "    nx, ny = 4096, 4096\n",
    "    dx, dy = 0.02, 0.03\n",
    "\n",
    "    n_snapshots = nstep // freq + 1\n",
    "    P_disso = np.zeros(n_snapshots)\n",
    "\n",
    "    for istep in range(0, nstep, freq):\n",
    "        wfc_dens = None\n",
    "\n",
    "        for ist in range(nst):\n",
    "            filename = f\"{pref}/wfcr_snap_{istep // freq}_state_{ist}_dens_rep_1\"\n",
    "            data_exact = np.loadtxt(filename)\n",
    "\n",
    "            if ist == 0:\n",
    "                wfc_dens = data_exact[:, 2]\n",
    "            else:\n",
    "                wfc_dens += data_exact[:, 2]\n",
    "\n",
    "        # Reshape data\n",
    "        x = data_exact[:, 0].reshape(nx, ny)\n",
    "        wfc_dens = wfc_dens.reshape(nx, ny)\n",
    "\n",
    "        # Calculate dissociation probability for the snapshot\n",
    "        mask = x > r_cut\n",
    "        P_disso[istep // freq] = np.sum(wfc_dens[mask]) * dx * dy\n",
    "\n",
    "    # Save and plot results\n",
    "    time = np.arange(0, nstep, freq)\n",
    "    np.savetxt(\"disso.txt\", P_disso)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826fc091",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_disso_DVR(\"phenol/DVR/wfc2-icond0-p15_0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82bc1efb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def phenol_all(dirs, labels, title=\"Populations\", onlegend=True, name_out=\"res.png\"):\n",
    "\n",
    "    with h5py.File(\"phenol/DVR/exact-model2-icond0-p15_0/data.hdf\", 'r') as f:\n",
    "        time_DVR = f[\"time/data\"][:]/41.0\n",
    "        pop_adi_DVR = np.array(f[\"pop_adi/data\"])\n",
    "        coh_adi_DVR = np.array(f[\"coherence_adi/data\"])\n",
    "    \n",
    "    MQC_se_pop = []\n",
    "    MQC_sh_pop = []\n",
    "    MQC_coh = []\n",
    "    MQC_Etot = []\n",
    "    \n",
    "    for idir in dirs:\n",
    "        with h5py.File(idir + \"/mem_data.hdf\") as f:\n",
    "            time = np.array(f[\"time/data\"])/41.0\n",
    "            se_pop = np.array(f[\"se_pop_adi/data\"])\n",
    "            sh_pop = np.array(f[\"sh_pop_adi/data\"])\n",
    "            coh_adi = np.array(f[\"coherence_adi/data\"])\n",
    "            Etot = np.array(f[\"Etot_ave/data\"])\n",
    "    \n",
    "        Etot0 = Etot[0]\n",
    "        Etot -= Etot0\n",
    "        Etot = np.abs(Etot)/Etot0\n",
    "    \n",
    "        MQC_se_pop.append(se_pop)\n",
    "        MQC_sh_pop.append(sh_pop)\n",
    "        MQC_coh.append(coh_adi)\n",
    "        MQC_Etot.append(Etot)\n",
    "    \n",
    "    nsteps = se_pop.shape[0]\n",
    "    nst = se_pop.shape[1]\n",
    "    \n",
    "    fig, axs = plt.subplots(nrows=4, ncols=2, figsize=(12,15) )\n",
    "    \n",
    "    ## Energy conservation\n",
    "    axs[0][0].margins(x=0)\n",
    "    axs[0][0].set_ylabel(r\"$|\\Delta E_{tot}|/E_{tot}(0)$\")\n",
    "    axs[0][0].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "    \n",
    "    for i, idir in enumerate(dirs):  \n",
    "        axs[0][0].plot(time, MQC_Etot[i], linewidth=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    # For label\n",
    "    for i, idir in enumerate(dirs):\n",
    "        axs[0][0].plot([],[], linewidth=5, color=colors[clrs_index[i]], label=labels[i])\n",
    "    axs[0][0].plot([],[], linewidth=5, color='k', label=\"DVR\")\n",
    "    \n",
    "    ## POP\n",
    "    for ist in range(nst):\n",
    "        axs[ist+1][0].margins(x=0)\n",
    "        axs[ist+1][0].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "        # DVR\n",
    "        axs[ist+1][0].plot(time_DVR, pop_adi_DVR[:,ist], linewidth=7, color=\"k\")\n",
    "        \n",
    "        # MQC\n",
    "        for i, idir in enumerate(dirs):\n",
    "            axs[ist+1][0].plot(time, MQC_sh_pop[i][:,ist], linewidth=5, color=colors[clrs_index[i]])\n",
    "        \n",
    "        for i, idir in enumerate(dirs):\n",
    "            axs[ist+1][0].plot(time, MQC_se_pop[i][:,ist], linewidth=5, \n",
    "                               linestyle='dotted', color=colors[clrs_index[i]], zorder=-10)\n",
    "                    \n",
    "        axs[ist+1][0].set_ylabel(rF\"$\\langle  \\rho_{{{ist}{ist}}} \\rangle $\")\n",
    "        \n",
    "    axs[-1][0].set_xlabel(\"time (fs)\")\n",
    "\n",
    "    ## Dissociation Prob.\n",
    "    r_cut = 2.6*units.Angst\n",
    "        \n",
    "    P_disso_DVR = np.loadtxt(\"disso.txt\")\n",
    "    time_DVR2 = np.array([100*x for x in range(41)])/41\n",
    "    \n",
    "    axs[0][1].plot(time_DVR2, P_disso_DVR, lw=7, color='k')\n",
    "    \n",
    "    axs[0][1].margins(x=0)\n",
    "    for i, idir in enumerate(dirs):\n",
    "        \n",
    "        P_disso = np.zeros(nsteps)    \n",
    "        with h5py.File(idir + \"/mem_data.hdf\") as f:\n",
    "            time = np.array(f[\"time/data\"])/41\n",
    "            states = np.array(f[\"states/data\"])\n",
    "            q = np.array(f[\"q/data\"][:,:,:])\n",
    " \n",
    "        ntraj = q.shape[1]\n",
    "        \n",
    "        # Calculate dissociation probabilities\n",
    "        P_disso = (q[:nsteps, :, 0] > r_cut).sum(axis=1) / ntraj\n",
    "\n",
    "        axs[0][1].plot(time, P_disso, lw=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    axs[0][1].set_ylabel(r\"$P^{disso}$\")\n",
    "    axs[0][1].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "    \n",
    "    ## COH\n",
    "    count = -1\n",
    "    for ist in [1,2,0]:\n",
    "        for jst in [2,1,0]:\n",
    "            if ist>=jst: continue\n",
    "            count+= 1\n",
    "            \n",
    "            axs[count+1][1].yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "            \n",
    "            axs[count+1][1].set_ylabel(rF\"$\\langle | \\rho_{{{ist}{jst}}}|^{2} \\rangle $\")\n",
    "            \n",
    "            axs[count+1][1].margins(x=0)\n",
    "            axs[count+1][1].plot(time_DVR, coh_adi_DVR[:,ist, jst], linewidth=7, color=\"k\")\n",
    "            \n",
    "            for i, idir in enumerate(dirs):\n",
    "                axs[count+1][1].plot(time, MQC_coh[i][:,ist, jst], linewidth=5, color=colors[clrs_index[i]])\n",
    "    \n",
    "    #if(onlegend):\n",
    "    #    axs[1][1].legend(frameon=False, fontsize=10)\n",
    "    axs[-1][1].set_xlabel(\"time (fs)\")\n",
    "    \n",
    "    for i in range(4):\n",
    "        for j in range(2):\n",
    "            axs[i][j].set_xticks([0, 20, 40, 60, 80])\n",
    "    \n",
    "    fig.legend(loc='upper center', ncol=3, fontsize=30, frameon=False, bbox_to_anchor=(0.565,1.01))\n",
    "    fig.tight_layout()\n",
    "    fig.subplots_adjust(top=0.9)  # Move subplots down to make space for the legend\n",
    "    \n",
    "    plt.savefig(name_out, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35b124ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "dirs = [\"model2-method0-icond0-p15_0\", \"model2-method1-icond0-p15_0\", \n",
    "        \"model2-method2-icond0-p15_0-eps1_0\", \"model2-method4-icond0-p15_0\"]\n",
    "\n",
    "dirs = [\"phenol/TSH/\" + idir for idir in dirs]\n",
    "\n",
    "labels=[\"FSSH\", \"QTSH\", \"QTSH-SDM\", \"QTSH-XF\"]\n",
    "\n",
    "phenol_all(dirs, labels, title=\"Populations\", onlegend=True, name_out=\"fig8_phenol.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ba81939",
   "metadata": {},
   "source": [
    "## DVR movie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c81648b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_DVR_wp(_model_indx, _icond_indx, _k, _snap):\n",
    "    dir_name = F\"superexchange/DVR/wfc{_model_indx}-icond{_icond_indx}-p{_k:.1f}\".replace(\".\",\"_\")\n",
    "    fn = dir_name + F\"/wfcr_snap_{_snap}_dens_rep_1\"\n",
    "    which_cols = [0]\n",
    "    nstates = 3\n",
    "    for state in [x for x in range(nstates)]:\n",
    "        which_cols.append(1+state)\n",
    "    data_exact = data_read.get_data_from_file2(fn, which_cols)\n",
    "    ngrids = len(data_exact[0])\n",
    "    \n",
    "    #find necessary region\n",
    "    grids = np.array(data_exact[0])\n",
    "    # Get wfc array\n",
    "    wfcs = []\n",
    "    for ist in range(nstates):\n",
    "        wfcs.append(np.array(data_exact[1+ist]))\n",
    "    wfcs = np.array(wfcs)\n",
    "    \n",
    "    return grids, wfcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf43bcd3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "# Create the figure and axes\n",
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "plt.subplots_adjust(left=0.15, bottom=0.25)  # Adjust layout to make space for slider\n",
    "\n",
    "# Initial value of k\n",
    "k = 5.0\n",
    "snap_init = 0\n",
    "\n",
    "grids, wfcs = read_DVR_wp(1, 0, k, snap_init)\n",
    "lines = [ax.plot(grids, wfcs[ist], color=colors[clrs_index[ist]], label=f\"WP {ist}\")[0] for ist in range(3)]\n",
    "\n",
    "# Add a legend\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xlim([-50,50])   \n",
    "\n",
    "title = ax.text(0.02,0.90, \"\", fontsize=20, transform=ax.transAxes)\n",
    "\n",
    "# Update function for animation\n",
    "def update(frame):\n",
    "    _snap = frame\n",
    "    grids, wfcs = read_DVR_wp(0, 0, k, _snap)\n",
    "    \n",
    "    # Update each line in the plot with new data\n",
    "    for ist, line in enumerate(lines):\n",
    "        line.set_ydata(wfcs[ist])\n",
    "    \n",
    "    ax.set_xlabel(r\"$R$ (Bohr)\")\n",
    "    ax.set_ylabel(\"nuclear density\", x=3.1)\n",
    "    ax.set_title(r\"Superexchange, $\\hbar k_{0}=$ 5.0 a.u.\")\n",
    "    \n",
    "    title.set_text(rF\"$t=$ {_snap*100*10/41:.1f} fs\")\n",
    "    \n",
    "    ax.margins(x=0)\n",
    "    return lines, title\n",
    "\n",
    "# Create the animation\n",
    "ani = FuncAnimation(fig, update, frames=range(41), interval=200, blit=True)\n",
    "\n",
    "movie_name = F'DVR_superexchange_{k:.0f}'\n",
    "ani.save(movie_name.replace(\".\",\"_\") + \".gif\", fps=80, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cdf393d",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}