import warnings
from collections import OrderedDict
from itertools import chain
import json
import math
import torch
import config as cfg
try:
import torch
from ipeps.ipess_kagome import IPESS_KAGOME_GENERIC
except ImportError as e:
warnings.warn("torch not available", Warning)
import yast.yast as yast
import ipeps.ipeps_kagome_abelian as ipeps_kagome
from ipeps.tensor_io import *
[docs]class IPESS_KAGOME_GENERIC_ABELIAN(ipeps_kagome.IPEPS_KAGOME_ABELIAN):
def __init__(self, settings, ipess_tensors, build_sites=True,
peps_args=cfg.peps_args, global_args=cfg.global_args):
r"""
:param settings: YAST configuration
:type settings: NamedTuple or SimpleNamespace (TODO link to definition)
:param ipess_tensors: dictionary of five tensors, which make up Kagome iPESS
ansatz
:param peps_args: ipeps configuration
:param global_args: global configuration
:type ipess_tensors: dict(str, yast.Tensor)
:type peps_args: PEPSARGS
:type global_args: GLOBALARGS
iPESS ansatz for Kagome lattice composes five tensors, specified
by dictionary::
ipess_tensors = {'T_u': yast.Tensor, 'T_d': ..., 'B_a': ..., 'B_b': ..., 'B_c': ...}
into a single rank-5 on-site tensor of parent IPEPS. These iPESS tensors can
be accessed through member ``ipess_tensors``.
The ``'B_*'`` are rank-3 tensors, with index structure [p,i,j] where the first index `p`
is for physical degree of freedom, while indices `i` and `j` are auxiliary.
These bond tensors reside on corners shared between different triangles of Kagome lattice.
Bond tensors are connected by rank-3 trivalent tensors ``'T_u'``, ``'T_d'`` on up and down triangles
respectively. Trivalent tensors have only auxiliary indices.
The on-site tensors of corresponding iPEPS is obtained by the following contraction::
2(d) 2(c) (-)a
\ / rot. pi |
0(w)==B_a B_b==0(v) clockwise (-)b--\
\ / => \
1(l) 1(k) s0--s2--d(+)
2(l) 1(k) | /
\ / |/ <- DOWN_T
T_d s1
| |
0(j) c(+)
1(j)
|
B_c==0(u)
|
2(i)
0(i)
|
T_u
/ \
1(a) 2(b)
where the signature convetions of trivalent and bond tensors are as follows::
(-)\ /(-) (-)| |(-)
T_d , T_u , (+)--B--(+)
|(-) (-)/ \(-)
This choice guarantees that the signature of on-site tensor of equivalent single-site
IPEPS_KAGOME_ABELIAN is compatible.
By construction, the degrees of freedom on down triangle are all combined into
a single on-site tensor of iPEPS. Instead, DoFs on the upper triangle have
to be accessed by construction of 2x2 patch (which is then embedded into environment)::
C T T C
a a
| |
T b--\ b--\
\ / \
s0--s2--d s0--s2--d T
| / | /
|/ |/
s1 s1
| |
c c
/ /
a a
| |
T b--\ b--\
\ / \
s0--s2--d s0--s2--d T
| / | /
|/ |/
s1 s1
| |
c c
C T T C
"""
#TODO verification?
assert ipess_tensors['T_d'].get_signature()==(-1,-1,-1),"Unexpected signature"
assert ipess_tensors['T_u'].get_signature()==(-1,-1,-1),"Unexpected signature"
assert ipess_tensors['B_a'].get_signature()==(-1,1,1),"Unexpected signature"
assert ipess_tensors['B_b'].get_signature()==(-1,1,1),"Unexpected signature"
assert ipess_tensors['B_c'].get_signature()==(-1,1,1),"Unexpected signature"
self.ipess_tensors= ipess_tensors
if build_sites:
sites= self.build_onsite_tensors()
else:
sites= {}
super().__init__(settings, sites, lX=1, lY=1, peps_args=peps_args,
global_args=global_args)
[docs] def get_parameters(self):
r"""
:return: variational parameters of IPESS_KAGOME_GENERIC_ABELIAN
:rtype: iterable
This function is called by optimizer to access variational parameters of the state.
In this case member ``ipess_tensors``.
"""
return list(self.ipess_tensors[ind].data for ind in self.ipess_tensors)
[docs] def get_checkpoint(self):
r"""
:return: serializable representation of IPESS_KAGOME_GENERIC_ABELIAN state
:rtype: dict
Return dict containing serialized ``ipess_tensors`` (block-sparse) tensors.
The individual blocks are serialized into Numpy ndarrays.
This function is called by optimizer to create checkpoints during
the optimization process.
"""
return {ind: self.ipess_tensors[ind].save_to_dict() for ind in self.ipess_tensors}
[docs] def load_checkpoint(self, checkpoint_file):
r"""
:param checkpoint_file: path to checkpoint file
:type checkpoint_file: str or file object
Initializes IPESS_KAGOME_GENERIC_ABELIAN from checkpoint file.
"""
checkpoint= torch.load(checkpoint_file, map_location=self.device)
self.ipess_tensors= {ind: yast.load_from_dict(config= self.engine, d=t_dict_repr) \
for ind,t_dict_repr in checkpoint["parameters"].items()}
for t in self.ipess_tensors.values(): t.requires_grad_(False)
self.sites = self.build_onsite_tensors()
self.sync_precomputed()
[docs] def to_dense(self, peps_args=cfg.peps_args, global_args=cfg.global_args):
r"""
:return: returns equivalent dense state with all on-site tensors in their dense
representation on PyTorch backend.
:rtype: IPESS_KAGOME_GENERIC_ABELIAN
Create an IPESS_KAGOME_GENERIC state with all ``ipess_tensors`` as dense possesing
no explicit block structure (symmetry). This operations preserves gradients
on the returned dense state.
"""
ipess_tensors_dense= {t_id: t.to_dense() for t_id,t in self.ipess_tensors.items()}
state_dense= IPESS_KAGOME_GENERIC(ipess_tensors_dense,peps_args=peps_args,\
global_args=global_args)
return state_dense
[docs] def build_onsite_tensors(self):
r"""
:return: elementary unit cell of underlying IPEPS
:rtype: dict[tuple(int,int): yast.Tensor]
Build on-site tensor of corresponding IPEPS_KAGOME_ABELIAN be
performing following contraction of ``ipess_tensors``::
2(-7) 2(-6) (-)(-4)
\ / rot. pi |
0(-3)==B_a B_b==0(-2) clockwise (-)(-5)--\
\ / => \
1(3) 1(2) s0(-1)--s2(-3)--(-7)(+)
2(3) 1(2) | /
\ / |/ <- DOWN_T
T_d s1(-2)
| |
0(1) (-6)(+)
1(1)
|
B_c==0(-1)
|
2(0)
0(0)
|
T_u
/ \
1(-4) 2(-5)
"""
A= yast.ncon([self.ipess_tensors['T_u'], self.ipess_tensors['B_c'],\
self.ipess_tensors['T_d'], self.ipess_tensors['B_b'], self.ipess_tensors['B_a']],\
[[1,-3,-4], [-0,2,1], [2,3,4], [-1,3,-5], [-2,4,-6]])
# [[0,-4,-5], [-1,1,0], [1,2,3], [-2,2,-6], [-3,3,-7]]
A= A.fuse_legs(axes=((0,1,2),3,4,5,6))
A= A/A.norm(p='inf')
sites= {(0, 0): A}
return sites
[docs] def add_noise(self, noise=0, peps_args=cfg.peps_args):
r"""
:param noise: magnitude of the noise
:type noise: float
:return: a copy of state with noisy ``ipess_tensors``. For default value of
``noise`` being zero ``self`` is returned.
:rtype: IPESS_KAGOME_GENERIC_ABELIAN
Create a new state by adding random uniform noise with magnitude ``noise`` to all
copies of ``ipess_tensors``. The noise is added to all allowed blocks making up
the individual tensors. If noise is 0, returns ``self``.
"""
if noise==0: return self
ipess_tensors= {}
for ind,t in self.ipess_tensors.items():
ts, Ds= t.get_leg_charges_and_dims(native=True)
t_noise= yast.rand(config= t.config, s=t.s, n=t.n, t=ts, D=Ds, isdiag=t.isdiag)
ipess_tensors[ind]= t + noise*t_noise
return IPESS_KAGOME_GENERIC_ABELIAN(self.engine, ipess_tensors,\
peps_args=peps_args)
[docs] def generate_weights(self):
r"""
:return: dictionary of weights on non-equivalent bonds of the iPESS ansatz
:rtype: dict( str: yast.Tensor )
Generate a set of identity tensors, weights W, for each non-equivalent link
in the iPESS ansatz::
2(d) 2(c) (-)a
\ / rot. pi |
0(w)==B_a B_b==0(v) clockwise (-)b--\
\ / => \
1(l) 1(k) s0--s2--d(+)
W_down_a W_down_b | /
2(l) 1(k) | /
\ / |/ <- DOWN_T
T_d s1
| |
0(j) c(+)
W_down_c
1(j)
|
B_c==0(u)
|
2(i)
W_up_c
0(i)
|
T_u
/ \
1(a) 2(b)
W_up_b W_up_a
"""
weights=dict()
for w_id, ind1, ind2 in [
('lambda_up_c', ('T_u',0), ('B_c',2)),
('lambda_up_a', ('T_u',2), ('B_a',2)),
('lambda_up_b', ('T_u',1), ('B_b',2)),
('lambda_dn_c', ('T_d',0), ('B_c',1)),
('lambda_dn_a', ('T_d',2), ('B_a',1)),
('lambda_up_b', ('T_u',1), ('B_b',1))
]:
W= yast.match_legs( tensors=[self.ipess_tensors[ind1[0]], self.ipess_tensors[ind2[0]]],\
legs=[ ind1[1] , ind2[1] ], isdiag=True ).flip_signature()
weights[w_id]= W
return weights
def __str__(self):
print(f"lX x lY: {self.lX} x {self.lY}")
for t_id,t in self.ipess_tensors.items():
print(f"{t_id}")
print(f"{t}")
print("")
for nid,coord,site in [(t[0], *t[1]) for t in enumerate(self.sites.items())]:
print(f"a{nid} {coord}: {site}")
# show tiling of a square lattice
coord_list = list(self.sites.keys())
mx, my = 3*self.lX, 3*self.lY
label_spacing = 1+int(math.log10(len(self.sites.keys())))
for y in range(-my,my):
if y == -my:
print("y\\x ", end="")
for x in range(-mx,mx):
print(str(x)+label_spacing*" "+" ", end="")
print("")
print(f"{y:+} ", end="")
for x in range(-mx,mx):
print(f"a{coord_list.index(self.vertexToSite((x,y)))} ", end="")
print("")
return ""
[docs] def write_to_file(self, outputfile, tol=None, normalize=False):
write_ipess_kagome_generic(self, outputfile, tol=tol, normalize=normalize)
# TODO verify global_args.device is compatible with device in settings
[docs]def read_ipess_kagome_generic(jsonfile, settings, peps_args=cfg.peps_args,\
global_args=cfg.global_args):
r"""
:param jsonfile: input file describing IPESS_KAGOME_GENERIC_ABELIAN in json format
:param settings: YAST configuration
:type settings: NamedTuple or SimpleNamespace (TODO link to definition)
:param peps_args: ipeps configuration
:param global_args: global configuration
:type jsonfile: str or Path object
:type peps_args: PEPSARGS
:type global_args: GLOBALARGS
:return: wavefunction
:rtype: IPESS_KAGOME_GENERIC_ABELIAN
Read IPESS_KAGOME_GENERIC_ABELIAN from file.
"""
with open(jsonfile) as j:
raw_state = json.load(j)
# Loop over non-equivalent tensor,coeffs pairs in the unit cell
ipess_tensors= OrderedDict()
assert "ipess_tensors" in raw_state.keys(),"Missing \"ipess_tensors\""
assert set(('T_u','T_d','B_a','B_b','B_c'))==set(list(raw_state["ipess_tensors"].keys())),\
"missing some ipess tensors"
for key,t in raw_state["ipess_tensors"].items():
# depending on the "format", read the bare tensor
assert "format" in t.keys(), "\"format\" not specified"
if t["format"]=="abelian":
X= read_json_abelian_tensor_legacy(t, settings)
# move all tensors to desired dtype and device
# TODO improve
if X.is_complex() and X.yast_dtype in ["float64", "float32"]:
warnings.warn("Downcasting complex tensor to real", RuntimeWarning)
ipess_tensors[key]= X.to(dtype=global_args.dtype, device=global_args.device)
else:
raise Exception("Unsupported format "+t["format"])
state = IPESS_KAGOME_GENERIC_ABELIAN(settings, ipess_tensors, peps_args=peps_args, \
global_args=global_args)
return state
[docs]def write_ipess_kagome_generic(state, outputfile, tol=None, normalize=False,\
peps_args=cfg.peps_args, global_args=cfg.global_args):
r"""
:param state: wavefunction to write out in JSON format
:param outputfile: target file
:param tol: minimum magnitude of tensor elements which are written out
:param normalize: if True, ipess tensors are normalized before writing
:type state: IPESS_KAGOME_GENERIC_ABELIAN
:type ouputfile: str or Path object
:type tol: float
:type normalize: bool
Write IPESS_KAGOME_GENERIC_ABELIAN to file.
"""
json_state = dict({"lX": state.lX, "lY": state.lY, \
"ipess_tensors": {}})
# write list of considered elementary tensors
for key, t in state.ipess_tensors.items():
tmp_t= t/t.norm(p='inf') if normalize else t
json_state["ipess_tensors"][key]= serialize_abelian_tensor_legacy(tmp_t)
with open(outputfile, 'w') as f:
json.dump(json_state, f, indent=4, separators=(',', ': '), cls=NumPy_Encoder)