Untitled

""" Test the expectation values of spin 1/2 fermions with analytical values of fermi sea """
import numpy as np
import pytest
import logging
import yastn
import yastn.tn.fpeps as fpeps
from yastn.tn.fpeps.operators.gates import gates_hopping, gate_local_Hubbard
from yastn.tn.fpeps.evolution import evolution_step_, gates_homogeneous
from yastn.tn.fpeps import initialize_diagonal_basis
from yastn.tn.fpeps.ctm import nn_exp_dict, ctmrg, one_site_dict

try:
    from .configs import config_U1xU1_R_fermionic as cfg
    # cfg is used by pytest to inject different backends and divices
except ImportError:
    from configs import config_U1xU1_R_fermionic as cfg

sym = 'U1xU1xZ2'
opt = yastn.operators.SpinfulFermions(sym=sym, backend=cfg.backend, default_device=cfg.default_device)
fid, fc_up, fc_dn, fcdag_up, fcdag_dn = opt.I(), opt.c(spin='u'), opt.c(spin='d'), opt.cp(spin='u'), opt.cp(spin='d')
n_up = fcdag_up @ fc_up
n_dn = fcdag_dn @ fc_dn
n_int = n_up @ n_dn
h_up, h_dn = fc_up @ fcdag_up, fc_dn @ fcdag_dn
nn_up, nn_dn, nn_do, nn_hole = n_up @ h_dn, n_dn @ h_up, n_up @ n_dn, h_up @ h_dn
projectors = [nn_up, nn_dn, nn_do, nn_hole]
lattice = 'square'
boundary = 'obc'
xx = 4
yy = 4
dbeta = 0.00625
coeff = 0.25 # for purification; 0.5 for ground state calculation and 1j*0.5 for real-time evolution
dims = (xx, yy)
tot_sites = (xx * yy)
mu_up, mu_dn = 0, 0 # chemical potential
t_up, t_dn = 1, 1 # hopping amplitude
U = 8
net = fpeps.Lattice(lattice, dims, boundary)  # shape = (rows, columns)
step = 'one-step'
tr_mode = 'optimal'
trotter_step = coeff * dbeta  

GA_nn_up, GB_nn_up = gates_hopping(t_up, trotter_step, fid, fc_up, fcdag_up)
GA_nn_dn, GB_nn_dn = gates_hopping(t_dn, trotter_step, fid, fc_dn, fcdag_dn)
g_loc = gate_local_Hubbard(mu_up, mu_dn, U, trotter_step, fid, fc_up, fc_dn, fcdag_up, fcdag_dn)
g_nn = [(GA_nn_up, GB_nn_up), (GA_nn_dn, GB_nn_dn)]

out_init = {1: {(0, 0): 0, (0, 1): 1, (0, 2): 0, (0, 3): 1,(1, 0): 1, (1, 1): 2, (1, 2): 1, (1, 3): 0,(2, 0): 0, (2, 1): 3, (2, 2): 3, (2, 3): 1,(3, 0): 3, (3, 1): 0, (3, 2): 1, (3, 3): 0},
            2: {(0, 0): 0, (0, 1): 1, (0, 2): 1, (0, 3): 1,(1, 0): 1, (1, 1): 3, (1, 2): 2, (1, 3): 0,(2, 0): 0, (2, 1): 3, (2, 2): 3, (2, 3): 0,(3, 0): 1, (3, 1): 0, (3, 2): 1, (3, 3): 0},
            3: {(0, 0): 0, (0, 1): 1, (0, 2): 1, (0, 3): 1,(1, 0): 0, (1, 1): 3, (1, 2): 3, (1, 3): 0,(2, 0): 0, (2, 1): 2, (2, 2): 3, (2, 3): 1,(3, 0): 1, (3, 1): 0, (3, 2): 1, (3, 3): 0}}

initial_pattern = 2
beta_end = 8
time_steps = round(beta_end / dbeta)

for D in range(4,11):

    print('bond dimension: ', D)
    opts_svd_ntu = {'D_total': D, 'tol_block': 1e-15}
    master_dict = {}
    psi = initialize_diagonal_basis(projectors, net, out_init[initial_pattern])
    gates = gates_homogeneous(psi, g_nn, g_loc)
    for nums in range(time_steps):
        beta = (nums + 1) * dbeta
        print(beta)
        
        psi, _ =  evolution_step_(psi, gates, step, tr_mode, env_type='NTU', opts_svd=opts_svd_ntu) 
        if beta % 1 == 0:  # Save psi at beta intervals of 1
            print(beta)
            chi = 40 # environmental bond dimension
            tol = 1e-10 # truncation of singular values of CTM projectors
            max_sweeps=50 
            tol_exp = 1e-7   # difference of some observable must be lower than tolernace
            ops = {'cdagc_up': {'l': fcdag_up, 'r': fc_up},
                   'ccdag_up': {'l': fc_up, 'r': fcdag_up},
                   'cdagc_dn': {'l': fcdag_dn, 'r': fc_dn},
                   'ccdag_dn': {'l': fc_dn, 'r': fcdag_dn}}

            cf_energy_old = 0
            opts_svd_ctm = {'D_total': chi, 'tol': tol}
            tot_energy_old = 0
            for step in ctmrg(psi, max_sweeps, iterator_step=1, AAb_mode=0, opts_svd=opts_svd_ctm):
        
                assert step.sweeps % 1 == 0 # stop every 2nd step as iteration_step=2
                obs_hor, obs_ver =  nn_exp_dict(psi, step.env, ops)
                cdagc_up = (sum(abs(val) for val in obs_hor.get('cdagc_up').values()) + sum(abs(val) for val in obs_ver.get('cdagc_up').values()))
                ccdag_up = (sum(abs(val) for val in obs_hor.get('ccdag_up').values()) + sum(abs(val) for val in obs_ver.get('ccdag_up').values()))
                cdagc_dn = (sum(abs(val) for val in obs_hor.get('cdagc_dn').values()) + sum(abs(val) for val in obs_ver.get('cdagc_dn').values()))
                ccdag_dn = (sum(abs(val) for val in obs_hor.get('ccdag_dn').values()) + sum(abs(val) for val in obs_ver.get('ccdag_dn').values()))
                site_exp_dict_int = one_site_dict(psi, step.env, n_int)

                mean_int = np.mean(list(site_exp_dict_int.values()))
                tot_energy =  U * mean_int - (cdagc_up + ccdag_up + cdagc_dn + ccdag_dn)/(xx*yy)
                if abs(tot_energy - tot_energy_old) < tol_exp:
                    break # here break if the relative differnece is below tolerance
                tot_energy_old = tot_energy

            print('energy :', tot_energy)
            site_exp_dict_up = one_site_dict(psi, step.env, n_up)        
            site_exp_dict_dn = one_site_dict(psi, step.env, n_dn)    
            site_exp_dict_magnetization = one_site_dict(psi, step.env, (n_up-n_dn))
            print('dn density :', site_exp_dict_dn)
            print('up density:', site_exp_dict_up)
            print('double occupancy :', site_exp_dict_int)

            op = {'SzSz': {'l': (n_up-n_dn), 'r': (n_up-n_dn)}, 'hd': {'l': (fid-n_up)@(fid-n_dn),'r': n_int}, 
                  'dh': {'l': n_int, 'r': (fid-n_up)@(fid-n_dn)},'hh': {'l': (fid-n_up)@(fid-n_dn), 'r': (fid-n_up)@(fid-n_dn)}}
            op_hor, op_ver = nn_exp_dict(psi, step.env, op)
            zz_dict = {**op_hor['SzSz'],**op_ver['SzSz']}
            hd_dict = {**op_hor['hd'],**op_ver['hd']}
            dh_dict = {**op_hor['dh'],**op_ver['dh']}
            hh_dict = {**op_hor['hh'],**op_ver['hh']}

            str_beta = str(beta)        
            # Initialize a master dictionary to hold all observables
            master_dict[str_beta] = {'Energy': np.array([tot_energy]), 'nup': np.array([site_exp_dict_up]), 'ndn': np.array([site_exp_dict_dn]),
            'z': np.array([site_exp_dict_magnetization]), 'nupndn': np.array([site_exp_dict_int]), 'zz_dict': np.array([zz_dict]), 'hd_dict': np.array([hd_dict]),
            'dh_dict': np.array([dh_dict]), 'hh_dict': np.array([hh_dict]), 'cdagc_up': np.array(obs_hor['cdagc_up']), 'ccdag_up': np.array(obs_hor['ccdag_up']),
            'cdagc_dn': np.array(obs_hor['cdagc_dn']), 'ccdag_dn': np.array(obs_hor['ccdag_dn'])}

    file_name = "init_pattern_%s_Nx_%1.0f_Ny_%1.0f_dbeta_%1.4f_%s_Ds_%s_U_%1.2f_%s" % (initial_pattern, xx, yy, dbeta, tr_mode, D, U, sym)
    with open("observables_hf_%s.txt"%(file_name), 'w') as f:
        for key in master_dict.keys():
            f.write(f"beta = {key}:\n")
            for subkey in master_dict[key]:
                f.write(f"  {subkey}:\n")
                item = master_dict[key][subkey].item()  # item() returns the scalar as a native Python number
                f.write(f"    {item}\n")
    np.savez("observables_doped_%s.npz"%(file_name), **master_dict)