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