{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "af606719",
   "metadata": {},
   "source": [
    "# Example - Tensor Renormalization Group (TRG)\n",
    "\n",
    "TRG[[1](https://arxiv.org/abs/cond-mat/0611687), [2](https://arxiv.org/abs/0806.3509), [3](https://tensornetwork.org/trg/)]\n",
    "is an tensor network algorithm for computing partition functions of 2D classical spin models, using real space renormalization.\n",
    "It is simple but quite powerful, and the basis for many more advanced algorithms.\n",
    "\n",
    "In its simplest form it only requires a manipulating a few tensors, so does not require any of the `quimb`\n",
    "functionality dealing with large and complex geometry networks. However, implementing it here does demonstrate:\n",
    "\n",
    "* the basic low-level tensor operations of contracting, decomposing and relabelling indices etc.\n",
    "* the more advanced feature of treating a small tensor network transparently as a 'lazy' tensor to enable more efficient iterative operations e.g.\n",
    "\n",
    "## Define the algorithm\n",
    "\n",
    "The following function runs the entire algorithm and is pretty extensively commented:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30af91eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import quimb.tensor as qtn\n",
    "from autoray import do\n",
    "from math import log, log1p, cosh, sinh, cos, pi\n",
    "\n",
    "\n",
    "def TRG(\n",
    "    beta,\n",
    "    chi,\n",
    "    iterations,\n",
    "    j=1.0,\n",
    "    h=0.0,\n",
    "    cutoff=0.0,\n",
    "    lazy=False,\n",
    "    to_backend=None,\n",
    "    progbar=False,\n",
    "    **split_opts\n",
    "):\n",
    "    \"\"\"Run the TRG algorithm on the square lattice.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    beta : float\n",
    "        Inverse temperature.\n",
    "    chi : int\n",
    "        The maximum bond dimension.\n",
    "    iterations : int\n",
    "        The number of iterations, the overall effective lattice size is then\n",
    "        ``(2**iterations, 2**iterations)``, with PBC.\n",
    "    j : float, optional\n",
    "        The coupling constant.\n",
    "    h : float, optional\n",
    "        The external field.\n",
    "    cutoff : float, optional\n",
    "        The cutoff for the bond truncations.\n",
    "    lazy : bool, optional\n",
    "        Whether to explicitly contract the effective site tensor at each\n",
    "        iteration (``False``), or treat it lazily as the loop from the last\n",
    "        iteration, allowing a more efficient iterative decomposition at large\n",
    "        ``chi``.\n",
    "    to_backend : callable, optional\n",
    "        A function that takes a numpy array and converts it to the desired\n",
    "        backend tensor.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    f : scalar\n",
    "        The free energy per site.\n",
    "    \"\"\"\n",
    "    if lazy and cutoff == 0.0:\n",
    "        # by default use a low-rank iterative decomposition\n",
    "        split_opts.setdefault('method', 'svds')\n",
    "\n",
    "    # setup the initial single site array, allowing custom backends\n",
    "    t = qtn.tensor_builder.classical_ising_T_matrix(beta, j=j, h=h, directions='lrud')\n",
    "    if to_backend is not None:\n",
    "        t = to_backend(t)\n",
    "\n",
    "    # This is the effective lattice\n",
    "    #\n",
    "    #      u          u\n",
    "    #      |          |\n",
    "    #   l--A--r .. l--A--r\n",
    "    #      |          |\n",
    "    #      d          d\n",
    "    #      :          :\n",
    "    #      u          u\n",
    "    #      |          |\n",
    "    #   l--A--r .. l--A--r\n",
    "    #      |          |\n",
    "    #      d          d\n",
    "    #\n",
    "    A = qtn.Tensor(t, ('d', 'l', 'u', 'r'))\n",
    "\n",
    "    # track the very large overall scalar in log with this\n",
    "    exponent = 0.0\n",
    "\n",
    "    if progbar:\n",
    "        import tqdm\n",
    "        its = tqdm.trange(2 * iterations)\n",
    "    else:\n",
    "        its = range(2 * iterations)\n",
    "\n",
    "    for i in its:\n",
    "\n",
    "        # split site tensor in two ways:\n",
    "        #     u                    u\n",
    "        #     |                    |\n",
    "        #  l--A--r  ->  l--AL~~b~~AU--r\n",
    "        #     |            |\n",
    "        #     d            d\n",
    "        AL, AU = A.split(\n",
    "            left_inds=['d', 'l'], get='tensors', bond_ind='b',\n",
    "            max_bond=chi, cutoff=cutoff, **split_opts)\n",
    "        #     u            u\n",
    "        #     |            |\n",
    "        #  l--A--r  ->  l--BU~~b~~BL--r\n",
    "        #     |                   |\n",
    "        #     d                   d\n",
    "        BU, BL = A.split(\n",
    "            left_inds=['l', 'u'], get='tensors', bond_ind='b',\n",
    "            max_bond=chi, cutoff=cutoff, **split_opts)\n",
    "\n",
    "        # reindex to form a plaquette\n",
    "        #             u\n",
    "        #  l ~~BL--AL~~\n",
    "        #      |    |       w/ inner loop indices: dp, lp, up, rp\n",
    "        #    ~~AU--BU~~ r\n",
    "        #    d\n",
    "        AU.reindex_({'b': 'd', 'r': 'dp', 'u': 'lp'})\n",
    "        BL.reindex_({'b': 'l', 'd': 'lp', 'r': 'up'})\n",
    "        AL.reindex_({'b': 'u', 'l': 'up', 'd': 'rp'})\n",
    "        BU.reindex_({'b': 'r', 'u': 'rp', 'l': 'dp'})\n",
    "\n",
    "        # we can just form the TN of this loop and treat like a tensor\n",
    "        A = (AU | BL | AL | BU)\n",
    "        if not lazy:\n",
    "            # ... or contract to dense A tensor explicitly\n",
    "            A = A.contract()\n",
    "\n",
    "        # bookeeping: move normalization into separate 'exponent'\n",
    "        nfact = A.largest_element()\n",
    "        A /= nfact\n",
    "        exponent *= 2  # first account for lattice doubling in size\n",
    "        exponent += do('log', nfact)\n",
    "\n",
    "    # perform the final periodic trace\n",
    "    mantissa = A.trace(['u', 'd'], ['l', 'r'])\n",
    "\n",
    "    # combine with the separately tracked  exponent\n",
    "    logZ = do('log', mantissa) + exponent\n",
    "    N = 2**(iterations * 2)\n",
    "\n",
    "    return - logZ / (N * beta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8383da5c-3aa3-4eb1-90f8-aa22f44e788b",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "tags": []
   },
   "source": [
    "Note we are mostly just are manipulating a few objects at the\n",
    "{class}`~quimb.tensor.tensor_core.Tensor` level. However, our main object `A` can actually be a\n",
    "{class}`~quimb.tensor.tensor_core.TensorNetwork` because many methods have exactly\n",
    "the same signature and usage, specifically here:\n",
    "\n",
    "- [`Tensor.reindex`](quimb.tensor.tensor_core.Tensor.reindex) / [`TensorNetwork.reindex`](quimb.tensor.tensor_core.TensorNetwork.reindex)\n",
    "- [`Tensor.split`](quimb.tensor.tensor_core.Tensor.split) / [`TensorNetwork.split`](quimb.tensor.tensor_core.TensorNetwork.split)\n",
    "- [`Tensor.largest_element`](quimb.tensor.tensor_core.Tensor.largest_element) / [`TensorNetwork.largest_element`](quimb.tensor.tensor_core.TensorNetwork.largest_element)\n",
    "- [`Tensor.trace`](quimb.tensor.tensor_core.Tensor.trace) / [`TensorNetwork.trace`](quimb.tensor.tensor_core.TensorNetwork.trace)\n",
    "\n",
    "## Run the algorithm\n",
    "\n",
    "We can run the function for pretty large `chi` if we use this lazy iterative\n",
    "feature, (which doesn't affect accuracy):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e2a07a81",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 32/32 [01:33<00:00,  2.92s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-2.1096509964871495"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chi = 64\n",
    "\n",
    "# the critical temperature is known analytically\n",
    "beta = log1p(2**0.5) / 2\n",
    "\n",
    "f = TRG(\n",
    "    beta=beta,\n",
    "    chi=chi,\n",
    "    iterations=16,  # L = 2**16\n",
    "    lazy=True,      # lazily treat loop TN as new tensor\n",
    "    progbar=True,\n",
    ")\n",
    "f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a1c0033",
   "metadata": {},
   "source": [
    "## Check against exact result\n",
    "\n",
    "\n",
    "The exact free energy is also known analytically in the thermodynamic\n",
    "limit[[4](https://journals.aps.org/pr/abstract/10.1103/PhysRev.65.117), [5](https://en.wikipedia.org/wiki/Ising_model#Onsager's_exact_solution)],\n",
    "which we can compute here as a check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "09439da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def free_energy_2d_exact(beta, j=1.0):\n",
    "    from scipy.integrate import quad\n",
    "\n",
    "    def inner1(theta1, theta2):\n",
    "        return log(\n",
    "            cosh(2 * beta * j)**2 -\n",
    "            sinh(2 * beta * j) * cos(theta1) -\n",
    "            sinh(2 * beta * j) * cos(theta2)\n",
    "        )\n",
    "\n",
    "    def inner2(theta2):\n",
    "        return quad(\n",
    "            lambda theta1: inner1(theta1, theta2),\n",
    "            0, 2 * pi,\n",
    "        )[0]\n",
    "\n",
    "    I = quad(inner2, 0, 2 * pi)[0]\n",
    "    return -(log(2) + I / (8 * pi**2)) / beta\n",
    "\n",
    "fex = free_energy_2d_exact(beta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f8ec2a9",
   "metadata": {},
   "source": [
    "So our relative error is given by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7603344c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7.02111621064816e-08"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "err = 1 - f / fex\n",
    "err"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "143476fb-ab2d-4172-bd57-ffc206567061",
   "metadata": {
    "raw_mimetype": "text/restructuredtext",
    "tags": []
   },
   "source": [
    "## Extensions\n",
    "\n",
    "Which is pretty decent, though methods which take into account the environement\n",
    "when truncating can do even better. Things you might try:\n",
    "\n",
    "- use a GPU backend (pass `to_backend`), this might require `method='svd'` and `lazy=False`\n",
    "- use other iterative SVD methods (e.g. `'isvd'` or `'rsvd'`) and play with `lazy`\n",
    "- using {meth}`~quimb.tensor.tensor_core.TensorNetwork.fit` to optimize the projectors at each step"
   ]
  }
 ],
 "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "39c10650315d977fb13868ea1402e99f3e10e9885c2c202e692ae90b8995050d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}