Tutorial-2D-Heisenberg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc47a3c8-b67d-4977-86f3-f12548972e6a",
   "metadata": {},
   "source": [
    "# Predicting Properties of Quantum Systems with Conditional Generative Models\n",
    "\n",
    "Tutorial by Maurice Weber, Xanadu Summer Resident, email: maurice.weber@inf.ethz.ch, Nov 2022"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ec5944-b2a0-4757-b8fb-058aa3cfe8b4",
   "metadata": {},
   "source": [
    "## Problem Setup\n",
    "In this tutorial, we learn how conditional generative models can be used to predict properties of families of quantum states. As a specific example, we will consider the ground states of the 2D anti-ferromagnetic random Heisenberg model with Hamiltonian\n",
    "\n",
    "$$\n",
    " H(\\mathbf{x}) = \\sum_{\\langle ij \\rangle} \\mathbf{x}_{ij} (X_i X_j + Y_i Y_j + Z_i Z_j),\n",
    "$$\n",
    "\n",
    "where the sum is over nearest neighbours in a 2D lattice and $\\mathbf{x} = (\\mathbf{x}_{ij})_{i,j=1}^n$ is a coupling matrix containing the coupling strengths between different sites. For a POVM $\\mathcal{M}$, the goal is to train a generative model $p_\\theta$ such that\n",
    "\n",
    "$$\n",
    "    p_\\theta(a_1,\\,\\ldots,\\,a_n~\\lvert~\\mathbf{x}) = \\langle \\psi_{\\mathbf x}\\rvert M_{\\vec a}\\lvert \\psi_{\\mathbf x} \\rangle ~~~~ M_{\\vec a} \\in \\mathcal{M},\n",
    "$$\n",
    "\n",
    "where $\\lvert \\psi_{\\mathbf x} \\rangle$ denotes the ground state of the Hamiltonian $H(\\mathbf x)$, and the vector $\\vec a = (a_1,~\\ldots,~a_n)^T$ contains the measurement outcomes. Here we consider the Pauli-6 POVM, which is given by the set of operators\n",
    "\n",
    "$$\n",
    "    \\mathcal{M}_{\\text{Pauli-6}}^{\\otimes n} = \\left\\{\n",
    "        \\frac{1}{3}\\lvert{+}\\rangle\\langle{+}\\rvert,\\,\n",
    "        \\frac{1}{3}\\lvert{-}\\rangle\\langle{-}\\rvert,\\,\n",
    "        \\frac{1}{3}\\lvert{l}\\rangle\\langle{l}\\rvert,\\,\n",
    "        \\frac{1}{3}\\lvert{r}\\rangle\\langle{r}\\rvert,\\,\n",
    "        \\frac{1}{3}\\lvert{0}\\rangle\\langle{0}\\rvert,\\,\n",
    "        \\frac{1}{3}\\lvert{1}\\rangle\\langle{1}\\rvert\n",
    "    \\right\\}^{\\otimes n}\n",
    "$$\n",
    "\n",
    "so that measurement outcomes take values $\\vec a \\in \\{+,~-,~l,~r,~0,~1\\}^n$. To predict properties of the ground state $\\lvert \\psi_{\\mathbf x} \\rangle$, we make use of classical shadows. In the following we briefly outline how classical shadows can be used to estimate local observables and subsystem entanglement entropies.\n",
    "\n",
    "### Post-Processing using Classical Shadows\n",
    "\n",
    "Classical shadows provide a way to estimate properties of a quantum state $\\rho$, given a collection of Pauli-6 measurements.\n",
    "Given $T$ measurement shots for the Pauli-6 POVM, the classical shadow of an arbitrary state $\\rho$ is given by the collection of snapshots\n",
    "\n",
    "$$\n",
    "    \\texttt{S}(\\rho;~T) = \\left\\{\\hat\\rho^{(1)}:=\\bigotimes_{i=1}^n\\left(3 \\lvert a_i^{(1)}\\rangle\\langle a_i^{(1)}\\rvert - \\mathbb{1}_2\\right),~\\ldots,~\\hat\\rho^{(T)}:=\\bigotimes_{i=1}^n\\left(3 \\lvert a_i^{(T)}\\rangle\\langle a_i^{(T)}\\rvert - \\mathbb{1}_2\\right)\\right\\}.\n",
    "$$\n",
    "\n",
    "This collection allows one to reconstruct the density matrix of $\\rho$ in expectation\n",
    "\n",
    "$$\n",
    "    \\hat\\rho_T := \\frac{1}{T}\\sum_{t=1}^T \\bigotimes_{i=1}^n\\left(3 \\lvert a_i^{(t)}\\rangle\\langle a_i^{(t)}\\rvert - \\mathbb{1}_2\\right) \\xrightarrow{T\\to\\infty} \\mathbb{E}_{\\vec a}\\left[\\bigotimes_{i=1}^n\\left(3 \\lvert a_i^{(t)}\\rangle\\langle a_i^{(t)}\\rvert - \\mathbb{1}_2\\right)\\right] = \\rho\n",
    "$$\n",
    "\n",
    "or directly estimate obervables. Specifically, to estimate Pauli observables $P = P_1\\otimes\\ldots\\otimes P_n$, with $P_i\\in\\{\\mathbb{1},\\,X,\\,Y,\\,Z\\}$, we use the identity\n",
    "\n",
    "$$\n",
    "    \\mathrm{Tr}\\left[P\\hat\\rho\\right] = \\frac{1}{T}\\sum_{t=1}^T \\prod_{i=1}^n \\left(3\\langle a_i^{(t)}\\rvert P_j\\lvert a_i^{(t)}\\rangle - \\mathrm{Tr}[P_i]\\right).\n",
    "$$\n",
    "\n",
    "Another important property that we will estimate using classical shadows is the second-order Rényi subsystem entanglement entropy $\\mathcal{S}(\\rho_A) = -\\ln \\mathrm{Tr}[\\rho_A^2]$ where $\\rho_A$ is the reduced density matrix for subsystem $A$. For two sites $i\\neq j$, we can approximate $\\mathcal{S}(\\rho_A)$ using the identity\n",
    "\n",
    "$$\n",
    "    \\mathrm{Tr}\\left[\\hat\\rho_{\\{ij\\}}^2\\right] = \\frac{1}{T^2}\\sum_{t,~s=1}^T \\left(9\\lvert\\langle a_i^{(t)}|a_i^{(s)}\\rangle\\rvert^2 - 4\\right)\\cdot\\left(9\\lvert\\langle a_j^{(t)}|a_j^{(s)}\\rangle\\rvert^2 - 4\\right).\n",
    "$$\n",
    "\n",
    "For further details about classical shadows, visit the two Pennylane tutorials: \n",
    "- [Tutorial on Classical Shadows](https://pennylane.ai/qml/demos/tutorial_classical_shadows.html): A Python re-implementation of Hsin-Yuan Huang's [_classical shadows code (C++)_](https://github.com/hsinyuan-huang/predicting-quantum-properties),\n",
    "- [Machine learning for quantum many-body problems](https://pennylane.ai/qml/demos/tutorial_ml_classical_shadows.html): A Python implementation of [_Provably efficient machine learning for quantum many-body problems_](https://arxiv.org/abs/2106.12627)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff7cb758-246f-40de-b496-ff4299f65f39",
   "metadata": {},
   "source": [
    "# A closer look at the 2D Heisenberg Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a3f46d8-3461-4d41-b728-e48d884cad8a",
   "metadata": {},
   "source": [
    "Here, we want to investigate the ground states of the 2D random Heisenber model. Specifically, here we compute 2-point correlation functions and subsystem entanglement entropies. These evaluations will later serve us a point of reference for the predictions made by our conditional generative model. Let us start by making some basic imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7034a3d7-c14c-440a-966b-80b50480a740",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools as it\n",
    "from matplotlib import gridspec\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import pennylane as qml\n",
    "import qutip\n",
    "import scipy as sp\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e92d1494-8937-4079-8016-79dabccc7ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f329ebcc-7c3d-4726-b5cd-55ea82de061c",
   "metadata": {},
   "source": [
    "We first setup a function which generates a 2D lattice connection weights sampled uniformly at random from the interval $[0,~2]$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75a7ad15-d7f1-4323-b2ca-d66b9b6c60ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_coupling_matrix(rows, cols):\n",
    "    qubits = rows * cols\n",
    "    \n",
    "    # Create a 2D Lattice\n",
    "    edges = [\n",
    "        (si, sj) for (si, sj) in it.combinations(range(qubits), 2)\n",
    "        if ((sj % cols > 0) and sj - si == 1) or sj - si == cols\n",
    "    ]\n",
    "    \n",
    "    # sample edge weights uniformly at random from [0, 2]\n",
    "    edge_weights = rng.uniform(0, 2, size=len(edges))\n",
    "    \n",
    "    coupling_matrix = np.zeros((qubits, qubits))\n",
    "    for (i, j), w in zip(edges, edge_weights):\n",
    "        coupling_matrix[i, j] = coupling_matrix[j, i] = w\n",
    "        \n",
    "    return coupling_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "179aead5-9b28-438d-86fe-dc4e142b3b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the system size and lattice geometry\n",
    "rows, cols = 4, 4\n",
    "wires = rows * cols\n",
    "\n",
    "# sample a coupling matrix\n",
    "J = sample_coupling_matrix(rows, cols)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecb0133d-db14-4369-a12c-605f56071446",
   "metadata": {},
   "source": [
    "To get an intuition about how our system looks like, we visualize the coupling graph corresponding to our Hamiltonian usig the `networkx` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c4e3c3d-ef37-49c5-a940-5d6d9bd4b8be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create graph object\n",
    "graph = nx.from_numpy_matrix(np.matrix(J), create_using=nx.DiGraph)\n",
    "graph = nx.relabel_nodes(graph, {i: i + 1 for i in graph.nodes})\n",
    "pos = {i: ((i-1) % cols, -((i-1) // cols)) for i in graph.nodes()}\n",
    "\n",
    "# make edge widths proportional to edge weight\n",
    "edge_widths = [\n",
    "    (x + 1.5) ** 2 for x in list(nx.get_edge_attributes(graph, \"weight\").values())\n",
    "]\n",
    "\n",
    "# extract edge weights for colouring\n",
    "edges, weights = zip(*nx.get_edge_attributes(graph,'weight').items())\n",
    "\n",
    "plt.figure(figsize=(cols / 1.5, rows / 1.5))\n",
    "nx.draw(\n",
    "    graph, pos, node_color=\"white\", with_labels=True, font_color=\"black\", edge_cmap=plt.cm.Blues,\n",
    "    node_size=400, width=edge_widths, horizontalalignment='center', edgecolors=\"black\", edgelist=edges, \n",
    "    edge_color=weights, arrows=False, verticalalignment='center_baseline', font_size=12\n",
    ")\n",
    "plt.title('Coupling Graph', fontsize=18)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38aafd16-e4d5-4458-b50b-da010919a9bf",
   "metadata": {},
   "source": [
    "Given the coupling graph, we can now setup the corresponding Hamiltonian, compute its ground state using exact diagonalization and compute different properties of the system. Let us first write a function which sets up the Hamiltonian. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3748084-ddc2-45d7-8b8b-89a9f233681b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_hamiltonian(coupling_matrix):\n",
    "    coeffs, ops = [], []\n",
    "    ns = coupling_matrix.shape[0]\n",
    "\n",
    "    for i, j in it.combinations(range(ns), r=2):\n",
    "        coeff = coupling_matrix[i, j]\n",
    "        if coeff:\n",
    "            for op in [qml.PauliX, qml.PauliY, qml.PauliZ]:\n",
    "                coeffs.append(coeff)\n",
    "                ops.append(op(i) @ op(j))\n",
    "\n",
    "    return qml.Hamiltonian(coeffs, ops)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298abe65-d1ba-4645-8050-5ff8a1a0c727",
   "metadata": {},
   "source": [
    "Next, we will compute its ground state with exact diagonalization and using a sparse representation of the Hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e25b464-29ee-4737-83e1-0144f81b20cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build sparse hamiltonian\n",
    "H = build_hamiltonian(J)\n",
    "H_sparse = qml.utils.sparse_hamiltonian(H)\n",
    "\n",
    "# diagonalize\n",
    "eigvals, eigvecs = sp.sparse.linalg.eigs(H_sparse, which='SR', k=1)\n",
    "eigvals = eigvals.real\n",
    "ground_state = eigvecs[:, np.argmin(eigvals)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0a4c2a4-c305-4f98-ba9f-64f2a76c79e0",
   "metadata": {},
   "source": [
    "As hinted at earlier, throughout this demo, we are mainly interested in two properties of the ground state. Specifically, these are the two point correlation function and the second-order Rényi subsystem entanglement entropy for subsystems of size at most two. The correlation function of the ground state $\\lvert\\psi_\\mathbf{J}\\rangle$ is given by\n",
    "\n",
    "$$\n",
    "    (i,~j) \\mapsto C_{ij} = \\frac{1}{3}\\langle\\psi_\\mathbf{x}\\rvert \\left(X_i X_j + Y_i Y_j + Z_i Z_j\\right)\\lvert\\psi_\\mathbf{x}\\rangle,\n",
    "$$\n",
    "\n",
    "and the entanglement entropy for subsystem $\\{ij\\}$ is\n",
    "\n",
    "$$\n",
    "    \\mathcal{S}_{i,j} = -\\ln \\mathrm{Tr}\\left[\\rho_{\\{ij\\}}^2\\right].\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18960f78-fdbf-429a-bd6f-f543621dade4",
   "metadata": {},
   "source": [
    "### Computing the exact two-point correlation function\n",
    "\n",
    "First, we write a function which computes the exact two-point correlation function for a given statevector. We can then use this function to compute the correlations for the ground state we have computed before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfdfedb4-5869-4258-88f5-2bdcddf5c4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this circuit measures observables for the provided ground state\n",
    "@qml.qnode(device=qml.device('default.qubit', wires=wires, shots=None))\n",
    "def circ(observables):\n",
    "    qml.QubitStateVector(ground_state, wires=range(wires))\n",
    "    return [qml.expval(o) for o in observables]\n",
    "\n",
    "\n",
    "def compute_exact_correlation_matrix(ground_state, wires):\n",
    "    # setup observables for correlation function\n",
    "    def corr_function(i, j):\n",
    "        ops = []\n",
    "        \n",
    "        for op in [qml.PauliX, qml.PauliY, qml.PauliZ]:\n",
    "            if i != j:\n",
    "                ops.append(op(i) @ op(j))\n",
    "            else:\n",
    "                ops.append(qml.Identity(i))\n",
    "\n",
    "        return ops\n",
    "    \n",
    "    # indices for sites for which correlations will be computed\n",
    "    coupling_pairs = list(it.product(range(wires), repeat=2))\n",
    "    \n",
    "    # compute exact correlation matrix\n",
    "    correlation_matrix = np.zeros((wires, wires))\n",
    "    for idx, (i, j) in tqdm(enumerate(coupling_pairs), total=len(coupling_pairs)):\n",
    "        observable = corr_function(i, j)\n",
    "\n",
    "        if i == j:\n",
    "            correlation_matrix[i][j] = 1.0\n",
    "        else:\n",
    "            correlation_matrix[i][j] = (\n",
    "                    np.sum(np.array([circ(observables=[o]) for o in observable]).T) / 3\n",
    "            )\n",
    "            correlation_matrix[j][i] = correlation_matrix[i][j]\n",
    "\n",
    "    return correlation_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52c7f345-4f88-49fb-9ada-5c0167870e3f",
   "metadata": {},
   "source": [
    "This lets us easily compute the correlation function for the ground state of the above system, by calling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388e4aa6-35dc-4f15-b875-9f9e5fc0c6c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "exact_correlation_matrix = compute_exact_correlation_matrix(ground_state, wires)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f0615e-4a8b-4f73-9984-0715c71e253a",
   "metadata": {},
   "source": [
    "To visualize this nicely, let us write a function which displays the the coupling matrix next to the computed correlation matrix. We can use this function later again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1070d4-4b06-4e9b-bdd2-17ea7adb2832",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_eval_matrix(coup_mat, eval_mat, title, cmap, true_mat=None, vmin=-1.0, vmax=1.0):\n",
    "    num_figures = 2 if true_mat is None else 3 \n",
    "    fig = plt.figure(figsize=(num_figures * cols + 2, rows))\n",
    "    gs = gridspec.GridSpec(rows, num_figures * cols, figure=fig)\n",
    "    \n",
    "    # plot coupling graph\n",
    "    ax = fig.add_subplot(gs[:, :cols])\n",
    "    graph = nx.from_numpy_matrix(np.matrix(coup_mat), create_using=nx.DiGraph)\n",
    "    graph = nx.relabel_nodes(graph, {i: i + 1 for i in graph.nodes})\n",
    "    pos = {i: ((i-1) % cols, -((i-1) // cols)) for i in graph.nodes()}\n",
    "    edge_widths = [(x + 1.5) ** 2 for x in list(nx.get_edge_attributes(graph, \"weight\").values())]\n",
    "    edges, weights = zip(*nx.get_edge_attributes(graph,'weight').items())\n",
    "    ax.set_title('Coupling Graph', fontsize=16)\n",
    "    ax.set_aspect('equal')\n",
    "\n",
    "    nx.draw(\n",
    "        graph, pos, node_color=\"white\", with_labels=True, font_color=\"black\", edge_cmap=plt.cm.Blues,\n",
    "        node_size=400, width=edge_widths, horizontalalignment='center', edgecolors=\"black\", edgelist=edges, edge_color=weights,\n",
    "        arrows=False, ax=ax, verticalalignment='center_baseline', font_size=10\n",
    "    )\n",
    "    \n",
    "    # plot correlation matrix\n",
    "    tick_locs = np.array([0] + [i for i in np.arange(5, rows * cols, 5) - 1])\n",
    "    tick_marks = tick_locs + 1\n",
    "    if true_mat is None:\n",
    "        ax = fig.add_subplot(gs[:, cols:])\n",
    "    else:\n",
    "        ax = fig.add_subplot(gs[:, cols:(2 * cols)])\n",
    "    im = ax.imshow(eval_mat, cmap=plt.cm.get_cmap(cmap), vmin=vmin, vmax=vmax)\n",
    "    ax = plt.gca()\n",
    "    ax.xaxis.tick_bottom()\n",
    "    ax.yaxis.tick_left()\n",
    "    ax.xaxis.set_ticks(tick_locs, tick_marks, fontsize=12)\n",
    "    ax.yaxis.set_ticks(tick_locs, tick_marks, fontsize=12)\n",
    "    ax.set_title(title, fontsize=16)\n",
    "    ax.set_aspect('equal')\n",
    "\n",
    "    for _, s in ax.spines.items():\n",
    "        s.set_color('white')\n",
    "    \n",
    "    if true_mat is not None:\n",
    "        ax = fig.add_subplot(gs[:, (2 * cols):])\n",
    "        im = ax.imshow(true_mat, cmap=plt.cm.get_cmap(cmap), vmin=vmin, vmax=vmax)\n",
    "        ax = plt.gca()\n",
    "        ax.xaxis.tick_bottom()\n",
    "        ax.yaxis.tick_left()\n",
    "        ax.xaxis.set_ticks(tick_locs, tick_marks, fontsize=12)\n",
    "        ax.yaxis.set_ticks(tick_locs, tick_marks, fontsize=12)\n",
    "        ax.set_title('Ground Truth', fontsize=16)\n",
    "        ax.set_aspect('equal')\n",
    "        \n",
    "        for _, s in ax.spines.items():\n",
    "            s.set_color('white')\n",
    "\n",
    "    # colorbar\n",
    "    bar = fig.colorbar(im, pad=0.01, ax=ax)\n",
    "    bar.set_label(r'$C_{ij}$', fontsize=12, rotation=0, labelpad=10)\n",
    "    bar.ax.tick_params(labelsize=12)\n",
    "\n",
    "    for _, s in bar.ax.spines.items():\n",
    "        s.set_color('white')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c4f761-b77a-4918-86ee-df330c91d83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_eval_matrix(J, exact_correlation_matrix, title='Correlation Function', cmap='RdBu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3041fed6-38e4-4167-b6f9-f348feca963d",
   "metadata": {},
   "source": [
    "### Computing the exact entanglement entropies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f4a2b4-4c2e-4150-bbda-50d6a1a33d55",
   "metadata": {},
   "source": [
    "Similar to the correlation functions, we can also compute subsystem entanglement entropies exactly, given the ground state above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b75417d-ef94-4928-b46c-5a5be354e1be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_exact_entropy_matrix(ground_state, wires):\n",
    "    ground_state_qobj = qutip.Qobj(ground_state, dims=[[2] * wires, [1] * wires])\n",
    "\n",
    "    # compute entropies\n",
    "    entropies = np.zeros(shape=(wires, wires), dtype=float)\n",
    "    for i in tqdm(range(wires)):\n",
    "        ptrace_diag = ground_state_qobj.ptrace(sel=[i])\n",
    "        entropies[i, i] = -np.log(np.trace(ptrace_diag * ptrace_diag).real)\n",
    "\n",
    "        for j in range(i + 1, wires):\n",
    "            ptrace = ground_state_qobj.ptrace(sel=[i, j])\n",
    "            e = -np.log(np.trace(ptrace * ptrace).real)\n",
    "            entropies[i, j] = entropies[j, i] = e\n",
    "\n",
    "    return entropies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2da3e1-0410-420e-8028-26c8737e6810",
   "metadata": {},
   "outputs": [],
   "source": [
    "exact_entropy_matrix = compute_exact_entropy_matrix(ground_state, wires)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0253710c-a8a9-48c3-8707-f4d81abc5cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_eval_matrix(J, exact_entropy_matrix, title='Subsystem Entanglement Entropies', cmap='Blues', vmin=0.0, vmax=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a12e8fa-b3fd-4bc3-a1d5-5f7a04dd217c",
   "metadata": {},
   "source": [
    "# Making predictions using a conditional transformer model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ddb54c-12ca-4228-8737-dbc5c9bf8ea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import jit\n",
    "import jax.numpy as jnp\n",
    "\n",
    "from heisenberg_tutorial.transformer import initialize, preprocess_coupling_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81cd70e-ba33-4c7d-bb21-bde2b3cc8ac5",
   "metadata": {},
   "source": [
    "Here, we use a conditional transformer model to predict 2-point correlations and subsystem entanglement entropies for the ground state of the 2D anti-ferromagnetic Heisenberg model. Our model consists of an encoder $g_\\phi\\colon \\mathcal{G} \\to \\mathbb{R}^d$, that encodes the coupling graph into a latent representation $G \\mapsto g_\\phi(G)$. The encoding $g_\\phi$ is parametrized as a graph convolutional neural network (GCNN), a particular type of neural network that is capable of processing graph-structured data. This representation is then fed into a transformer network, making use of a self-attention mechanism, allowing the model to learn long-distance correlations between pairs of qubits. The following is an illustration of the architecture of our generative model:\n",
    "\n",
    "<div align=\"center\">\n",
    "<img src=\"resources/figures/tutorial-transformer.png\" width=\"400\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b97eb055-0920-4789-abad-57f94ae8cfce",
   "metadata": {},
   "source": [
    "In this section we will use a model which has already been trained on a dataset with samples of ground states of 80 random Hamiltonians. For this trained model, we generate new samples in an autorgressive manner, making use of the decomposition\n",
    "\n",
    "$$\n",
    "     p_{\\theta,\\phi}(a_1,\\,\\ldots,\\,a_n~\\lvert~\\mathbf{x}) = \\prod_{i=1}^n p_\\theta(a_i~\\lvert~a_{i-1},\\,\\ldots,\\,a_1,\\, g_\\phi(\\mathbf{x})).\n",
    "$$\n",
    "\n",
    "Let us first initialize the generative model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ec38be-6228-461d-a6db-d069ff170f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize transformer\n",
    "transformer = initialize()\n",
    "\n",
    "# process coupling matrix\n",
    "coupling_graph = preprocess_coupling_matrix(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccec0946-2359-4db6-a202-aa188bd90ae4",
   "metadata": {},
   "source": [
    "Now, we can generate a number of samples which we will use to estimate properties of the ground state. Note that generating these 10'000 takes around 2 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64011155-25ee-4374-bfa7-10e37b0f85a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample from model\n",
    "samples = transformer.sample(\n",
    "    samples=10000, coupling_graph=coupling_graph, qubits=wires, batch_size=1000, print_progress=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1763306-a19c-4c43-8e78-41736b2bc41e",
   "metadata": {},
   "source": [
    "Let us have a quick look at a few samples which our model has generated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5af3fb8e-3f49-4d6d-91a1-e2cd9e067324",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(samples[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b8f9bb-b97a-4d44-849b-4233d9f8e985",
   "metadata": {},
   "source": [
    "Each row in this array corresponds to a particular realization of measuring the Pauli-6 POVM. The encoding we have used here corresponds to the following mapping:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    +,~- &\\mapsto 0, 1 ~~~~ \\text{Pauli-X Basis}\\\\\n",
    "    r,~l &\\mapsto 2, 3 ~~~~ \\text{Pauli-Y Basis}\\\\\n",
    "    0,~1 &\\mapsto 4, 5 ~~~~ \\text{Pauli-Z Basis}\n",
    "\\end{aligned}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc009cc-2a49-4fbe-bfe0-3b4687bddf3f",
   "metadata": {},
   "source": [
    "To construct a classical shadow from this string of measurements, we convert it to two numpy arrays; the first contains the binary measurement outcomes, and the second contains the bases in which the qubits were measured:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed4492e2-970a-4d8e-94f2-ffc342d12342",
   "metadata": {},
   "source": [
    "### Predicting Properties using Classical Shadows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44500879-b3cf-4aa3-adcc-db725392b099",
   "metadata": {},
   "source": [
    "Given our model, we can sample from it and generate a shadow state, based on which we compute entanglement entropies and correlation functions. In the following, we use the classical shadow implementation from PennyLane to predict these properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9538da7-ab66-41a3-b89b-00df43b04ac2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# splits samples into bits (i.e., measurement outcomes) and recipes (measurement bases)\n",
    "recipes = samples // 2\n",
    "bits = samples - 2 * recipes\n",
    "\n",
    "# instantiate classical shadow\n",
    "shadow = qml.ClassicalShadow(bits=bits, recipes=recipes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113cc878-bb22-44ad-9686-e0dee8e75ceb",
   "metadata": {},
   "source": [
    "The following function uses jax to compute entanglement entropies based on the classical shadow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc17ed59-51a0-4551-9722-7533bd5dadf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def _jax_compute_size_one_entropies(x):\n",
    "    return -jnp.log(jnp.einsum('timl,silm->i', x, x))\n",
    "\n",
    "@jit\n",
    "def _jax_compute_size_two_entropies(x):\n",
    "    return -jnp.log(jnp.einsum('tilm,siml,tjrk,sjkr->ij', x, x, x, x))\n",
    "\n",
    "def compute_entropies_from_shadow(shadow):\n",
    "    \"\"\"\n",
    "    compute second-order Rényi entanglement entropies for all subsystems of size at most two, using the classical shadow\n",
    "    protocol\n",
    "    \"\"\"\n",
    "    local_snapshots = shadow.local_snapshots()\n",
    "    shadow_size = shadow.snapshots\n",
    "\n",
    "    # compute size two entropies\n",
    "    entropies = np.array(_jax_compute_size_two_entropies(local_snapshots) + 2 * np.log(shadow_size))\n",
    "\n",
    "    # compute size one entropies\n",
    "    entropies_size_one = np.array(_jax_compute_size_one_entropies(local_snapshots) + 2 * np.log(shadow_size))\n",
    "    np.fill_diagonal(entropies, entropies_size_one)\n",
    "\n",
    "    return entropies.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64f5df8-2145-4c1d-b46d-c7eb967a559e",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_entropy_matrix = compute_entropies_from_shadow(shadow=shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b749e6-7b14-443d-9662-8335f4cd02ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_eval_matrix(\n",
    "    J, predicted_entropy_matrix, true_mat=exact_entropy_matrix, title='Model Prediction', cmap='Blues', vmin=0.0, vmax=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3bfcd7e-41b0-4aac-b098-bc05063d3cbf",
   "metadata": {},
   "source": [
    "We can also estimate the correlation matrix with our classical shadow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb629ad2-7c7a-4033-8bea-c3f0542228f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_correlation_matrix_from_shadow(shadow):\n",
    "    wires = shadow.bits.shape[1]\n",
    "\n",
    "    qubit_pairs = list(it.combinations(range(wires), r=2))\n",
    "\n",
    "    correlations = np.zeros((wires, wires))\n",
    "    np.fill_diagonal(correlations, 1.0)\n",
    "\n",
    "    for idx, (i, j) in enumerate(qubit_pairs):\n",
    "        obs = qml.PauliX(i) @ qml.PauliX(j) + qml.PauliY(i) @ qml.PauliY(j) + qml.PauliZ(i) @ qml.PauliZ(j)\n",
    "        correlations[i, j] = correlations[j, i] = shadow.expval(H=obs, k=1) / 3\n",
    "\n",
    "    return correlations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf25ce55-6e89-4fd2-9ec8-0fdf9631c15b",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_correlation_matrix = compute_correlation_matrix_from_shadow(shadow=shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98bfdbf-f94d-42ba-ab7a-3b20d5069383",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_eval_matrix(\n",
    "    J, predicted_correlation_matrix, true_mat=exact_correlation_matrix, title='Model Prediction', cmap='RdBu'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4faa1f-c3aa-4dfd-8b7d-2d9ca58b0b64",
   "metadata": {},
   "source": [
    "# Training a transformer for a single Hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4dcbe66-fb91-4f6d-a11a-9f4b32c0b8bd",
   "metadata": {},
   "source": [
    "We have seen how a trained conditional generative model can be used to make predictions about different quantum states. Here we will go through the training process and train a transformer on measurements. To train a full-fledged conditional generative model on a big enough number of ground states one needs access to GPUs and -- even then -- training the model to convergence takes around 2 hours of GPU time. For this reason, here we will train a model on a single Hamiltonian as this can be done in reasonable time on a modern laptop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6185ff06-cb30-48f3-92e1-180213a9f52a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch_geometric.data import Data as GraphData\n",
    "from torch_geometric.data import Batch as GraphBatch\n",
    "\n",
    "from src.models.graph_encoder import get_graph_encoder\n",
    "from src.models.gctransformer import init_gctransformer, get_sample_structure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "061d28bd-2725-4c36-a653-b70a5d826db4",
   "metadata": {},
   "source": [
    "Let us first collect measurements from the above ground state. We do this by making measurements in random Pauli bases, as it is done in the classical shadow protocol."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6bedb93-ba54-4b69-a8e6-2a413a3bfd6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_pauli6_samples(statevector, wires, shots, device_name='default.qubit'):\n",
    "    \"\"\" generate samples from the Pauli-6 POVM. The resulting samples are encoded as integers according to\n",
    "    the rules\n",
    "    0,1 -> +,- in the X Basis\n",
    "    2,3 -> r,l in the Y Basis\n",
    "    4,5 -> 0,1 in the Z Bassi\n",
    "    \"\"\"\n",
    "    @qml.qnode(device=qml.device(device_name, wires=wires, shots=shots), diff_method=None, interface=None)\n",
    "    def shadow_measurement():\n",
    "        qml.QubitStateVector(statevector, wires=range(wires))\n",
    "        return qml.classical_shadow(wires=range(wires))\n",
    "\n",
    "    bits, recipes = shadow_measurement()\n",
    "\n",
    "    # encode measurements and bases as integers\n",
    "    data = 2 * recipes + bits\n",
    "    data = np.array(data, dtype=int)\n",
    "\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff747989-55fa-4535-b9b1-a4534106cb7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate data\n",
    "train_samples = 1000\n",
    "measurements = generate_pauli6_samples(ground_state, wires=wires, shots=train_samples)\n",
    "\n",
    "# next, we bring the measurements into the format expected by the transformer \n",
    "measurements = measurements + 1\n",
    "measurements = np.concatenate([0 * np.ones(shape=(train_samples, 1)), measurements], axis=1)\n",
    "measurements = measurements.astype(int)\n",
    "\n",
    "# convert measurements to torch tensor\n",
    "measurements_tensor = torch.from_numpy(measurements)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d172e4b-a15f-4f18-83b9-e56119a15e6d",
   "metadata": {},
   "source": [
    "#### Setting up the generative model\n",
    "Now, we setup our conditional generative model, consisting of a graph encoder and a transformer. We use a graph encoder with 3 layers and final node dimension 16:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5bdbedc-b9a6-45d4-a4ea-65e70137df2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_encoder = get_graph_encoder(\n",
    "        'gcn_proj_3_16', in_node_dim=wires, gcn_dim=16, gcn_layers=3, d_model=128, qubits=wires\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09a620e5-a1b6-441a-8f43-61d5ddd2c36e",
   "metadata": {},
   "source": [
    "Next, we intialize the transformer, consisting of 4 layers, model dimension 128 and 4 attention heads in each self-attention module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a7410a7-8c95-4c6c-ab4b-1fd550a0b14a",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = init_gctransformer(\n",
    "        n_outcomes=6, # Number of distinct Pauli-6 measurements\n",
    "        encoder=graph_encoder,\n",
    "        n_layers=4, # number of sublayers\n",
    "        d_model=128, # hidden dimension\n",
    "        d_ff=4 * 128, # number of hidden units in the fully connected layers\n",
    "        n_heads=4, # number of attention heads\n",
    "        dropout=0.0,\n",
    "        pad_token=-1, # no padding\n",
    "        start_token=0, # start token marks the beginning of a sequence\n",
    "        end_token=-1, # no end token in alphabet\n",
    "        token_shift=1 # shift tokens by 1 to account for the start token\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbdbcc9a-5567-4bf9-82e3-af702ec1f610",
   "metadata": {},
   "source": [
    "In the next step, we normalize the coupling matrix and convert it into a graph structure that can be fed to the graph encoder using the `pytorch_geometric` libary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26ef2e82-2d61-4050-b127-868aa8190b02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize\n",
    "cmax = np.max(J, axis=(0, 1)).reshape((1, 1))\n",
    "cmin = np.min(J, axis=(0, 1)).reshape((1, 1))\n",
    "J_normalized = (J - cmin) / (cmax - cmin)\n",
    "\n",
    "coupling_matrices = [J_normalized for _ in range(train_samples)]\n",
    "\n",
    "# setup edge indices\n",
    "graphs_edge_indices = [\n",
    "    [\n",
    "        [(i, j), (j, i)] for i, j in it.combinations(range(wires), 2) if mat[i, j] != 0.0\n",
    "    ] for mat in coupling_matrices\n",
    "]\n",
    "\n",
    "# write graph edge weights to torch tensor in COO format and convert to tensor\n",
    "graphs_edge_weights = [\n",
    "    torch.from_numpy(np.array([mat[i, j] for i, j in edges], dtype=np.float32).flatten())\n",
    "    for edges, mat in zip(graphs_edge_indices, coupling_matrices)\n",
    "]\n",
    "\n",
    "# convert to tensor in COO format for torch_geometric\n",
    "graphs_edge_indices = [\n",
    "    torch.from_numpy(np.array([x for x in it.chain(*edge_indices)]).T)\n",
    "    for edge_indices in graphs_edge_indices\n",
    "]\n",
    "\n",
    "\n",
    "graphs_list = [GraphData(\n",
    "    x=torch.eye(wires, dtype=torch.float32), \n",
    "    edge_index=edges, edge_weight=weights\n",
    ") for edges, weights in zip(graphs_edge_indices, graphs_edge_weights)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f0b0bb-3ce2-4a0b-a36b-0318eb7f2de1",
   "metadata": {},
   "source": [
    "Now that we have our training data ready, let su create an iterator that allows us to loop through the data in batches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4faa443d-d5f7-4686-9633-fa0b227901d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TgtBatch:\n",
    "    \"\"\" object for holding a batch of data with mask \"\"\"\n",
    "    def __init__(self, tgt, pad):\n",
    "        self.tgt = tgt[:, :-1]\n",
    "        self.tgt_y = tgt[:, 1:]\n",
    "        self.tgt_mask = make_std_mask(self.tgt, pad)\n",
    "        self.ntokens = (self.tgt_y != pad).data.sum()  # noqa\n",
    "        \n",
    "def subsequent_mask(size):\n",
    "    \"\"\" mask out subsequent positions \"\"\"\n",
    "    attn_shape = (1, size, size)\n",
    "    mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)\n",
    "    return mask == 0\n",
    "\n",
    "\n",
    "def make_std_mask(tgt, pad):\n",
    "    \"\"\" Create a mask to hide padding and future words. \"\"\"\n",
    "    tgt_mask = (tgt != pad).unsqueeze(-2)  # noqa\n",
    "    tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)\n",
    "    return tgt_mask\n",
    "\n",
    "\n",
    "def batch_iterator(data, graphs_list, iterations, batch_size, shuffle=True):\n",
    "    for step in range(1, iterations + 1):\n",
    "        # shuffle data\n",
    "        indices = rng.choice(len(data), batch_size, replace=True)\n",
    "        tgt_batch = TgtBatch(data[indices], pad=-1)\n",
    "        graphs_batch = GraphBatch.from_data_list(data_list=[graphs_list[i] for i in indices])\n",
    "        \n",
    "        yield tgt_batch, graphs_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97c57c98-4c81-4fcd-aaab-f37bee45ca08",
   "metadata": {},
   "source": [
    "Next, we set up the KL-divergence loss function. This loss function is more general than the standard KL divergence loss functions as it allows for padding token which are ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7878169f-a7b6-45d5-bc62-eb0cec659c32",
   "metadata": {},
   "outputs": [],
   "source": [
    "class KLDivLoss(nn.Module):\n",
    "    def __init__(self, padding_idx):\n",
    "        super(KLDivLoss, self).__init__()\n",
    "        self.criterion = nn.KLDivLoss(reduction='sum')\n",
    "        self.padding_idx = padding_idx\n",
    "\n",
    "    def forward(self, x: torch.Tensor, target: torch.Tensor):\n",
    "        target_onehot = torch.zeros_like(x)\n",
    "        target_onehot.scatter_(1, target.unsqueeze(1), 1.0)\n",
    "\n",
    "        if self.padding_idx >= 0:\n",
    "            target_onehot[:, self.padding_idx] = 0\n",
    "\n",
    "        mask = torch.nonzero(target.data == self.padding_idx)  # noqa\n",
    "\n",
    "        if mask.dim() > 0:\n",
    "            target_onehot.index_fill_(0, mask.squeeze(), 0.0)\n",
    "\n",
    "        loss = self.criterion(x, target_onehot)\n",
    "        return loss\n",
    "    \n",
    "criterion = KLDivLoss(padding_idx=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9843d743-d3e3-4006-b867-04c9c5375fb7",
   "metadata": {},
   "source": [
    "Finally, we also need an optimizer and a learning rate scheduler. We use the AdamW optimizer (which optionally would allow for weight decay regularization, in contrast to the standard Adam optimizer), and a cosine annealing scheduler to decay the learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f300d583-2bb2-4454-ae63-fc6147673e8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_steps = 2000\n",
    "optimizer = torch.optim.AdamW(transformer.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-9)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae6852f7-a9bd-4518-b35d-2ae21d514448",
   "metadata": {},
   "source": [
    "We can now start training our model, using the batch iterator we have setup above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5ae1087-1cf1-4a46-8054-28992c64d6eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "pbar = tqdm(\n",
    "    enumerate(batch_iterator(measurements_tensor, graphs_list, iterations=total_steps, batch_size=100)),\n",
    "    total=total_steps\n",
    ")\n",
    "\n",
    "# set transformer to train mode\n",
    "transformer.train()\n",
    "\n",
    "losses = []\n",
    "\n",
    "# iterate over data\n",
    "for step, (tgt_batch, graphs_batch) in pbar:\n",
    "    graph_embed, out = transformer.forward(tgt_batch.tgt, tgt_batch.tgt_mask, coupling_graph=graphs_batch)\n",
    "    log_probs = transformer.generator(out)\n",
    "    \n",
    "    total_loss = criterion(\n",
    "        log_probs.contiguous().view(-1, log_probs.size(-1)), tgt_batch.tgt_y.contiguous().view(-1)\n",
    "    ) / tgt_batch.ntokens.item()\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    total_loss.backward()\n",
    "    optimizer.step()\n",
    "    scheduler.step()\n",
    "    \n",
    "    # record loss\n",
    "    losses.append(total_loss.item())\n",
    "    \n",
    "    pbar.set_postfix_str(f'[{step+1}/{total_steps}] loss: {total_loss.item():.6f}, lr: {scheduler.get_last_lr()[0]:.6f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2401a18-7abb-4280-9f9c-ef55816cb334",
   "metadata": {},
   "source": [
    "Now that we have trained our model, let us first visualize the loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2733463-57fc-46b6-bb0f-686b787ea9a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style('whitegrid')\n",
    "plt.plot(np.arange(total_steps), losses);\n",
    "plt.xlabel('Step', fontsize=16);\n",
    "plt.ylabel('Loss', fontsize=16);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98020333-6677-41e5-9dfa-5808634b2130",
   "metadata": {},
   "source": [
    "Next, let us use the model the generate a number of samples which we can use to predict the correlation function and entanglement entropies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "779c53e4-dc61-4f3d-b01e-f1ebc18697d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = transformer.sample(\n",
    "    samples=10000, coupling_graph=coupling_graph, qubits=wires, batch_size=1000, print_progress=False\n",
    ")\n",
    "\n",
    "# splits sampels into measurement outcomes and bases\n",
    "recipes = samples // 2\n",
    "bits = samples - 2 * recipes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfce3380-63a6-4b67-9074-97af470caccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate shadow\n",
    "shadow = qml.ClassicalShadow(bits=bits, recipes=recipes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc1a94c-02d7-4a42-afd1-50bdbd232cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute renyi-2 entropies using classical shadow\n",
    "model_entropy_matrix = compute_entropies_from_shadow(shadow=shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a87dc20-9091-47fd-8fa8-96dc879f022c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style('white')\n",
    "plot_eval_matrix(\n",
    "    J, model_entropy_matrix, true_mat=exact_entropy_matrix, title='Model Prediction', cmap='Blues', vmin=0.0, vmax=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fed0465-9e21-4b48-9564-842485ed29a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute correlation matrix using classical shadow\n",
    "model_correlation_matrix = compute_correlation_matrix_from_shadow(shadow=shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac4cd7a4-0415-4439-86f9-4248cb0b0679",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style('white')\n",
    "plot_eval_matrix(\n",
    "    J, model_correlation_matrix, true_mat=exact_correlation_matrix, title='Model Prediction', cmap='RdBu'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9beae7c-f510-49bc-98c6-9926bcc58c69",
   "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.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "8757b7934aa9b77b37b5993bfa43adcaa32fcacf6406e599a8158e6eb26bf7b6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}