Source code for ctm.generic_abelian.env_abelian

import warnings
import config as cfg
import yast.yast as yast
try:
    import torch
    from ctm.generic.env import ENV
except ImportError as e:
    warnings.warn("torch not available", Warning)

[docs]class ENV_ABELIAN(): def __init__(self, chi=1, state=None, settings=None, init=False,\ init_method=None, ctm_args=cfg.ctm_args, global_args=cfg.global_args): r""" :param chi: environment bond dimension :math:`\chi` :param state: wavefunction :param settings: YAST configuration :type settings: NamedTuple or SimpleNamespace (TODO link to definition) :param init: initialize environment tensors :type init: bool :param init_method: choice of environment initialization. See :meth:`init_env`. :type init_method: str :param ctm_args: CTM algorithm configuration :param global_args: global configuration :type chi: int :type state: IPEPS_ABELIAN :type ctm_args: CTMARGS :type global_args: GLOBALARGS For each pair of (vertex, on-site tensor) in the elementary unit cell of ``state``, create corresponding environment tensors: Half-row/column tensors T's and corner tensors C's. The corner tensors have dimensions :math:`\chi \times \chi` and the half-row/column tensors have dimensions :math:`\chi \times \chi \times D^2` (D might vary depending on the corresponding dimension of on-site tensor). The environment of each double-layer tensor (A) is composed of eight different tensors:: y\x -1 0 1 -1 C T C 0 T A T 1 C T C The individual tensors making up the environment of a site are defined by four directional vectors :math:`d = (x,y)_{\textrm{environment tensor}} - (x,y)_\textrm{A}` as follows:: C(-1,-1) T (1,-1)C |(0,-1) T--(-1,0)--A(0,0)--(1,0)--T |(0,1) C(-1,1) T (1,1)C These environment tensors of some ENV object ``e`` are accesed through its members ``C`` and ``T`` by providing a tuple of coordinates and directional vector to the environment tensor:: coord=(0,0) # tuple(x,y) identifying vertex on the square lattice rel_dir_vec_C=(-1,-1) # tuple(rx,ry) identifying one of the four corner tensors rel_dir_vec_T=(-1,0) # tuple(rx,ry) identifying one of the four half-row/column tensors C_upper_left= e.C[(coord,rel_dir_vec_C)] # return upper left corner tensor of site at coord T_left= e.T[(coord,rel_dir_vec_T)] # return left half-row tensor of site at coord The index-position convention is as follows: Start from the index in the **direction "up"** <=> (0,-1) and continue **anti-clockwise**. The reference symmetry signatures are shown on the right:: C--1 0--T--2 0--C C(+1) (-1)T(+1) (-1)C | | | (+1) (+1) (+1) 0 1 1 0 0 | | (-1) (-1) T--2 1--T T(+1) (-1)T | | (+1) (+1) 1 2 0 0 0 | | | (-1) (-1) (-1) C--1 1--T--2 1--C C(+1) (-1)T(+1) (-1)C .. note:: The structure of fused double-layer legs, which are carried by T-tensors, is obtained by fusing on-site tensor (`ket`) with its conjugate (`bra`). The leg of `ket` always preceeds `bra` when fusing. """ if state: self.engine= state.engine self.dtype= state.dtype self.nsym = state.nsym self.sym= state.sym elif settings: self.engine= settings self.dtype= settings.default_dtype self.nsym = settings.sym.NSYM self.sym= settings.sym.SYM_ID else: raise RuntimeError("Either state or settings must be provided") self.device= global_args.device self.chi= chi # initialize environment C,T dictionaries self.C = dict() self.T = dict() if init or init_method: if not init_method: init_method= ctm_args.ctm_env_init_type if state and init_method in ["CTMRG"]: init_env(state, self, init_method) else: raise RuntimeError("Cannot perform initialization for desired"\ +" ctm_env_init_type "+init_method+"."\ +" Missing state.") def __str__(self): s=f"ENV_abelian chi={self.chi}\n" s+=f"dtype {self.dtype}\n" s+=f"device {self.device}\n" s+=f"nsym {self.nsym}\n" s+=f"sym {self.sym}\n" if len(self.C)==0: s+="C is empty\n" for cr,t in self.C.items(): s+=f"C({cr[0]} {cr[1]}): {t}\n" if len(self.T)==0: s+="T is empty\n" for cr,t in self.T.items(): s+=f"T({cr[0]} {cr[1]}): {t}\n" return s def extend(self, new_chi, ctm_args=cfg.ctm_args, global_args=cfg.global_args): raise NotImplementedError
[docs] def to_dense(self, state, ctm_args=cfg.ctm_args, global_args=cfg.global_args): r""" :param state: abelian-symmetric iPEPS :type state: IPEPS_ABELIAN :return: returns equivalent of the environment with all C,T tensors in their dense representation on PyTorch backend. :rtype: ENV Create a copy of environment with all on-site tensors as dense possesing no explicit block structure (symmetry). The virtual spaces of on-site tensors in ``state`` are added to corresponding spaces of environment's T tensors to guarantee consistency. .. note:: This operations preserves gradients on the returned dense environment. """ vts= state.vertexToSite dir_to_leg= {(0,-1): 1, (0,1): 0, (-1,0): 2, (1,0): 1} C_lss= { cid: dict() for cid in self.C.keys() } T_lss= dict() # 0) compute correct leg structure of T's. Unfuse the pair of # auxiliary legs connecting T's to on-site tensors. Merge the # environment virtual spaces on connected legs tmp_T= { tid: t.unfuse_legs(dir_to_leg[tid[1]]) for tid,t in self.T.items() } for tid in tmp_T.keys(): t_xy,t_dir= tid if t_dir==(0,-1): #UP # all legs of current T # 0--T(x-1,y)--3 0--T(x,y)--3 0--T(x+1,y)--3 # 1,2 1,2 1,2 T_lss[tid]= { 1: tmp_T[tid].get_legs(axis=1), 2: tmp_T[tid].get_legs(axis=2), \ 0: yast.leg_union(tmp_T[(vts((t_xy[0]-1, t_xy[1])), t_dir)].get_legs(axis=3),\ self.C[((t_xy),(-1,-1))].get_legs(axis=1)).conj(),\ 3: yast.leg_union(tmp_T[(vts((t_xy[0]+1, t_xy[1])), t_dir)].get_legs(axis=0),\ self.C[((t_xy),(1,-1))].get_legs(axis=0)).conj() } # upper-left corner # C--1 # 0 C_lss[(tid[0],(-1,-1))][1]= T_lss[tid][0].conj() # upper-right corner # 0--C # 1 C_lss[(tid[0],(1,-1))][0]= T_lss[tid][3].conj() elif t_dir==(0,1): #DOWN # 0,1 0,1 0,1 # 2--T(x-1,y)--3 2--T(x,y)--3 2--T(x+1,y)--3 T_lss[tid]= { 0: tmp_T[tid].get_legs(axis=0), 1: tmp_T[tid].get_legs(axis=1), \ 2: yast.leg_union(tmp_T[(vts((t_xy[0]-1, t_xy[1])), t_dir)].get_legs(axis=3),\ self.C[((t_xy),(-1,1))].get_legs(axis=1)).conj(),\ 3: yast.leg_union(tmp_T[(vts((t_xy[0]+1, t_xy[1])), t_dir)].get_legs(axis=2),\ self.C[((t_xy),(1,1))].get_legs(axis=1)).conj() } # lower-left corner # 0 # C--1 C_lss[(tid[0],(-1,1))][1]= T_lss[tid][2].conj() # lower-right corner # 0 # 1--C C_lss[(tid[0],(1,1))][1]= T_lss[tid][3].conj() elif t_dir==(-1,0): #LEFT # 0 # T--2,3 # 1 T_lss[tid]= { 2: tmp_T[tid].get_legs(axis=2), 3: tmp_T[tid].get_legs(axis=3),\ 0: yast.leg_union(tmp_T[(vts((t_xy[0], t_xy[1]-1)), t_dir)].get_legs(axis=1),\ self.C[((t_xy),(-1,-1))].get_legs(axis=0)).conj(),\ 1: yast.leg_union(tmp_T[(vts((t_xy[0], t_xy[1]+1)), t_dir)].get_legs(axis=0),\ self.C[((t_xy),(-1,1))].get_legs(axis=0)).conj() } # upper-left corner # C--1 # 0 C_lss[(tid[0],(-1,-1))][0]= T_lss[tid][0].conj() # lower-left corner # 0 # C--1 C_lss[(tid[0],(-1,1))][0]= T_lss[tid][1].conj() elif t_dir==(1,0): #RIGHT # 0 # 1,2--T # 3 T_lss[tid]= { 1: tmp_T[tid].get_legs(axis=1), 2: tmp_T[tid].get_legs(axis=2),\ 0: yast.leg_union(tmp_T[(vts((t_xy[0], t_xy[1]-1)), t_dir)].get_legs(axis=3),\ self.C[((t_xy),(1,-1))].get_legs(axis=1)).conj(),\ 3: yast.leg_union(tmp_T[(vts((t_xy[0], t_xy[1]+1)), t_dir)].get_legs(axis=0),\ self.C[((t_xy),(1,1))].get_legs(axis=0)).conj() } # upper-right corner # 0--C # 1 C_lss[(tid[0],(1,-1))][1]= T_lss[tid][0].conj() # lower-right corner # 0 # 1--C C_lss[(tid[0],(1,1))][0]= T_lss[tid][3].conj() else: raise RuntimeError("Invalid T-tensor id "+str(tid)) # 1) convert to dense representation. Reshape T's into double-layer form C_torch= {cid: c.to_dense(legs=C_lss[cid]) for cid,c in self.C.items()} T_torch= dict() for tid,t in tmp_T.items(): t_xy,t_dir= tid t= t.to_dense(legs=T_lss[tid]) if t_dir==(0,-1): T_torch[tid]= t.view(t.size(0), t.size(1)*t.size(2), t.size(3)) elif t_dir==(0,1): T_torch[tid]= t.view(t.size(0)*t.size(1), t.size(2), t.size(3)) elif t_dir==(-1,0): T_torch[tid]= t.view(t.size(0), t.size(1), t.size(2)*t.size(3)) else: T_torch[tid]= t.view(t.size(0), t.size(1)*t.size(2), t.size(3)) max_chi= max(self.chi, max([max(c.size()) for c in C_torch.values()])) if max_chi>self.chi: warnings.warn("Increasing chi. Equivalent chi ("+str(max_chi)+") of symmetric"\ +" environment is higher than original chi ("+str(self.chi)+").", Warning) # 2) Fill the dense environment with dimension chi by dense representations of # symmetric environment tensors env_torch= ENV(max_chi, ctm_args=ctm_args, global_args=global_args) for cid,c in C_torch.items(): env_torch.C[cid]= torch.zeros(max_chi,max_chi,dtype=c.dtype,device=c.device) env_torch.C[cid][:c.size(0),:c.size(1)]= c for tid,t in T_torch.items(): t_site, t_dir= tid if t_dir==(0,-1): env_torch.T[tid]= torch.zeros(max_chi,t.size(1),max_chi,dtype=t.dtype,device=t.device) env_torch.T[tid][:t.size(0),:,:t.size(2)]= t elif t_dir==(-1,0): env_torch.T[tid]= torch.zeros(max_chi,max_chi,t.size(2),dtype=t.dtype,device=t.device) env_torch.T[tid][:t.size(0),:t.size(1),:]= t elif t_dir==(0,1): env_torch.T[tid]= torch.zeros(t.size(0),max_chi,max_chi,dtype=t.dtype,device=t.device) env_torch.T[tid][:,:t.size(1),:t.size(2)]= t elif t_dir==(1,0): env_torch.T[tid]= torch.zeros(max_chi,t.size(1),max_chi,dtype=t.dtype,device=t.device) env_torch.T[tid][:t.size(0),:,:t.size(2)]= t return env_torch
[docs] def clone(self): r""" :return: returns a clone of the environment :rtype: ENV_ABELIAN Create a clone of environment with all tensors (their blocks) attached to the computational graph. .. note:: This operation preserves gradient tracking. """ e= ENV_ABELIAN(self.chi, settings=self.engine) e.C= {cid: c.clone() for cid,c in self.C.items()} e.T= {tid: t.clone() for tid,t in self.T.items()} return e
[docs] def detach(self): r""" :return: returns a view of the environment with all C,T tensors detached from computational graph. :rtype: ENV_ABELIAN In case of using PyTorch backend, get a detached "view" of the environment. See `torch.Tensor.detach <https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html>`_. .. note:: This operation does not preserve gradient tracking. """ e= ENV_ABELIAN(self.chi, settings=self.engine) e.C= {cid: c.detach() for cid,c in self.C.items()} e.T= {tid: t.detach() for tid,t in self.T.items()} return e
def detach_(self): for c in self.C.values(): c.detach() for t in self.T.values(): t.detach()
[docs]def init_env(state, env, init_method=None, ctm_args=cfg.ctm_args): """ :param state: wavefunction :param env: CTM environment :param init_method: desired initialization method :param ctm_args: CTM algorithm configuration :type state: IPEPS_ABELIAN :type env: ENV_ABELIAN :type init_method: str :type ctm_args: CTMARGS Initializes the environment `env` according to one of the supported options specified by :class:`CTMARGS.ctm_env_init_type <config.CTMARGS>` * CTMRG - tensors C and T are built from the on-site tensors of `state` """ if not init_method: init_method= ctm_args.ctm_env_init_type if init_method=='CTMRG': init_from_ipeps_pbc(state, env, ctm_args.verbosity_initialization) else: raise ValueError("Invalid environment initialization: "+init_method)
def init_from_ipeps_pbc(state, env, verbosity=0): if verbosity>0: print("ENV: init_from_ipeps_pbc") # corners for coord,site in state.sites.items(): # Left-upper corner # # i = C--(1,3)->1(+1) # j--a*--4->3 | # /\ (0,2)->0(+1) # 2<-3 m # \ i # j--a--4->1 # / # 3->0 vec = (-1,-1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,1,2), (0,1,2)), conj=(0,1)) # mijef,mijab->efab; efab->eafb a= a.fuse_legs( axes=((0,2),(1,3)) ) a= a/a.norm(p='inf') env.C[(coord,vec)]= a # right-upper corner # # i = (-1)0<-(0,2)--C # 2<-2--a*--j | # /\ (+1)1<-(1,3) # 3<-3 m # \ i # 0<-2--a--j # / # 1<-3 vec = (1,-1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,1,4), (0,1,4)), conj=(0,1)) # miefj,miabj->efab; efab->eafb a= a.fuse_legs( axes=((0,2),(1,3)) ) a= a/a.norm(p='inf') env.C[(coord,vec)]=a # right-lower corner # # 1->1 = (-1)0<-(0,2) # 3<-2--a*--j | # /\ (-1)1<-(1,3)--C # i m # \ 1->0 # 1<-2--a--j # / # i vec = (1,1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,3,4), (0,3,4)), conj=(0,1)) # miefj,miabj->efab; efab->eafb a= a.fuse_legs( axes=((0,2),(1,3)) ) a= a/a.norm(p='inf') env.C[(coord,vec)]=a # left-lower corner # # 1->2 = (0,2)->2(-1) # i--a*--4->3 | # /\ C--(1,3)->1(+1) # j m # \ 1->0 # i--a--4->1 # / # j vec = (-1,1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,2,3), (0,2,3)), conj=(0,1)) # miefj,miabj->efab; efab->eafb a= a.fuse_legs( axes=((0,2),(1,3)) ) a= a/a.norm(p='inf') env.C[(coord,vec)]=a # half-row/-column transfer tensor for coord,site in state.sites.items(): # upper transfer matrix # # i = (-1)0<-(0,3)--T--(2,5)->2(+1) # 3<-2--a*--4->5 | # /\ (+1)1<-(1,4) # 4<-3 m # \ i # 0<-2--a--4->2 # / # 1<-3 vec = (0,-1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,1), (0,1)), conj=(0,1)) # miefg,miabc->efgabc ; efgabc->eafbgc a= a.fuse_legs( axes=((0,3),(1,4),(2,5)) ) a= a/a.norm(p='inf') env.T[(coord,vec)]=a # left transfer matrix # # 1->3 = (0,3)->0(-1) # i--a*--4->5 T--(2,5)->2(+1) # /\ (1,4)->1(+1) # 4<-3 m # \ 1->0 # i--a--4->2 # / # 1<-3 vec = (-1,0) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,2), (0,2)), conj=(0,1)) # meifg,maibc->efgabc ; efgabc->eafbgc a= a.fuse_legs( axes=((0,3),(1,4),(2,5)) ) a= a/a.norm(p='inf') env.T[(coord,vec)]=a # lower transfer matrix # # 1->3 = (-1)0<-(0,3) # 4<-2--a*--4->5 | # /\ (-1)1<-(1,4)--T--(2,5)->2(+1) # i m # \ 1->0 # 1<-2--a--4->2 # / # i vec = (0,1) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,3), (0,3)), conj=(0,1)) # mefig,mabic->efgabc; efgabc->eafbgc a= a.fuse_legs( axes=((0,3),(1,4),(2,5)) ) a= a/a.norm(p='inf') env.T[(coord,vec)]=a # right transfer matrix # # 1->3 = (-1)0<-(0,3) # 4<-2--a*--i (-1)1<-(1,4)--T # /\ (+1)2<-(2,5) # 5<-3 m # \ 1->0 # 1<-2--a--i # / # 2<-3 vec = (1,0) A = state.site((coord[0]+vec[0],coord[1]+vec[1])) a= A.tensordot(A, ((0,4), (0,4)), conj=(0,1)) # mefig,mabic->efgabc; efgabc->eafbgc a= a.fuse_legs( axes=((0,3),(1,4),(2,5)) ) a= a/a.norm(p='inf') env.T[(coord,vec)]=a