peps-exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e61218b5-796d-45bb-8bb4-5d98e23070f5",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Parition function of classical 2D stat-mech model\n",
    "\n",
    "2D Ferromagnetic Ising model\n",
    "\n",
    "    H = -\\sum_{\\langle ij \\rangle} \\sigma_i \\sigma_j\n",
    "\n",
    "Partition function - sum of Boltzman weights for all configurations\n",
    "\n",
    "    Z(\\beta) = \\sum_{\\mathbf{\\sigma}} exp(-\\Beta H) = \\sum_{\\mathbf{\\sigma}} \\prod_{\\langle ij \\rangle} exp(\\Beta \\sigma_i \\sigma_j)\n",
    "    \n",
    "Partition function = contraction of a tensor network\n",
    "\n",
    "    x--h--x--h--x--h--x\n",
    "    |     |     |     |\n",
    "    h     h     h     h\n",
    "    |     |     |     |\n",
    "    x--h--x--h--x--h--x\n",
    "    |     |     |     |\n",
    "    h     h     h     h\n",
    "    |     |     |     |\n",
    "    x--h--x--h--x--h--x\n",
    "    \n",
    "where rank-2 tensor h<=>exp(\\Beta \\sigma_i \\sigma_j) is a Boltzman weight located on the edge between two neighbouring Ising spins ij \n",
    "and x<=>\\delta_{uldr} is a rank-4 tensor enforcing the same Ising spin state to enter the Boltzman weights on the adjacent edges.\n",
    "\n",
    "By writing out h=\\sqrt(h)\\sqrt(h), i.e. as two weights with temperature \\Beta/2, on each edge and absorbing these weights symmetrically\n",
    "into x's on each vertex we get familiar PEPS form\n",
    "\n",
    "                |\n",
    "              \\sqrt(h)   \n",
    "                |                 |\n",
    "    --\\sqrt(h)--x--\\sqrt(h)-- = --A--\n",
    "                |                 |\n",
    "             \\sqrt(h)\n",
    "                |\n",
    "\n",
    "The partition function Z(A(\\beta)) is given by the contraction\n",
    "\n",
    "    A--A--A\n",
    "    A--A--A\n",
    "    A--A--A"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff04980c-1cd1-46b3-89ac-c19a39442c45",
   "metadata": {},
   "source": [
    "## 0) preliminary\n",
    "\n",
    "Install prerequisites using conda, pip, ...\n",
    "\n",
    "``conda install -c pytorch -c anaconda -c conda-forge pytorch cpuonly scipy pytest opt_einsum``\n",
    "\n",
    "or\n",
    "\n",
    "``pip3 install torch --index-url https://download.pytorch.org/whl/cpu scipy pytest opt_einsum``\n",
    "\n",
    "Custom definitions (might be useful later)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4617f2ef-2c7c-4a3c-aa66-ec29c36868d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.linalg import sqrtm\n",
    "\n",
    "def f_sqrth(beta):\n",
    "    r\"\"\"\n",
    "    Square root of rank-2 Boltzman weight from Ising interaction\n",
    "    \"\"\"\n",
    "    return sqrth\n",
    "    \n",
    "def f_A(beta):\n",
    "    r\"\"\"\n",
    "    On-site tensor A generating the partition function Z(\\Beta)\n",
    "    \"\"\"\n",
    "    return A"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "743c9319-71a5-4701-bb93-7b0750088810",
   "metadata": {},
   "source": [
    "### i.1) Check symmetries of A\n",
    "\n",
    "Convince yourself, that tensor A obeys point-group symmetries of square lattice.\n",
    "\n",
    "A_uldr = A_ldru rotation by pi/2 clockwise\n",
    "\n",
    "A_uldr = A_ruld rotation by pi/2 anti-clockwise\n",
    "\n",
    "A_uldr = A_dlur reflection along x-axis\n",
    "\n",
    "A_uldr = A_urdl reflection along y-axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0240287b-a9cf-45c5-9434-3f97df1b2457",
   "metadata": {},
   "outputs": [],
   "source": [
    "..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f7d9c3-e0c2-4ac4-b202-d7d1f0e6cf54",
   "metadata": {},
   "source": [
    "## i) Build initial CTM boundary \n",
    "\n",
    "Consider two options for initial boundary\n",
    "\n",
    "Option a): Initial C,T are created from Boltzman weights on the open-boundary edges \n",
    "              \n",
    "    C-- = x--\\sqrt(h)\n",
    "    |     |\n",
    "          \\sqrt(h) \n",
    "        \n",
    "    --T-- = \\sqrt(h)--x--\\sqrt(h)\n",
    "      |               | \n",
    "                     \\sqrt(h)\n",
    "                     \n",
    "Option b): Take random (symmetric) C,T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9324e776-75fd-4798-aa29-c01de40d4105",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_obc_CT(beta,chi):\n",
    "    r\"\"\"\n",
    "    Build initial boundaries from Boltzman weights of open boundary system\n",
    "    \"\"\"\n",
    "    return C,T\n",
    "\n",
    "def init_rand_CT(A,chi):\n",
    "    r\"\"\"\n",
    "    Build random initial boundaries C, T\n",
    "    \"\"\"\n",
    "    return C,T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3793f04a-f01d-4f6d-ae21-0ec727aa675c",
   "metadata": {},
   "source": [
    "## ii) The CTMRG move\n",
    "\n",
    "First, enlarge the system\n",
    "\n",
    "    C_i--T_i--C_i      C_i--T_i--T_i--T_i--C_i\n",
    "    T_i--A----T_i  =>  T_i--A----A----A----C_i\n",
    "    C_i--T_i--C_i      T_i--A----A----A----C_i\n",
    "                       T_i--A----A----A----C_i\n",
    "                       C_i--T_i--T_i--T_i--C_i\n",
    "\n",
    "### ii.1) Expand\n",
    "Define and accumulate enlarged quadrant of the system into enlarged corner tensor\n",
    "of size (\\chi*D)^2\n",
    "\n",
    "    C_i--T_i-- = C'_i==\n",
    "    |    |       ||\n",
    "    T_i--A----\n",
    "    |    |\n",
    "    \n",
    "Afterwards, compute its optimal compression up to rank \\chi and retrieve corresponding projectors\n",
    "\n",
    "    C'== =  U'--D'--U'^\\dag==  => compress =>  U--\\chi--D--\\chi--U^\\dag==    \n",
    "    ||      ||                                 ||"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188f4dd3-38ba-432d-aa08-ae5cb12cbd81",
   "metadata": {},
   "outputs": [],
   "source": [
    "def enlarged_C(A,C,T):\n",
    "    r\"\"\"\n",
    "    Build enlarged corner tensor\n",
    "    \n",
    "        C----1 1(0)---T--2(1) = C'==2,7\n",
    "        0             3(2)      ||\n",
    "        0(1)          3(0)      4,6\n",
    "        T--5(2) 5(1)--A--7(3)\n",
    "        4(0)          6(2) \n",
    "    \"\"\"\n",
    "    return C2x2\n",
    "\n",
    "def compress(chi,C):\n",
    "    r\"\"\"\n",
    "    Compute optimal projectors U, which compress C' to at most rank-\\chi corner tensor C.\n",
    "    Return leading chi part D of the spectrum (by magnitude) and corresponding projector.\n",
    "    \"\"\"\n",
    "    return D,U"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f19ec80-b63d-44e4-8a41-e3b57bc54f2a",
   "metadata": {},
   "source": [
    "### ii.2) Absorb\n",
    "\n",
    "Renormalize the network, be absorbing and compressing enlarged corners and half-row/-column transfer matrices\n",
    "\n",
    "    C_i--T_i--T_i--T_i--C_i      C_{i+1}--T_{i+1}--C_{i+1}     \n",
    "    T_i--A----A----A----C_i  =>  T_{i+1}--A--------T_{i+1}\n",
    "    T_i--A----A----A----C_i      C_{i+1}--T_{i+1}--C_{i+1} \n",
    "    T_i--A----A----A----C_i\n",
    "    C_i--T_i--T_i--T_i--C_i\n",
    "\n",
    "Using optimal projectors, extend half-row/-column transfer matrix and compress it\n",
    "\n",
    "                                     |\n",
    "                                     U\n",
    "    |    |     ||                    ||       |\n",
    "    T_i--A-- = T'_i-- => compress => T'_i-- = T_{i+1}\n",
    "    |    |     ||                    ||       |\n",
    "                                     U^\\dag\n",
    "                                     |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbb5b78-6642-49b2-9b0f-ce561d88b859",
   "metadata": {},
   "outputs": [],
   "source": [
    "def absorb(A,T,U):\n",
    "    r\"\"\"\n",
    "    Aborb and compress half-row/-column tensor given optimal projector of enlarged corner\n",
    "    \n",
    "        4(2)\n",
    "        U\n",
    "        |---------\\   \n",
    "        1(0)       3(1)\n",
    "        1          3(0)      1\n",
    "        T--2 2(1)--A--6(3) = T--2\n",
    "        0          5(2)      0\n",
    "        0          5(1)\n",
    "        |----------/\n",
    "        U^\\dag\n",
    "        7(2)\n",
    "    \"\"\"\n",
    "    return T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9044aea-86b7-4b16-b288-9bdda41b65c4",
   "metadata": {},
   "source": [
    "### ii.3) Compose enlarge, compress, and absorb operations into a full CTM move\n",
    "\n",
    "The compressed corner is given by leading-chi kept eigenvalues C_{i+1} of enlarged corner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e91027-bbbf-43ed-9d23-947efcb25f01",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ctm_iter(A,C,T,chi):\n",
    "    r\"\"\"\n",
    "    Performs a single CTM interation renormalizing the boundary tensors\n",
    "    to environment dimension \\chi. Return new C_n,T_n\n",
    "    \"\"\"\n",
    "    return C_n,T_n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407a82ba-50cd-4a51-961d-ae906c17a351",
   "metadata": {},
   "source": [
    "## iii) CTMRG algorithm\n",
    "\n",
    "The CTMRG algorithm requires an initial choice of the boundary tensors.\n",
    "Afterwards, it is iterated until the fixed point is reached.\n",
    "\n",
    "* Normalize numerical values of the boundary tensors by making magnitude of the largest element equal to 1.\n",
    "* To establish convergence, let's use spectra of the corners."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e84bfc4-7f2b-4ccf-94b6-d733880045bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ctmrg(A,C,T,chi,ctm_max_iter=100,ctm_conv_tol=1.0e-8,callback=None):\n",
    "    r\"\"\"\n",
    "    Return a fixed point of CTMRG process for PEPS generated by tensor A with initial boundaries C,T\n",
    "    \"\"\" \n",
    "    # setup\n",
    "    \n",
    "    # main ctmrg loop\n",
    "    for i in range(ctm_max_iter):\n",
    "        C,T= ctm_iter(A,C,T,chi)\n",
    "        \n",
    "        # normalize numerical values of boundary tensors by making magnitude of largest element equal to 1\n",
    "        # and symmetrize T tensor (suppres finite-precision error)\n",
    "        \n",
    "        # check convergence\n",
    "        \n",
    "        # (optional) monitor through callback\n",
    "        if converged: break\n",
    "        \n",
    "    return C,T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd24d0f0-cc40-4a89-a138-21446793acbe",
   "metadata": {},
   "source": [
    "## iv) Evaluate on-site observable: magnetization\n",
    "\n",
    "Thermal expectation value of magnetization is given as\n",
    "\n",
    "    m = \\langle \\sigma \\rangle = Tr[\\sigma exp(-\\Beta H)] / Z(\\Beta)\n",
    "    \n",
    "which can be reinterpreted as a ratio of two tensor networks, which we approximate\n",
    "through converged boundary tensors\n",
    "          \n",
    "              ...              ...\n",
    "           A--A--A     /  ..A--A--A..          CTC   / CTC\n",
    "    m = ...A--B--A..  /   ..A--A--A.. \\approx  TBT  /  TAT\n",
    "           A--A--A   /    ..A--A--A            CTC /   CTC\n",
    "              ...              ...\n",
    "          \n",
    "where tensor B carries an on-site observable, \\sigma         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e719100d-7364-45cb-a71e-f6dc3b44a80c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_B(beta):\n",
    "    r\"\"\"\n",
    "    On-site tensor B carrying, obtained from A by modifying the weights with observable \\sigma\n",
    "    \"\"\"\n",
    "    return B\n",
    "\n",
    "def aux_rho_1x1(C,T):\n",
    "    r\"\"\"\n",
    "    Build environment of single site as rank-4 tensor\n",
    "    \n",
    "        i) C--1 1(0)--T--2(1) = C--T--1\n",
    "           0          3(2)      0  2\n",
    "\n",
    "        ii) C--T--1 1(0)--C   = C--T-----C\n",
    "            0  2          |     0  2  3--T\n",
    "                    4(2)--T              1\n",
    "                          3(1)\n",
    "\n",
    "        iii) C------T-------C\n",
    "             0      2       |       C----T----C\n",
    "             0(1)           |       |    0    |\n",
    "             T---4(3)    3--T     = T--1   3--T = rho\n",
    "             |              1       |    2    |\n",
    "             |      5(2)    1(0)    C----T----C\n",
    "             C------T-------C       \n",
    "    \"\"\"\n",
    "    return rho\n",
    "\n",
    "def eval_1x1(B,A,rho):\n",
    "    r\"\"\"\n",
    "    Given 1x1 environment rho, compute the normalized expectation value \n",
    "    of on-site tensor B\n",
    "    \"\"\"\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f99b009-afb9-4e32-a55a-7eea5f916199",
   "metadata": {},
   "source": [
    "## v) Wrap everything up to compute phase diagram\n",
    "\n",
    "The exact solution gives critical temperature and magnetization profile\n",
    "\n",
    "    \\Beta_c = log(1 + \\sqrt(2))/2\n",
    "    \n",
    "    m=(1-[sinh(2)]^{-4})^{1/8}\n",
    "    \n",
    "Pick environment dimension \\chi and compute magnetization profile for a range of \\Beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddef2e87-477b-4ab7-b3c2-7e72b2657dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi=8\n",
    "betas=np.linspace(0.1,0.6,20)\n",
    "\n",
    "for beta in betas:\n",
    "    # initialize PEPS tensor\n",
    "    \n",
    "    # converge boundaries\n",
    "    \n",
    "    # compute magnetization\n",
    "    \n",
    "# plot results across the range of temperatures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3e1c046-0271-4d8d-a708-3b74bda1a608",
   "metadata": {},
   "source": [
    "## vi) What about correlation lengths ?\n",
    "\n",
    "Compute PEPS correlation length from the leading part of the spectrum of widht-1 PEPS transfer matrix\n",
    "\n",
    "    --T--\n",
    "    --A-- = \\sum_i \\lambda_i |v_i><v_i|\n",
    "    --T--"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e25c06-13f9-4d74-bd56-eec9f7fc3d87",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse.linalg import LinearOperator, eigs\n",
    "\n",
    "def xi(A,T,n=2,v0=None):\n",
    "    r\"\"\"\n",
    "    Compute leading n eigenvalues of width-1 transfer operator\n",
    "    \n",
    "    Use dominant eigensolver supplying a matrix-vector product\n",
    "    \n",
    "      --0(chi) 0-----T--3(1)     ---0\n",
    "     |               4(2)       |\n",
    "     |               4(0)       |\n",
    "     v--1(D) 1(1)----A---6(3) = v'--1\n",
    "     |               5(2)       |\n",
    "     |               5(2)       |\n",
    "      --2(chi) 2(1)--T--7(0)     ---2\n",
    "      \n",
    "    Return correlation length xi\n",
    "    \"\"\"\n",
    "    def _mv(v):\n",
    "        # apply transfer matrix to vector v\n",
    "        return v\n",
    "\n",
    "    # call scipy's eigs\n",
    "\n",
    "    # post-process and evaluate correlation length\n",
    "    \n",
    "    return xi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f2a116-06ce-447e-baf4-4fd593ccf4d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi=8\n",
    "betas=np.linspace(0.1,0.6,20)\n",
    "\n",
    "for beta in betas:\n",
    "    # initialize PEPS tensor\n",
    "    \n",
    "    # converge boundaries\n",
    "    \n",
    "    # compute magnetization\n",
    "    # compute correlation length\n",
    "    \n",
    "# plot results across the range of temperatures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3db00cc3-dcf1-4f1d-99ee-efa4f89ab576",
   "metadata": {},
   "source": [
    "# AFM Heisenberg model\n",
    "\n",
    "> This exercise implements a basic form of \n",
    "> https://github.com/jurajHasik/peps-torch/blob/master/examples/j1j2/optim_j1j2_c4v.py\n",
    "\n",
    "We want to find PEPS ansatz, which minimizes the energy of **2D spin-1/2 AFM Heisenberg model on a square lattice**\n",
    "\n",
    "    H = -\\sum_{\\langle ij \\rangle} \\mathbf{S}_i \\cdot \\mathbf{S}_j\n",
    "    \n",
    "where \\mathbf{S}_i is a vector of spin-1/2 generators \\mathbf{S} = (S^z,S^x,S^y)\n",
    "\n",
    "This model characterizes antiferromagnets, arising as a strongly-interacting limit of a single-band Hubbard model\n",
    "at half-filling.\n",
    "\n",
    "Although we expect the ground state to posses Q=(pi,pi), i.e. Neel order, we can still get away with \n",
    "single-site translationally invariant PEPS ansatz. We can \n",
    "\n",
    "    a) obtain Q=(pi,pi) state from a ferromagnet by appropriate unitary rotation -2iS^y on each site \n",
    "       of one of the two sublattices of square lattice\n",
    "     \n",
    "or using complementary point of view\n",
    "     \n",
    "    b) transform the Hamiltonian by unitary, such that its ground state becomes \n",
    "       a translationally invariant correlated ferromagnet\n",
    "\n",
    "We parametrize the PEPS ansatz by a single real and C4v-symmetric rank-5 tensor\n",
    "\n",
    "       u s\n",
    "       |/\n",
    "    l--a--r <=> a^s_uldr with s being physical index (at position 0), followed by auxiliary indices u,l,d,r \n",
    "       |                 one for each edge of site\n",
    "       d\n",
    "       \n",
    "We want this tensor to transform trivially (as an A_1 irrep) under the action of square lattice point group (C_{4v}), i.e.\n",
    "    \n",
    "    * rotations by pi/2 clock-wise and counter-clockwise\n",
    "    * vertical and horizontal reflections"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce1e91d-e271-433f-b463-8072e41f6acf",
   "metadata": {},
   "source": [
    "## 0) preliminary\n",
    "\n",
    "Custom \"vjp\" for eigendecomposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef36d21b-9941-49a8-aed7-6f4902ab60ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "class SYMEIG(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(self, A):\n",
    "        r\"\"\"\n",
    "        :param A: square symmetric matrix\n",
    "        :type A: torch.Tensor\n",
    "        :return: eigenvalues values D, eigenvectors vectors U\n",
    "        :rtype: torch.Tensor, torch.Tensor\n",
    "\n",
    "        Computes symmetric decomposition :math:`M= UDU^\\dagger`.\n",
    "        \"\"\"\n",
    "        # is input validation (A is square and symmetric) provided by torch.linalg.eigh ?\n",
    "\n",
    "        D, U = torch.linalg.eigh(A)\n",
    "        # torch.symeig returns eigenpairs ordered in the ascending order with \n",
    "        # respect to eigenvalues. Reorder the eigenpairs by abs value of the eigenvalues\n",
    "        # abs(D)\n",
    "        absD,p= torch.sort(torch.abs(D),descending=True)\n",
    "        D= D[p]\n",
    "        U= U[:,p]\n",
    "\n",
    "        self.save_for_backward(D,U)\n",
    "        return D,U\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(self, dD, dU):\n",
    "        r\"\"\"\n",
    "        :param dD: gradient on D\n",
    "        :type dD: torch.Tensor\n",
    "        :param dU: gradient on U\n",
    "        :type dU: torch.Tensor\n",
    "        :return: gradient\n",
    "        :rtype: torch.Tensor\n",
    "\n",
    "        Computes backward gradient for ED of symmetric matrix with regularization\n",
    "        of :math:`F_{ij}=1/(D_i - D_j)`\n",
    "        \"\"\"\n",
    "        epsilon=1.0e-12\n",
    "        D, U= self.saved_tensors\n",
    "        Uh = U.t().conj()\n",
    "        D_scale= D[0].abs() # D is ordered in descending fashion by abs val\n",
    "\n",
    "        F = (D - D[:, None])\n",
    "        F = F/(F**2 + epsilon)\n",
    "        F.diagonal().fill_(0)\n",
    "        \n",
    "        dA = U @ (torch.diag(dD) + F*(Uh@dU)) @ Uh\n",
    "        return dA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d9f9d6-94f8-4013-81b0-49c1b61676d2",
   "metadata": {},
   "source": [
    "## i) build initial CTM boundary\n",
    "\n",
    "Consider two options:\n",
    "    \n",
    "    * a) by partial contraction of on-site tensor and it's adjoint\n",
    "    * b) a simple random \"product state-like\" boundary with environment dimension \\chi=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebe367c-a070-48ac-8dc1-2b4a31a85ade",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_pbc_CT(a,chi):\n",
    "    r\"\"\"\n",
    "    Build initial boundaries from on-site tensor\n",
    "\n",
    "            u           \n",
    "        l--a*--y   = C--1\n",
    "          /\\         0\n",
    "         x  m\n",
    "             \\ u\n",
    "           l--a--r\n",
    "             /\n",
    "            d\n",
    "        \n",
    "            x      =   1     \n",
    "        l--a*--z       T--2\n",
    "          /\\           0\n",
    "         y  m\n",
    "             \\ u\n",
    "           l--a--r\n",
    "             /\n",
    "            d\n",
    "    \"\"\"\n",
    "    return C,T\n",
    "\n",
    "def init_random_product_CT(a,chi):\n",
    "    r\"\"\"\n",
    "    Build initial boundaries as trivial C and random T\n",
    "    \"\"\"\n",
    "    return C,T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70219b6-3c26-42b2-907f-32b3ae26d3b0",
   "metadata": {},
   "source": [
    "## ii) redefine CTMRG \n",
    "\n",
    "Using the components from 2D Ising, define CTMRG process now \n",
    "\n",
    "    * utilizing regularized eigendecompostion SYMEIG as D,U= SYMEIG.apply(M)\n",
    " \n",
    "    * renormalizing double-layer networks "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbed710f-d244-405a-92b5-cbf3d3b166db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dl_aadag(a):\n",
    "    r\"\"\"\n",
    "    Construct double-layer tensor \\sum_s a^s_uldr a*^s_u'l'd'r' = A_(uu')(ll')(dd')(rr')\n",
    "    \"\"\"\n",
    "    return A\n",
    "\n",
    "def compress(chi,C):\n",
    "    r\"\"\"\n",
    "    Compute optimal projectors U, which compress C' to at most rank-\\chi corner tensor C\n",
    "    \"\"\"\n",
    "    D,U= SYMEIG.apply(enlarged_corner)\n",
    "    \n",
    "    # reorder by magnitude and truncate\n",
    "    \n",
    "    return D,U\n",
    "\n",
    "def ctmrg_dl(a,C,T,chi,**kwargs):\n",
    "    r\"\"\"\n",
    "    Converge CTMRG and returned fixed point boundary tensors C,T\n",
    "    \"\"\"\n",
    "    return C,T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c609f55-9131-4486-acef-e96da60e8ab2",
   "metadata": {},
   "source": [
    "## iii) Test CTMRG procedure\n",
    "\n",
    "Benchmark your implementation against optimized PEPS\n",
    "(see https://github.com/jurajHasik/j1j2_ipeps_states)\n",
    "\n",
    "1. clone the repo with dataset containing optimized PEPS states for J1-J2 model\n",
    "2. load on-site tensor and build double-layer tensor\n",
    "3. run CTMRG until convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8133d5e6-a1c3-4aca-a192-afa26f9e18cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! git clone https://github.com/jurajHasik/j1j2_ipeps_states.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbcae51f-427c-49bf-8022-2d88233c6979",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os,sys\n",
    "if not any(['j1j2_ipeps_states' in path for path in sys.path]):\n",
    "    sys.path.insert(1, os.path.join(os.getcwd(), 'j1j2_ipeps_states'))\n",
    "    \n",
    "from ipeps_io import load_from_pepstorch_json_dense\n",
    "\n",
    "a= torch.asarray(load_from_pepstorch_json_dense(\"j1j2_ipeps_states/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D3_chi_opt108.json\"))\n",
    "\n",
    "chi=18\n",
    "# run CTMRG and inspect converged boundary tensors\n",
    "\n",
    "# Look at spectrum of the corner. For chi=18, you should obtain following data\n",
    "#\n",
    "# spectrum(C)= [1.0000, -0.5220,  0.5215, -0.4195,  0.4189,  0.3171, -0.2569,  0.2569,\n",
    "#        -0.2108, -0.1329,  0.1295, -0.1295, -0.1285,  0.1282,  0.1096, -0.0842,\n",
    "#         0.0842,  0.0760]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e55bf99-f5d1-418a-acc5-aa6d0dea462e",
   "metadata": {},
   "source": [
    "## iv) Evaluation of variational energy\n",
    "\n",
    "Construct reduced density matrix of nearest-neighbour pair\n",
    "\n",
    "        C--T--------------T--------------C = rho_s0s1,s'0s'1\n",
    "        |  |              |              |     \n",
    "        T--a[s'0]^+a[s0]--a[s'1]^+a[s1]--T   \n",
    "        |  |              |              |\n",
    "        C--T--------------T--------------C\n",
    "        \n",
    "Define exhange spin vector \\mathbf{S}, interaction operator \\mathbf{S}_i \\cdot \\mathbf{S}_j, and the unitary rotation operator -i\\sigma^y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68a36c58-1f0d-4576-ac87-8de581aad00f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rho_2x1(a,C,T):\n",
    "    r\"\"\"\n",
    "    Build reduced density matrix of nearest-neighbours\n",
    "    \n",
    "        i) C--1 1(0)--T--2(1) = C--T--1\n",
    "           0          3(2)      0  2\n",
    "\n",
    "        ii) C--T--1 1(0)--C   = C--T-----C\n",
    "            0  2          |     0  2  3--T\n",
    "                    4(2)--T              1 \n",
    "                          3(1)\n",
    "        \n",
    "        Reshape auxiliary indices of T's and append a, a*, and last T\n",
    "        \n",
    "        iii) C-----------T------------------1\n",
    "             |           |----\\\n",
    "             |           4     5\n",
    "             |           4(1)  5(1)                  C--T-------0\n",
    "             T---2 2(2)--a[6(0)]------------8(4)   = T--A[4,5]==2,3\n",
    "             | \\-3 3(2)--|-----a*[9(0)]-----11(4)    C--T-------1\n",
    "             |           7(3)  10(2)\n",
    "             |           7(2)  10(3)\n",
    "             C--0 0(1)---T------------------12(0)\n",
    "             \n",
    "        iv) C--T-------0   0(1)--T-------C\n",
    "            T--A[4,5]==2,3 2,3===A[6,7]--T = rho_2,6;5,7\n",
    "            C--T-------1   1(0)--T-------C\n",
    "    \"\"\"\n",
    "    # contract TN corresponding to 2x1 patch embedded in environment\n",
    "    \n",
    "    # symmetrize and normalize\n",
    "    return rho\n",
    "\n",
    "def Svec():\n",
    "    r\"\"\"\n",
    "    Return \\mathbf{S} = (S^z, S^x, S^y)\n",
    "    \"\"\"\n",
    "    return result\n",
    "\n",
    "def SS():\n",
    "    r\"\"\"\n",
    "    Return \\mathbf{S}_i \\cdot \\mathbf{S}_j\n",
    "    \"\"\"\n",
    "    return result\n",
    "\n",
    "def rot_op():\n",
    "    r\"\"\"\n",
    "    Return unitary -i\\sigma^y\n",
    "    \"\"\"\n",
    "    return unitary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce877580-8ea8-44ba-95ee-8fadfa3c8b75",
   "metadata": {
    "tags": []
   },
   "source": [
    "## v) Test observables\n",
    "\n",
    "Evaluate energy and on-site magnetization of the state\n",
    "\n",
    "    \\langle \\mathbf{S}_i \\cdot \\mathbf{S}_j \\rangle = Tr(rho_2x1 \\mathbf{S}_i \\cdot \\mathbf{\\tilde{S}}_j\n",
    "                                                         \n",
    "    \\langle \\mathbf{S} \\rangle = Tr(rho_2x1 \\mathbf{S}\\otimes\\mathbf{1})\n",
    "                                                         \n",
    "where \\mathbf{\\tilde{S}} is a vector of spin operators appropriately conjugated by unitary rot_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d845fc-a72c-473b-a7da-4770fce35629",
   "metadata": {},
   "outputs": [],
   "source": [
    "def observables(a,C,T):\n",
    "    r\"\"\"\n",
    "    Given boundary tensors C,T and on-site tensor a, evaluate\n",
    "    and return expectation values of NN spin exchange and components of spin vectors\n",
    "    \"\"\"\n",
    "    return e, m\n",
    "\n",
    "# here we recompute environment and evaluate observables for selected state\n",
    "a= torch.asarray(load_from_pepstorch_json_dense(\"j1j2_ipeps_states/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D3_chi_opt108.json\"))\n",
    "\n",
    "\n",
    "chi=18\n",
    "# converge CTMRG and get boundaries\n",
    "\n",
    "# compute and compare observables\n",
    "\n",
    "# you should find \n",
    "# <S.S>_NN, m^z, m^x, m^y\n",
    "# -0.33359724591894546, -0.19644642869439194, 0.28794605609487134, 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "984b5597-57af-46bd-8039-eaf04d451c1d",
   "metadata": {},
   "source": [
    "## vi) Optimization\n",
    "\n",
    "### vi.1) Objective function\n",
    "\n",
    "Lets create a full objective function, which takes on-site tensor a, projected to A_1 irrep\n",
    "of C_{4v} group, and returns the observables as well as converged boundaries C,T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "329350a1-db4d-4c08-88d5-ce7f4adae0c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_c4v_symm_A1(a):\n",
    "    r\"\"\"\n",
    "    Project on-site tensor `a` on :math:`A_1` irrep of :math:`C_{4v}` group.\n",
    "    \"\"\"\n",
    "    # left-right reflection\n",
    "    # up-down reflection\n",
    "    # pi/2 anti-clockwise\n",
    "    # pi/2 clockwise\n",
    "    return a\n",
    "\n",
    "def obj(a):\n",
    "    r\"\"\"\n",
    "    Objective function computes fixed point of CTMRG and then evaluates observables, in particular energy\n",
    "    \"\"\"\n",
    "    # project to A_1 irrep\n",
    "    \n",
    "    # normalize the on-site tensor\n",
    "    \n",
    "    # converge CTMRG\n",
    "    \n",
    "    # evaluate observables and return them together with boundary tensors\n",
    "    return e,m,C,T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b32d46d-5131-4532-a4eb-a3242636e8ac",
   "metadata": {},
   "source": [
    "### vi.2) Gradient descent\n",
    "\n",
    "Starting from a random symmetric tensor, perform gradient descent optimization. \n",
    "To evalute gradients, use reverse-mode automatic differentiation provided by PyTorch.\n",
    "\n",
    "See basics of optimization with PyTorch - see https://pytorch.org/tutorials/beginner/basics/optimization_tutorial.html\n",
    "\n",
    "Experiment with different optimizers - see https://pytorch.org/docs/stable/optim.html\n",
    "\n",
    "Hint: As a generic non-linear optimization, you might get stuck in local minima.\n",
    "Try different random seeds / initializations of tensor a."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5689871e-bda4-4ff0-a372-29792b85f0e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "bond_dim=2\n",
    "chi=16\n",
    "learning_rate=1.0\n",
    "max_opt_iter=100\n",
    "\n",
    "# Take random initial states\n",
    "#\n",
    "# a0= ...\n",
    "\n",
    "# Choose an optimizer\n",
    "#\n",
    "# optimizer= torch.optim.SGD([a0,], lr=learning_rate)\n",
    "optimizer= torch.optim.LBFGS([a0,], lr=1, max_iter=1, max_eval=None, tolerance_grad=1e-07, tolerance_change=1e-09, history_size=100, line_search_fn=None)\n",
    "\n",
    "# Define closure, which evaluates objective and computes the gradient with backpropagation (reverse-mode automatic differentiation).\n",
    "def closure():\n",
    "    optimizer.zero_grad()\n",
    "    e,m,C,T= obj(a0)\n",
    "        \n",
    "    # Backpropagation\n",
    "    e.backward()\n",
    "    \n",
    "    # (optionally) report observables\n",
    "    return e\n",
    "\n",
    "for i in range(max_opt_iter):\n",
    "    # update tensor\n",
    "    optimizer.step(closure)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89d409e6-76ef-4c9a-9587-10c5bd5e4f2b",
   "metadata": {},
   "source": [
    "## vii) What about correlation lengths\n",
    "\n",
    "We have obtained some (optimal?) PEPS approximating GS of AFM Heisenberg model.\n",
    "How large are correlation lengths of such PEPS states ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5346796-f31a-4334-9693-334db4cd0683",
   "metadata": {},
   "outputs": [],
   "source": [
    "# take optimal state, stop gradient tracking (see https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html) \n",
    "# and compute its environment\n",
    "a= a0.detach().clone()\n",
    "chi=18\n",
    "\n",
    "# compute CTMRG\n",
    "\n",
    "# compute correlation length"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dac98193-40fa-4317-821a-fd8ca30bac3d",
   "metadata": {},
   "source": [
    "# Analysis of S=2 AKLT state and spin-1/2 NN RVB\n",
    "\n",
    "    1. get peps-torch\n",
    "    \n",
    "    2. use AKLT S=2 example: examples/akltS2/ctmrg_akltS2_c4v.py\n",
    "       (or see https://github.com/jurajHasik/peps-torch/blob/master/examples/akltS2/ctmrg_akltS2_c4v.py)\n",
    "       \n",
    "    3. use spin-1/2 J1-J2 example: examples/j1j2/ctmrg_j1j2_c4v.py\n",
    "       (or see https://github.com/jurajHasik/peps-torch/blob/master/examples/j1j2/ctmrg_j1j2_c4v.py)\n",
    "\n",
    "Explore spin-spin correlation functions, dimer correlation function, and spectrum of the corner or transfer matrix as a function of chi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34523890-a46c-4e65-8556-0a29e13d1eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "! git clone https://github.com/jurajHasik/peps-torch.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9c478f8-26bd-4c36-bbab-60afb1c39ac2",
   "metadata": {},
   "source": [
    "## i) Evaluate environment and observables of S=2 AKLT state\n",
    "\n",
    "This example reads in the on-site tensor a of AKLT state, converges the CTMRG to obtain\n",
    "environment tensors C,T and then computes various observables:\n",
    "Hamiltonian term, components of magnetization, spin-spin and dimer-dimer correlation functions\n",
    "\n",
    "This states shows a rapid decay of all correlation functions.\n",
    "\n",
    "(https://github.com/jurajHasik/peps-torch/blob/master/examples/akltS2/ctmrg_akltS2_c4v.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81db7552-752d-40f7-91d2-943abe1f220b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# corrf_r dictates range over which correlation functions are evaluated\n",
    "\n",
    "! cd peps-torch; python examples/akltS2/ctmrg_akltS2_c4v.py --instate test-input/AKLT-S2_1x1.in --chi 16 --corrf_r 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9730ee8-04dc-4ecf-afa0-ca92f88dcaca",
   "metadata": {},
   "source": [
    "## ii) Evaluate environment and observables of spin-1/2 NN RVB state\n",
    "\n",
    "This example reads in the on-site tensor a of AKLT state, converges the CTMRG to obtain\n",
    "environment tensors C,T and then computes various observables:\n",
    "energy of J1-J2 Hamiltonina, components of magnetization, spin-spin and dimer-dimer correlation functions\n",
    "\n",
    "This states shows a rapid decay of spin-spin correlations, but increasingly slower decay (with larger chi) of dimer-dimer correlations.\n",
    "\n",
    "    * As chi is increased, the CTMRG starts to take larger amount of iterations to converge, easily beyond 1000\n",
    "    * You can explore how variational energy behaves as frustration J2 is increased \n",
    "    \n",
    "(https://github.com/jurajHasik/peps-torch/blob/master/examples/j1j2/ctmrg_j1j2_c4v.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "616a49fc-42e0-4021-a31f-ccf2792ecf6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cd peps-torch; python examples/j1j2/ctmrg_j1j2_c4v.py --instate test-input/RVB_1x1.in --j2 0.0 --chi 18 --corrf_r 20 --top_n 10 --CTMARGS_ctm_max_iter 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f5f12e-861d-4a6a-80ca-5b71b377e183",
   "metadata": {},
   "source": [
    "# U(1)-symmetric Simple update [peps-torch & YASTN]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7db9416-42f8-42b6-9f80-9e430332126e",
   "metadata": {},
   "source": [
    "## i) AFM Heisenberg model on a square lattice \n",
    "\n",
    "Optimize AFM Heisenberg model on a square lattice using PEPS with 2x2 unit cell\n",
    "    \n",
    "    1. get submodules - YASTN\n",
    "    2. simulate coupled spin-1/2 ladders in isotropic Heisenberg limit using examples/ladders/abelian/SU_ladders_u1.py\n",
    "       (see https://github.com/jurajHasik/peps-torch/blob/master/examples/ladders/abelian/SU_ladders_u1.py)\n",
    "    \n",
    "Starting from mean-field Neel state.\n",
    "\n",
    "You can compare the energies between these D=4 optimizations and the gradient ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7696880c-3026-475e-a153-5a5450b3df17",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cd peps-torch; git submodule update --init --recursive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f8c022-7aac-4fa5-a4b5-03bb0de7c417",
   "metadata": {},
   "outputs": [],
   "source": [
    "# alpha=1 represents (isotropic) 2D AFM Heisenberg model on a square lattice\n",
    "\n",
    "! cd peps-torch; python examples/ladders/abelian/SU_ladders_u1.py --alpha 1. --bond_dim 4 --chi 16 --instate test-input/abelian/NEEL_D1_2x2_abelian-U1_state.json"
   ]
  }
 ],
 "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.10.9"
  },
  "toc-autonumbering": false,
  "toc-showcode": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}