exercise_toycodes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "16d4aab6",
   "metadata": {},
   "source": [
    "# Exercises based on the toy codes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39301469",
   "metadata": {},
   "source": [
    "Uncomment and run the cell below (removing the `#`) when running this notebook on https://colab.research.google.com.\n",
    "\n",
    "Alternatively, you can run this notebook locally with jupyter, provided that you have the `*.py` python files from \n",
    "\n",
    "https://github.com/jhauschild/WinterschoolBarcelona2021\n",
    "\n",
    "in the same folder as your notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d42e939",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install git+git://github.com/jhauschild/WinterschoolBarcelona2021\n",
    "\n",
    "# use `pip uninstall tenpy-toycodes` to remove it a gain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d5c720",
   "metadata": {},
   "source": [
    "This tutorial focuses on a set of toy codes (using only python with numpy + scipy) that provide a simple implementation of the various MPS algorithms.\n",
    "\n",
    "You can add your code below by inserting additional cells as neccessary and running them (press Shift+Enter).\n",
    "\n",
    "**DISCLAIMER**: the toy codes used are not optimized, and we only use very small bond dimensions here. For state-of-the-art MPS calculations (especially for cylinders towards 2D), `chi` should be significantly larger, often on the order of several 1000s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd020922",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7d6aedd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.set_printoptions(precision=5, suppress=True, linewidth=100)\n",
    "plt.rcParams['figure.dpi'] = 150"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "337e0452",
   "metadata": {},
   "source": [
    "# toy codes: a_mps.py\n",
    "\n",
    "The file [a_mps.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/a_mps.py) defines a `SimpleMPS` class, that provides methods for expectation values and the entanglement entropy. You can initialize an inital product state MPS with the provided functions `init_FM_MPS` or `init_Neel_MPS`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b2cfaee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.a_mps import init_FM_MPS, init_Neel_MPS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f0ecc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_FM = init_FM_MPS(L=10, d=2, bc='finite')\n",
    "print(psi_FM)\n",
    "SigmaZ = np.diag([1., -1.])\n",
    "print(psi_FM.site_expectation_value(SigmaZ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b77ba98",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "Initialize a Neel state MPS. Print the `SigmaZ` expectation values.\n",
    "\n",
    "\n",
    "Print the entanglement entropy. What do you expect? Why do you get so many numbers, and not just one?\n",
    "Tipp: read the code ;-)\n",
    "\n",
    "Extract the half-chain entanglement entropy, i.e., the entropy when cutting the chain into two equal-length halves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667882c2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a255d8e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0966eb59",
   "metadata": {},
   "source": [
    "# toy codes: b_model.py\n",
    "\n",
    "The file [b_model.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/b_model.py) defines a `TFIModel` class representing the transverse field Ising model        \n",
    "$$H = - J \\sum_{i} \\sigma^z_i \\sigma^z_{i+1} - g \\sum_{i} \\sigma^x_i$$\n",
    "\n",
    "\n",
    "that provides both in the form of bond-terms `H_bonds` (as required for TEBD) and in the form of an MPO `H_mpo`.\n",
    "You can use `H_bonds` with `SimpleMPS.bond_expectation_values` to evalue the energy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d439f112",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.b_model import TFIModel\n",
    "\n",
    "model = TFIModel(L=10, J=1., g=1.2, bc='finite')\n",
    "\n",
    "print(\"<H_bonds> = \", psi_FM.bond_expectation_value(model.H_bonds))\n",
    "print(\"energy:\", np.sum(psi_FM.bond_expectation_value(model.H_bonds)))\n",
    "# (make sure the model and state have the same length and boundary conditions!)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "017d4d67",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Check the energies for the Neel state and make sure it matches what you expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0ef22db",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "bd218561",
   "metadata": {},
   "source": [
    "# toy codes: c_tebd.py\n",
    "\n",
    "The file [c_tebd.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/c_tebd.py) implements the TEBD algorithm, that we already learned about yesterday. \n",
    "\n",
    "\n",
    "**You don't have to read the code now; we will only use it for comparison with the other algorithms**.\n",
    "\n",
    "\n",
    "You can run TEBD with imaginary time evolutinon to find the ground state like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6c0d701",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.c_tebd import calc_U_bonds, run_TEBD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e8f7a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assuming you defined a `model`\n",
    "\n",
    "chi_max = 15\n",
    "\n",
    "psi = init_FM_MPS(model.L, model.d, model.bc)\n",
    "for dt in [0.1, 0.01, 0.001, 1.e-4, 1.e-5]:\n",
    "    U_bonds = calc_U_bonds(model.H_bonds, dt)\n",
    "    run_TEBD(psi, U_bonds, N_steps=100, chi_max=chi_max, eps=1.e-10)\n",
    "    E = np.sum(psi.bond_expectation_value(model.H_bonds))\n",
    "    print(\"dt = {dt:.5f}: E = {E:.13f}\".format(dt=dt, E=E))\n",
    "# the `run_TEBD` modified `psi`, so we can now calculate the expectation values from it\n",
    "print(\"final bond dimensions: \", psi.get_chi())\n",
    "mag_x = np.sum(psi.site_expectation_value(model.sigmax))\n",
    "mag_z = np.sum(psi.site_expectation_value(model.sigmaz))\n",
    "print(\"magnetization in X = {mag_x:.5f}\".format(mag_x=mag_x))\n",
    "print(\"magnetization in Z = {mag_z:.5f}\".format(mag_z=mag_z))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7b318b",
   "metadata": {},
   "source": [
    "# toy codes: d_dmrg.py\n",
    "\n",
    "The file [d_dmrg.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/d_dmrg.py) implements the DMRG algorithm.\n",
    "It can be called like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3196f562",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.d_dmrg import SimpleDMRGEngine, SimpleHeff2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffc67843",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_max = 15\n",
    "\n",
    "psi = init_FM_MPS(model.L, model.d, model.bc)\n",
    "eng = SimpleDMRGEngine(psi, model, chi_max=chi_max, eps=1.e-10)\n",
    "for i in range(10):\n",
    "    E_dmrg = eng.sweep()\n",
    "    E = np.sum(psi.bond_expectation_value(model.H_bonds))\n",
    "    print(\"sweep {i:2d}: E = {E:.13f}\".format(i=i + 1, E=E))\n",
    "print(\"final bond dimensions: \", psi.get_chi())\n",
    "mag_x = np.mean(psi.site_expectation_value(model.sigmax))\n",
    "mag_z = np.mean(psi.site_expectation_value(model.sigmaz))\n",
    "print(\"magnetization in X = {mag_x:.5f}\".format(mag_x=mag_x))\n",
    "print(\"magnetization in Z = {mag_z:.5f}\".format(mag_z=mag_z))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0eee6bf",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Read the code of [d_dmrg.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/d_dmrg.py) and try to understand the general structure of how it works.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbd12c50",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Compare running the code cells for DMRG and imaginary time evolution TEBD abovce for various parameters \n",
    "of `L`, `J`, `g`, and `bc`. Which one is faster? Do they always agree?\n",
    "\n",
    "Note: \n",
    "The transverse field Ising model has a quantum phase transition at $g/J = 1.$,\n",
    "from a ferro-magnetic phase (`g < J `) to a paramagnetic phase (`g > J`). We will map out the full phase diagram below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e7e445",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbf84fcc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8717f6a4",
   "metadata": {},
   "source": [
    "# Infinite DMRG\n",
    "\n",
    "The `SimpleDMRG` code also allows to run infinite DMRG, simply by replacing the `bc='finite'` for both the model and the MPS. Look at the implementation of `d_dmrg.py` to see where the differences are.\n",
    "\n",
    "The `L` parameter now just indices the number of tensors insite the unit cell of the infinite MPS.\n",
    "It has to be at least `2`, since we optimize 2 tensors at once in our DMRG code.\n",
    "Further, note that we now use the `mean` to calculate densities of observables instead of extensive quantities (like L)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76ab2cc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TFIModel(L=2, J=1., g=0.2, bc='infinite')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "717abeee",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_max = 10\n",
    "\n",
    "psi = init_FM_MPS(model.L, model.d, model.bc)\n",
    "eng = SimpleDMRGEngine(psi, model, chi_max=chi_max, eps=1.e-7)\n",
    "for i in range(10):\n",
    "    E_dmrg = eng.sweep()\n",
    "    E = np.mean(psi.bond_expectation_value(model.H_bonds))\n",
    "    print(\"sweep {i:2d}: E/L = {E:.13f}\".format(i=i + 1, E=E))\n",
    "print(\"final bond dimensions: \", psi.get_chi())\n",
    "mag_x = np.mean(psi.site_expectation_value(model.sigmax))\n",
    "mag_z = np.mean(psi.site_expectation_value(model.sigmaz))\n",
    "print(\"magnetization in X = {mag_x:.5f}\".format(mag_x=mag_x))\n",
    "print(\"magnetization in Z = {mag_z:.5f}\".format(mag_z=mag_z))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b260440",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Extend the following function `run_DMRG` to initialize an MPS compatible with the model, run DMRG on that MPS and return the corresponding state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b76969aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_DMRG(model, chi_max=50):\n",
    "    print(f\"runnning DMRG for L={model.L:d}, g={model.g:.2f}, bc={model.bc}, chi_max={chi_max:d}\")\n",
    "    \n",
    "    raise NotImplementedError(\"TODO: this is an Exercise!\")\n",
    "    \n",
    "    return psi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9b9796c",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Use that function to check that the average energy density and expectation values for infinite DMRG are independent of the unit cell length `L`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f7df7fc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8f7b39ac",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Use the `run_DMRG` function above to plot expectation values as a function of `g` for fixed `J=1` with the following code.\n",
    "\n",
    "\n",
    "Modify it to also plot the average expectation values of `sigmaz` and `sigmax` and the \n",
    "`sigmaz` correlation function between spin 0 and 20."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9a35a7b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#gs = [0.1, 0.5, 1.0, 1.5, 2.0]\n",
    "gs = np.linspace(0., 2., 21)\n",
    "entropies = []\n",
    "\n",
    "for g in gs:\n",
    "    model = TFIModel(L=2, J=1., g=g, bc='infinite')\n",
    "    psi = run_DMRG(model, chi_max=50)\n",
    "    entropies.append(np.max(psi.entanglement_entropy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ce49bc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(gs, entropies, marker='o', label='max entropy $S$')\n",
    "# plot expecation values of sigmax and sigmaz as well\n",
    "\n",
    "\n",
    "plt.xlabel('$g/J$')\n",
    "plt.ylabel('expectation values')\n",
    "plt.legend(loc='best')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17bee63e",
   "metadata": {},
   "source": [
    "# Defining a new model: the XX chain\n",
    "\n",
    "For the time evolution, we want to consider the XX chain in a staggered field, given by the Hamiltonian\n",
    "\n",
    "$$ H = \\sum_{i=0}^{N-2} (\\sigma^x_i \\sigma^x_{i+1} + \\sigma^y_i \\sigma^Y_{i+1})  - h_s \\sum_{i=0}^{N-1} (-1)^i \\sigma^z_i \n",
    "     = 2 \\sum_{i=0}^{N-2} (\\sigma^+_i \\sigma^-_{i+1} + \\sigma^+_i \\sigma^-_{i+1})  - h_s \\sum_{i=0}^{N-1} (-1)^i \\sigma^z_i\n",
    "$$\n",
    "\n",
    "for the usual Pauli matrices $\\sigma^z$, $\\sigma^x$.\n",
    "\n",
    "A Jordan-Wigner transformation maps the XX Chain to free fermions, which we can diagonalize exactly with a few lines of python codes that are given in [free_fermions_exact.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/free_fermions_exact.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a23c7236",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.free_fermions_exact import XX_model_ground_state_energy\n",
    "\n",
    "print(\"E_exact = \", XX_model_ground_state_energy(L=10, h_staggered=0.))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b8919e4",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "The following code implements the model for the XX Chain hamiltonian, but the MPO lacks some terms. Fill them in!\n",
    "\n",
    "\n",
    "Tip: In Python, `(-1)**i` represents $(-1)^i$.\n",
    "\n",
    "\n",
    "Compare the energies of DMRG and (imaginary) TEBD to check that you got it correct.\n",
    "**Note** that you have to use the Neel state as initial state to be in the right charge sector:\n",
    "the `|up up up ... up >` state already is an eigenstate of this Hamiltonian, so neither DMRG nor TEBD will modify it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f562a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class XXChain:\n",
    "    \"\"\"Simple class generating the Hamiltonian of the \n",
    "    The Hamiltonian reads\n",
    "    .. math ::\n",
    "        H = - J \\\\sum_{i} \\\\sigma^x_i \\\\sigma^x_{i+1} - g \\\\sum_{i} \\\\sigma^z_i\n",
    "    \"\"\"\n",
    "    def __init__(self, L, hs, bc='finite'):\n",
    "        assert bc in ['finite', 'infinite']\n",
    "        self.L, self.d, self.bc = L, 2, bc\n",
    "        self.hs = hs\n",
    "        self.sigmax = np.array([[0., 1.], [1., 0.]])  # Pauli X\n",
    "        self.sigmay = np.array([[0., -1j], [1j, 0.]]) # Pauli Y\n",
    "        self.sigmaz = np.array([[1., 0.], [0., -1.]]) # Pauli Z\n",
    "        self.id = np.eye(2)\n",
    "        self.init_H_bonds()\n",
    "        self.init_H_mpo()\n",
    "\n",
    "    def init_H_bonds(self):\n",
    "        \"\"\"Initialize `H_bonds` hamiltonian.\"\"\"\n",
    "        sx, sy, sz, id = self.sigmax, self.sigmay, self.sigmaz, self.id\n",
    "        d = self.d\n",
    "        nbonds = self.L - 1 if self.bc == 'finite' else self.L\n",
    "        H_list = []\n",
    "        for i in range(nbonds):\n",
    "            hL = hR = 0.5 * self.hs\n",
    "            if self.bc == 'finite':\n",
    "                if i == 0:\n",
    "                    hL = self.hs\n",
    "                if i + 1 == self.L - 1:\n",
    "                    hR = self.hs\n",
    "            H_bond = np.kron(sx, sx) + np.kron(sy, sy)\n",
    "            H_bond = H_bond - hL * (-1)**i * np.kron(sz, id) - hR * (-1)**(i+1) * np.kron(id, sz)\n",
    "            # H_bond has legs ``i, j, i*, j*``\n",
    "            H_list.append(np.reshape(H_bond, [d, d, d, d]))\n",
    "        self.H_bonds = H_list\n",
    "\n",
    "    # (note: not required for TEBD)\n",
    "    def init_H_mpo(self):\n",
    "        \"\"\"Initialize `H_mpo` Hamiltonian.\"\"\"\n",
    "        w_list = []\n",
    "        for i in range(self.L):\n",
    "            w = np.zeros((4, 4, self.d, self.d), dtype=complex)\n",
    "            w[0, 0] = w[3, 3] = self.id\n",
    "            \n",
    "            raise NotImplementedError(\"add further entries here\")\n",
    "            \n",
    "            w_list.append(w)\n",
    "        self.H_mpo = w_list\n",
    "        \n",
    "model = XXChain(9, 4., bc='finite')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5021c8b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"E_exact = \", XX_model_ground_state_energy(model.L, model.hs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd03222",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi_max = 100\n",
    "\n",
    "psi = init_Neel_MPS(model.L, model.d, model.bc)  # important: Neel\n",
    "for dt in [0.1, 0.01, 0.001, 1.e-4, 1.e-5]:\n",
    "    U_bonds = calc_U_bonds(model.H_bonds, dt)\n",
    "    run_TEBD(psi, U_bonds, N_steps=100, chi_max=chi_max, eps=1.e-10)\n",
    "    E = np.sum(psi.bond_expectation_value(model.H_bonds))\n",
    "    print(\"dt = {dt:.5f}: E = {E:.13f}\".format(dt=dt, E=E))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f5fc2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DMRG again\n",
    "chi_max = 100\n",
    "psi = init_Neel_MPS(model.L, model.d, model.bc)  # important: Neel!\n",
    "eng = SimpleDMRGEngine(psi, model, chi_max=chi_max, eps=1.e-10)\n",
    "for i in range(10):\n",
    "    eng.sweep()\n",
    "    E = np.sum(psi.bond_expectation_value(model.H_bonds))\n",
    "    #print(psi.get_chi())\n",
    "    print(\"sweep {i:2d}: E = {E:.13f}\".format(i=i + 1, E=E))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42fe4054",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4969359f",
   "metadata": {},
   "source": [
    "# Time evolution of the Neel state under the XX chain\n",
    "\n",
    "We will now consider the Neel state as an initial state for (real-)time evolution, i.e., calculate\n",
    "$$ \\vert \\psi(t)\\rangle =e^{-i H t} \\vert\\uparrow \\downarrow \\uparrow \\downarrow \\dots\\rangle$$\n",
    "under the Hamiltonian of the XX chain (with staggered field).\n",
    "\n",
    "Again, we can use the mapping to free fermions to calculate some quantities exactly.\n",
    "We will focus on the half chain entanglement entropy and compare it to TEBD:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d9fc268",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.free_fermions_exact import XX_model_time_evolved_entropies\n",
    "\n",
    "dt = 0.1\n",
    "tmax = 4.\n",
    "L = 50\n",
    "hs = 0.\n",
    "model = XXChain(L, hs)\n",
    "\n",
    "times_exact = np.arange(0., tmax, dt)\n",
    "S_exact = XX_model_time_evolved_entropies(L=L, h_staggered=hs, time_list=times_exact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "777211cf",
   "metadata": {},
   "source": [
    "For reference, will will also compare to a TEBD run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12f349e4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "chi_max = 30\n",
    "\n",
    "psi = init_Neel_MPS(L, 2, bc='finite')\n",
    "U_bonds = calc_U_bonds(model.H_bonds, 1.j * dt) # here, imaginary dt = real time evolution\n",
    "t = 0.\n",
    "times_tebd = []\n",
    "S_tebd = []\n",
    "while t < tmax: \n",
    "    times_tebd.append(t)\n",
    "    S_tebd.append(psi.entanglement_entropy()[(L-1)//2])\n",
    "    print(f\"t={t:.2f}\")\n",
    "    t += dt\n",
    "    run_TEBD(psi, U_bonds, N_steps=1, chi_max=chi_max, eps=1.e-7)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e3a0072",
   "metadata": {},
   "source": [
    "# toy code: e_tdvp.py\n",
    "\n",
    "The file [e_tdvp.py](https://github.com/jhauschild/WinterschoolBarcelona2021/blob/main/toycodes/e_tdvp.py) implements the TDVP time evolution for finite MPS with the `SimpleTDVPEngine`.\n",
    "It implements both the true TDVP with one-site updates in `sweep_one_site`, and `sweep_two_site` that allows to grow the bond dimension:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7814c1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from toycodes.e_tdvp import SimpleTDVPEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c6f5f2e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "chi_max = 30\n",
    "t_switch = 1.\n",
    "\n",
    "psi = init_Neel_MPS(L, 2, bc='finite')\n",
    "eng = SimpleTDVPEngine(psi, model, chi_max=chi_max, eps=1.e-7)\n",
    "\n",
    "t = 0.\n",
    "times_tdvp = []\n",
    "S_tdvp = []\n",
    "while t < tmax: \n",
    "    times_tdvp.append(t)\n",
    "    S_tdvp.append(psi.entanglement_entropy()[(L-1)//2])\n",
    "    print(f\"t={t:.2f}, chi_max = {max(psi.get_chi()):d}\")\n",
    "    t += dt\n",
    "    if t < t_switch:\n",
    "        eng.sweep_two_site(dt)\n",
    "    else:\n",
    "        eng.sweep_one_site(dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8139c6a8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667f930e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(times_exact, S_exact, label=\"exact\")\n",
    "plt.plot(times_tebd, S_tebd, label=\"TEBD\")\n",
    "plt.plot(times_tdvp, S_tdvp, label=\"TDVP\")\n",
    "\n",
    "plt.xlabel('time $t$')\n",
    "plt.ylabel('half-chain entropy $S$')\n",
    "plt.legend(loc='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a83ffdd6",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Read the code of `e_tdvp.py` and try to see the general structure of how it works."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9553b647",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Add curves for smaller (or larger) chi to the plot above. (It might take a little while to run for large chi...).\n",
    "\n",
    "Try switching earlier or later from two-site to one-site TDVP.\n",
    "\n",
    "What happens if you switch on the staggered field?\n",
    "\n",
    "\n",
    "**NOTE**: You can avoid having to rerun the time evolution if you make copies of the results `S_tdvp` and similar lists.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad4d0c74",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0989680",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d16df3aa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1e99ab49",
   "metadata": {},
   "source": [
    "### Advanced exercises - if you're an expert and have time left ;-)\n",
    "\n",
    "- Obtain the ground state of the transverse field ising model at the critical point with DMRG.\n",
    "  Try to plot the corrlation function as a function of `j-i`.\n",
    "  What form does it have? Is an MPS a good ansatz for that?\n",
    "- Calling `SimpleMPS.correlation_function` in a loop over `j` for fixed `i` is inefficient for large `j-i`. \n",
    "  Try to rewrite the correlation function to obtain the values for all `j>i` up to some cutoff in a single loop.\n",
    "- Look at the light-cone after a local quench of the XX chain, similarly as is done in  `e_tdvp.example_TDVP_tf_ising_lightcone`. \n",
    "  How does the speed of the light cone change with the staggered field?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f2001a",
   "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}