2.5_Wormhole_Experiments.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Make a Wormhole\n",
    "\n",
    "## Part 6: Wormhole Experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/freed_operator_and_state.png\">\n",
    "\n",
    "In our first example, we choose a random Hamiltonian and implement the state transfer protocol. We start with a message initialized to \\\\( \\begin{pmatrix} 1 \\\\ 0 \\end{pmatrix} \\\\), swap it into the left black hole, and then evolve. We then consider the expectation value of the Z operator on the corresponding qubit on the right side. Because of the tranpose, this should be -1 if all goes well (as opposed to +1). We still need a g value: a strength of the V coupling. So we try different g's to see what works and graph the expectation value of the Z operator against the value of g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RANDOM HAMILTONIAN MODEL\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip.qip.operations import swap\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def construct_tfd(E, beta=0):\n",
    "    El, Ev = E.eigenstates()\n",
    "    return (1/np.sqrt((-beta*E).expm().tr()))*\\\n",
    "             sum([np.exp(-(1/2)*beta*l)*\\\n",
    "                 qt.tensor(Ev[i].conj(), Ev[i])\\\n",
    "                    for i, l in enumerate(El)])\n",
    "\n",
    "def construct_coupling(n):\n",
    "    return\\\n",
    "    (1/(n-1))*sum([qt.tensor(*[qt.sigmaz()\\\n",
    "                       if j == i else qt.identity(2) \\\n",
    "                            for j in range(2*n)])*\n",
    "                   qt.tensor(*[qt.sigmaz()\\\n",
    "                        if j == i+n else qt.identity(2) \\\n",
    "                            for j in range(2*n)])\n",
    "                                for i in range(1, n)])\n",
    "\n",
    "##############################################################\n",
    "\n",
    "n = 7\n",
    "beta = 0\n",
    "\n",
    "IDn = qt.identity(2**n)\n",
    "IDn.dims = [[2]*n, [2]*n]\n",
    "\n",
    "E = qt.rand_herm(2**n)\n",
    "E.dims = IDn.dims\n",
    "tfd = construct_tfd(E, beta=beta)\n",
    "tfd.dims = [[2]*(2*n), [1]*(2*n)]\n",
    "\n",
    "msg = qt.basis(2,0)\n",
    "state = qt.tensor(msg, tfd)\n",
    "\n",
    "V = construct_coupling(n)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def worm(g, t, state, E, V, IDn):\n",
    "    state2 = qt.tensor(qt.identity(2), (1j*E*t).expm(), IDn)*state\n",
    "    state3 = swap(N=2*n+1, targets=[0,1])*state2\n",
    "    state4 = qt.tensor(qt.identity(2), (-1j*E*t).expm(), IDn)*state3\n",
    "    state5 = qt.tensor(qt.identity(2), (1j*g*V).expm())*state4\n",
    "    state6 = qt.tensor(qt.identity(2), IDn, (-1j*E.trans()*t).expm())*state5\n",
    "    return qt.expect(qt.sigmaz(), state6.ptrace(n+1))\n",
    "\n",
    "t = 10\n",
    "G = np.linspace(-10, 25, 40)\n",
    "Z = [worm(g, t, state, E, V, IDn) for g in G]\n",
    "\n",
    "best_g = G[np.argmin(np.array(Z))]\n",
    "print(best_g)\n",
    "\n",
    "plt.plot(G, Z, linewidth=2.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our second example, we use a specific Hamiltonian--actually two!: the kicked Ising model. Its time evolution can be described in two parts:\n",
    "\n",
    "$$ U = U_{K}U_{I} $$\n",
    "$$ U_{K} = e^{ib \\sum_{i} X_{i}} $$\n",
    "$$ U_{I} = e^{iJ\\sum_{i} Z_{i}Z_{i+1} + i \\sum_{i} h_{i}Z_{i}} $$\n",
    "   \n",
    "The model parameters are \\\\(J, b, \\\\) and \\\\(h_{i}\\\\). It turns out that if \\\\( J = b = \\frac{\\pi}{4} \\\\) and the \\\\( h_{i} \\\\) drawn uniformly from a Gaussian distribution, then the model is \"maximally chaotic\" if you look at how the entanglement entropy of subsystems grows (if you start with a product state). The first term is the kick, the second term is the Ising. \n",
    "\n",
    "Otherwise, the same as before. Compare the produced plot with that included in [arXiv:1911.06314](https://arxiv.org/abs/1911.06314)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# KICKED ISING MODEL\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip.qip.operations import swap\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def construct_tfd(E, beta=0):\n",
    "    El, Ev = E.eigenstates()\n",
    "    return (1/np.sqrt((-beta*E).expm().tr()))*\\\n",
    "             sum([np.exp(-(1/2)*beta*l)*\\\n",
    "                 qt.tensor(Ev[i].conj(), Ev[i])\\\n",
    "                    for i, l in enumerate(El)])\n",
    "\n",
    "def construct_coupling(n):\n",
    "    return\\\n",
    "    (1/(n-1))*sum([qt.tensor(*[qt.sigmaz()\\\n",
    "                       if j == i else qt.identity(2) \\\n",
    "                            for j in range(2*n)])*\n",
    "                   qt.tensor(*[qt.sigmaz()\\\n",
    "                        if j == i+n else qt.identity(2) \\\n",
    "                            for j in range(2*n)])\n",
    "                                for i in range(1, n)])\n",
    "\n",
    "##############################################################\n",
    "\n",
    "n = 7\n",
    "beta = 0\n",
    "\n",
    "IDn = qt.identity(2**n)\n",
    "IDn.dims = [[2]*n, [2]*n]\n",
    "\n",
    "##############################################################\n",
    "\n",
    "J = b = np.pi/4\n",
    "h = np.random.normal(scale=0.5, size=n)\n",
    "HK = b*sum([\\\n",
    "        qt.tensor(*[qt.sigmax()\\\n",
    "            if i == j else qt.identity(2)\\\n",
    "                for j in range(n)])\\\n",
    "                    for i in range(n)])\n",
    "HI = J*sum([\\\n",
    "        qt.tensor(*[qt.sigmaz()\\\n",
    "            if j == i else qt.identity(2)\\\n",
    "                for j in range(n)])*\\\n",
    "        qt.tensor(*[qt.sigmaz()\\\n",
    "            if j == i+1 else qt.identity(2)\\\n",
    "                for j in range(n)])\n",
    "                for i in range(n-1)])+\\\n",
    "       sum([h[i]*\\\n",
    "               qt.tensor(*[qt.sigmaz()\\\n",
    "                   if j == i else qt.identity(2)\\\n",
    "                       for j in range(n)])\\\n",
    "                           for i in range(n)])\n",
    "UK = (1j*HK).expm()\n",
    "UI = (1j*HI).expm()\n",
    "U = UK*UI\n",
    "\n",
    "##############################################################\n",
    "\n",
    "E = HI + HK\n",
    "tfd = construct_tfd(E, beta=beta)\n",
    "tfd.dims = [[2]*(2*n), [1]*(2*n)]\n",
    "\n",
    "msg = qt.basis(2,0)\n",
    "state = qt.tensor(msg, tfd)\n",
    "\n",
    "V = construct_coupling(n)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def worm(g, t, state, E, V, IDn):\n",
    "    state2 = qt.tensor(qt.identity(2), U.dag()**t, IDn)*state\n",
    "    state3 = swap(N=2*n+1, targets=[0,1])*state2\n",
    "    state4 = qt.tensor(qt.identity(2), U**t, IDn)*state3\n",
    "    state5 = qt.tensor(qt.identity(2), (1j*g*V).expm())*state4\n",
    "    state6 = qt.tensor(qt.identity(2), IDn, U.trans()**t)*state5\n",
    "    return qt.expect(qt.sigmaz(), state6.ptrace(n+1))\n",
    "\n",
    "t = 10\n",
    "G = np.linspace(-10, 25, 40)\n",
    "Z = [worm(g, t, state, E, V, IDn) for g in G]\n",
    "\n",
    "best_g = G[np.argmin(np.array(Z))]\n",
    "print(best_g)\n",
    "\n",
    "plt.plot(G, Z, linewidth=2.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that we can cement the analogy with quantum teleportation. The ZZ coupling is a classical coupling: it's actually diagonal! So we can replace it with a measurement and correction analogously to the teleportation case. We make Z measurements on all the \"carrier qubits\": in other words, the B qubits on the left side. For each we get an eigenvalue +1 or -1. We then act on the B qubits on the right with:\n",
    "\n",
    "$$ e^{ig\\sum_{i} z_{i}Z_{i}/(a-b)} $$\n",
    "\n",
    "In the above, i ranges over the B qubits on the right, and little z is the obtained eigenvalue. a is the number of A qubits on a side (in this case 1), b is the number of B qubits on a side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RANDOM HAMILTONIAN MODEL / MEASUREMENT\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip.qip.operations import swap\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def construct_tfd(E, beta=0):\n",
    "    El, Ev = E.eigenstates()\n",
    "    return (1/np.sqrt((-beta*E).expm().tr()))*\\\n",
    "             sum([np.exp(-(1/2)*beta*l)*\\\n",
    "                 qt.tensor(Ev[i].conj(), Ev[i])\\\n",
    "                    for i, l in enumerate(El)])\n",
    "\n",
    "def construct_coupling(n):\n",
    "    return\\\n",
    "    (1/(n-1))*sum([qt.tensor(*[qt.sigmaz()\\\n",
    "                       if j == i else qt.identity(2) \\\n",
    "                            for j in range(2*n)])*\n",
    "                   qt.tensor(*[qt.sigmaz()\\\n",
    "                        if j == i+n else qt.identity(2) \\\n",
    "                            for j in range(2*n)])\n",
    "                                for i in range(1, n)])\n",
    "\n",
    "def optimize_g(n, t, state, E, V, ID2, IDn):\n",
    "    def prep_worm(n, t, E, ID2, IDn):\n",
    "        U1 = qt.tensor(qt.identity(2), (1j*E*t).expm(), IDn)\n",
    "        U2 = swap(N=2*n+1, targets=[1,2])\n",
    "        U3 = qt.tensor(qt.identity(2), (-1j*E*t).expm(), IDn)\n",
    "        U5 = qt.tensor(qt.identity(2), IDn, (-1j*E.trans()*t).expm())\n",
    "        return [U5, U3*U2*U1]\n",
    "    def worm(n, g, V, state, ops, ID2):\n",
    "        A, B = ops\n",
    "        U4 = qt.tensor(qt.identity(2), (1j*g*V).expm())\n",
    "        state = A*U4*B*state\n",
    "        return qt.expect(qt.sigmaz(), state.ptrace(n+1))\n",
    "    G = np.linspace(0, 15, 60)\n",
    "    worm_prep = prep_worm(n, t, E, ID2, IDn)\n",
    "    Z = [worm(n, g, V, state, worm_prep, ID2) for g in G]\n",
    "    g = G[np.argmin(np.array(Z))]\n",
    "    return g\n",
    "\n",
    "##############################################################\n",
    "\n",
    "n = 4\n",
    "beta = 0\n",
    "\n",
    "IDn = qt.identity(2**n)\n",
    "IDn.dims = [[2]*n, [2]*n]\n",
    "\n",
    "E = qt.rand_herm(2**n)\n",
    "E.dims = IDn.dims\n",
    "tfd = construct_tfd(E, beta=beta)\n",
    "tfd.dims = [[2]*(2*n), [1]*(2*n)]\n",
    "\n",
    "msg = qt.basis(2,0)\n",
    "state = qt.tensor(msg, tfd)\n",
    "g = optimize_g(n, t, state, E, V, ID2, IDn)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def measure(op, i, state):\n",
    "    n = len(state.dims[0])\n",
    "    L, V = op.eigenstates()\n",
    "    rho = state.ptrace(i)\n",
    "    probabilities = np.array([(rho*V[j]*V[j].dag()).tr() for j in range(len(V))])\n",
    "    probabilities = probabilities/sum(probabilities)\n",
    "    choice = np.random.choice(list(range(len(L))), p=probabilities)\n",
    "    projector = qt.tensor(*[V[choice]*V[choice].dag()\\\n",
    "                            if j == i else qt.identity(2)\\\n",
    "                                for j in range(n)])\n",
    "    projector.dims = [state.dims[0], state.dims[0]]\n",
    "    state = (projector*state).unit()\n",
    "    return L[choice], state\n",
    "\n",
    "def worm(n, g, t, state, E, IDn):\n",
    "    state2 = qt.tensor(qt.identity(2), (1j*E*t).expm(), IDn)*state\n",
    "    state3 = swap(N=2*n+1, targets=[0,1])*state2\n",
    "    state4 = qt.tensor(qt.identity(2), (-1j*E*t).expm(), IDn)*state3\n",
    "    \n",
    "    #print(\"measuring left B qubits in the z-basis...\")\n",
    "    measurements = []\n",
    "    for i in range(2, n+1):\n",
    "        l, state4 = measure(qt.sigmaz(), i, state4)\n",
    "        measurements.append(l)\n",
    "        #print(\"got %.2f on qubit %d\" % (l, i))\n",
    "    #print(\"applying  unitary correction...\") \n",
    "    correction = (1j*g*sum([z*qt.tensor(*[qt.sigmaz()\\\n",
    "                              if j == 2+n+i else qt.identity(2)\\\n",
    "                              for j in range(2*n+1)])/(n-1)\\\n",
    "                                     for i, z in enumerate(measurements)])).expm()\n",
    "    state5 = (correction*state4).unit()\n",
    "    state6 = qt.tensor(qt.identity(2), IDn, (-1j*E.trans()*t).expm())*state5\n",
    "    return qt.expect(qt.sigmaz(), state6.ptrace(n+1))\n",
    "\n",
    "print(worm(n, g, t, state, E, IDn))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we consider evolving the wormwhole system continously in time, step by step through the state transfer protocol. First we optimize our g value. Then we construct the wormhole, this time using an entangled qubit as a reference. We track the mutual information between the reference qubit and those in the wormhole system instead of the Z expectation values. In a sense, using the entangled qubit trick allows us to defer our choice of which state we actually through into the wormhole. Using the \"jump\" method, we can project the reference state into some desired state, and see which state comes out the other side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TIME STEP RANDOM HAMILTONIAN MODEL\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip.qip.operations import swap\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def construct_tfd(E, beta=0):\n",
    "    El, Ev = E.eigenstates()\n",
    "    return (1/np.sqrt((-beta*E).expm().tr()))*\\\n",
    "             sum([np.exp(-(1/2)*beta*l)*\\\n",
    "                 qt.tensor(Ev[i].conj(), Ev[i])\\\n",
    "                    for i, l in enumerate(El)])\n",
    "\n",
    "def construct_coupling(n):\n",
    "    return\\\n",
    "    (1/(n-1))*sum([qt.tensor(*[qt.sigmaz()\\\n",
    "                       if j == i else qt.identity(2) \\\n",
    "                            for j in range(2*n)])*\n",
    "                   qt.tensor(*[qt.sigmaz()\\\n",
    "                        if j == i+n else qt.identity(2) \\\n",
    "                            for j in range(2*n)])\n",
    "                                for i in range(1, n)])\n",
    "\n",
    "def optimize_g(n, t, state, E, V, ID2, IDn):\n",
    "    def prep_worm(n, t, E, ID2, IDn):\n",
    "        U1 = qt.tensor(ID2, (1j*E*t).expm(), IDn)\n",
    "        U2 = swap(N=2*n+2, targets=[1,2])\n",
    "        U3 = qt.tensor(ID2, (-1j*E*t).expm(), IDn)\n",
    "        U5 = qt.tensor(ID2, IDn, (-1j*E.trans()*t).expm())\n",
    "        return [U5, U3*U2*U1]\n",
    "    def worm(n, g, V, state, ops, ID2):\n",
    "        A, B = ops\n",
    "        U4 = qt.tensor(ID2, (1j*g*V).expm())\n",
    "        state = A*U4*B*state\n",
    "        e = qt.entropy_mutual(state.ptrace((0, n+2)), 0, 1)\n",
    "        #print(\"g: %.3f e: %.10f\" % (g, e))\n",
    "        return e\n",
    "    G = np.linspace(0, 15, 60)\n",
    "    worm_prep = prep_worm(n, t, E, ID2, IDn)\n",
    "    ents = [worm(n, g, V, state, worm_prep, ID2) for g in G]\n",
    "    g = G[np.argmax(np.array(ents))]\n",
    "    return g\n",
    "\n",
    "##############################################################\n",
    "\n",
    "dt = 0.05\n",
    "t = 10\n",
    "n = 4\n",
    "beta = 0\n",
    "\n",
    "ID2 = qt.identity(4)\n",
    "ID2.dims = [[2,2], [2,2]]\n",
    "IDn = qt.identity(2**n)\n",
    "IDn.dims = [[2]*n, [2]*n]\n",
    "\n",
    "E = qt.rand_herm(2**n)\n",
    "E.dims = IDn.dims\n",
    "tfd = construct_tfd(E, beta=beta)\n",
    "tfd.dims = [[2]*(2*n), [1]*(2*n)]\n",
    "\n",
    "ref_msg = qt.bell_state(\"00\")\n",
    "state = qt.tensor(ref_msg, tfd)\n",
    "\n",
    "V = construct_coupling(n)\n",
    "g = optimize_g(n, t, state, E, V, ID2, IDn)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "entropies = [qt.entropy_mutual(\\\n",
    "                state.ptrace((0, i)), 0, 1)\\\n",
    "                     for i in range(1, 2*n+2)]\n",
    "entropies = entropies[:1+n] + entropies[1+n:][::-1]\n",
    "\n",
    "##############################################################\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_ylabel('entanglement')\n",
    "plt.ylim([0,0.5])\n",
    "plt.xticks(np.arange(2*n+1), [\"O\"]+\\\n",
    "                   [\"L%d\" % (i) for i in range(n)]+\\\n",
    "                   [\"R%d\" % (i) for i in range(n, 0, -1)])\n",
    "bar = ax.bar(list(range(2*n+1)), entropies)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def update(n, state, fig, bar):\n",
    "    entropies = [qt.entropy_mutual(\\\n",
    "                    state.ptrace((0, i)), 0, 1)\\\n",
    "                        for i in range(1, 2*n+2)]\n",
    "    entropies = entropies[:1+n] + entropies[1+n:][::-1]\n",
    "    for b, e in zip(bar, entropies):\n",
    "        b.set_height(e)\n",
    "    fig.canvas.draw()\n",
    "    plt.pause(0.00001)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "state = qt.tensor(ref_msg, tfd)\n",
    "\n",
    "print(\"evolving left back in time...\")\n",
    "U1 = qt.tensor(ID2, (1j*E*dt).expm(), IDn)\n",
    "for i in range(int(t/dt)):\n",
    "    state = U1*state\n",
    "    update(n, state, fig, bar)\n",
    "\n",
    "print(\"swapping in msg...\")\n",
    "state = swap(N=2*n+2, targets=[1,2])*state\n",
    "\n",
    "print(\"evolving left forward in time...\")\n",
    "U2 = qt.tensor(ID2, (-1j*E*dt).expm(), IDn)\n",
    "for i in range(int(t/dt)):\n",
    "    state = U2*state\n",
    "    update(n, state, fig, bar)\n",
    "    \n",
    "print(\"coupling...\")\n",
    "U3 = qt.tensor(ID2, (1j*V*dt).expm())\n",
    "for i in range(int(g/dt)):\n",
    "    state = U3*state\n",
    "    update(n, state, fig, bar)\n",
    "    \n",
    "print(\"evolving right forward in time\")\n",
    "U4 = qt.tensor(ID2, IDn, (-1j*E.trans()*dt).expm())\n",
    "for i in range(int(t/dt)):\n",
    "    state = U4*state\n",
    "    update(n, state, fig, bar)\n",
    "\n",
    "print(\"done!\")\n",
    "\n",
    "def jump(n, guy, state):\n",
    "    proj = qt.tensor(guy*guy.dag(), *[qt.identity(2) for i in range(2*n+1)])\n",
    "    return (proj*state).unit().ptrace(n+2)\n",
    "\n",
    "guy = qt.basis(2,0)\n",
    "guy2 = jump(n, guy, state)\n",
    "print(\"in:\")\n",
    "print(guy)\n",
    "print(\"out:\")\n",
    "print(guy2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we consider evolving the whole system continuously in time without steps, just under the following hamiltonian:\n",
    "\n",
    "$$ \\eta = E_{L} + E_{R}^{T} - gV $$\n",
    "$$ H = \\eta - \\langle TFD \\mid \\eta \\mid TFD \\rangle $$\n",
    "\n",
    "The point of the expectation value is that the TFD should be an approximate ground state of E."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONTINUOUS TIME RANDOM HAMILTONIAN MODEL\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import qutip as qt\n",
    "import numpy as np\n",
    "from qutip.qip.operations import swap\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def construct_tfd(E, beta=0):\n",
    "    El, Ev = E.eigenstates()\n",
    "    return (1/np.sqrt((-beta*E).expm().tr()))*\\\n",
    "             sum([np.exp(-(1/2)*beta*l)*\\\n",
    "                 qt.tensor(Ev[i].conj(), Ev[i])\\\n",
    "                    for i, l in enumerate(El)])\n",
    "\n",
    "def construct_coupling(n):\n",
    "    return\\\n",
    "    (1/(n-1))*sum([qt.tensor(*[qt.sigmaz()\\\n",
    "                       if j == i else qt.identity(2) \\\n",
    "                            for j in range(2*n)])*\n",
    "                   qt.tensor(*[qt.sigmaz()\\\n",
    "                        if j == i+n else qt.identity(2) \\\n",
    "                            for j in range(2*n)])\n",
    "                                for i in range(1, n)])\n",
    "\n",
    "def optimize_g(n, t, state, E, V, ID2, IDn):\n",
    "    def prep_worm(n, t, E, ID2, IDn):\n",
    "        U1 = qt.tensor(ID2, (1j*E*t).expm(), IDn)\n",
    "        U2 = swap(N=2*n+2, targets=[1,2])\n",
    "        U3 = qt.tensor(ID2, (-1j*E*t).expm(), IDn)\n",
    "        U5 = qt.tensor(ID2, IDn, (-1j*E.trans()*t).expm())\n",
    "        return [U5, U3*U2*U1]\n",
    "    def worm(n, g, V, state, ops, ID2):\n",
    "        A, B = ops\n",
    "        U4 = qt.tensor(ID2, (1j*g*V).expm())\n",
    "        state = A*U4*B*state\n",
    "        e = qt.entropy_mutual(state.ptrace((0, n+2)), 0, 1)\n",
    "        #print(\"g: %.3f e: %.10f\" % (g, e))\n",
    "        return e\n",
    "    G = np.linspace(0, 15, 60)\n",
    "    worm_prep = prep_worm(n, t, E, ID2, IDn)\n",
    "    ents = [worm(n, g, V, state, worm_prep, ID2) for g in G]\n",
    "    g = G[np.argmax(np.array(ents))]\n",
    "    return g\n",
    "\n",
    "##############################################################\n",
    "\n",
    "dt = 0.05\n",
    "t = 10\n",
    "n = 4\n",
    "beta = 0\n",
    "\n",
    "ID2 = qt.identity(4)\n",
    "ID2.dims = [[2,2], [2,2]]\n",
    "IDn = qt.identity(2**n)\n",
    "IDn.dims = [[2]*n, [2]*n]\n",
    "\n",
    "E = qt.rand_herm(2**n)\n",
    "E.dims = IDn.dims\n",
    "tfd = construct_tfd(E, beta=beta)\n",
    "tfd.dims = [[2]*(2*n), [1]*(2*n)]\n",
    "\n",
    "ref_msg = qt.bell_state(\"00\")\n",
    "state = qt.tensor(ref_msg, tfd)\n",
    "\n",
    "V = construct_coupling(n)\n",
    "g = optimize_g(n, t, state, E, V, ID2, IDn)\n",
    "\n",
    "ETA = qt.tensor(E, IDn) + qt.tensor(IDn, E) - g*V\n",
    "H = ETA - qt.expect(ETA, tfd)\n",
    "U = qt.tensor(ID2, (-1j*H*dt).expm())\n",
    "\n",
    "##############################################################\n",
    "\n",
    "entropies = [qt.entropy_mutual(\\\n",
    "                state.ptrace((0, i)), 0, 1)\\\n",
    "                     for i in range(1, 2*n+2)]\n",
    "entropies = entropies[:1+n] + entropies[1+n:][::-1]\n",
    "\n",
    "##############################################################\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_ylabel('entanglement')\n",
    "plt.ylim([0,0.5])\n",
    "plt.xticks(np.arange(2*n+1), [\"O\"]+\\\n",
    "                   [\"L%d\" % (i) for i in range(n)]+\\\n",
    "                   [\"R%d\" % (i) for i in range(n, 0, -1)])\n",
    "bar = ax.bar(list(range(2*n+1)), entropies)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "def update(n, state, fig, bar):\n",
    "    entropies = [qt.entropy_mutual(\\\n",
    "                    state.ptrace((0, i)), 0, 1)\\\n",
    "                        for i in range(1, 2*n+2)]\n",
    "    entropies = entropies[:1+n] + entropies[1+n:][::-1]\n",
    "    for b, e in zip(bar, entropies):\n",
    "        b.set_height(e)\n",
    "    fig.canvas.draw()\n",
    "    plt.pause(0.00001)\n",
    "\n",
    "##############################################################\n",
    "\n",
    "state = qt.tensor(ref_msg, tfd)\n",
    "state = swap(N=2*n+2, targets=[1,2])*state\n",
    "\n",
    "for i in range(1000):\n",
    "    state = U*state\n",
    "    update(n, state, fig, bar)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "VPython",
   "language": "python",
   "name": "vpython"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}