Download this notebook (.ipynb + data files)

8. Emery model on the Lieb lattice#

Status:

  • we can reproduce Adam’s energy for the 4x4 “tilted” Lieb lattice (2 Cu per unit cell)!!

from pathlib import Path

from tutorial_utils import run_afqmc

scratch_dir = Path("data")
scratch_dir.mkdir(parents=True, exist_ok=True)

8.1. Set up the lattice#

from afqmctools.systems.lattice import get_lattice
import afqmctools.utils.visualize as vis

lattice_params = dict(
    L1 = 4,
    L2 = 4,
    type = 'square',
    basis = [
        [0.,0.],    # Cu 3d_{x^2 - y^2}
        [0.5,0.0],  # O 2p_x
        [0.0,0.5]   # O 2p_y
    ]
)

lattice = get_lattice(
    params=lattice_params
)

vis.plot_lattice(lattice,title="Lieb Lattice",show_labels=False)

8.2. Set up the Hamiltonian#

import numpy as np

# definning model parameters
tdp = 1.2
tpp = 0.7
epsilon_d = -7.6
epsilon_p = -3.2
Ud = 8.4
Up = 2.0

# for the Emery model, we have to define sign conventions related to the
#   overlap of the +/- lobes of the Co 3d_{x^2 - y^2} orbital and the O 2p_x / 2p_y orbitals
#
#   For tpd:
#       we'll use a '+' sign for the -x, +y direction
#       we'll use a '-' sign for the +x, -y direction
#   For tpp:
#       we'll use a '+' sign for the 1/sqrt(2) (-x, +y), and 1/sqrt(2) (+x,-y) direction
#       we'll use a '-' sign for the 1/sqrt(2) (-x, -y), and 1/sqrt(2) (+x,+y) direction

# define unit vectors
x_hat = np.array([1.0,0.0])
y_hat = np.array([0.0,1.0])

r_x_minus_y_hat = (1/np.sqrt(2))*np.array([1.0,-1.0])  # \hat{r}_{x-y} = \hat{x} - \hat{y}
r_xy_hat = (1/np.sqrt(2))*np.array([1.0,1.0])         # \hat{r}_{x+y} = \hat{x} + \hat{y}

t_pd_signs = { () }

# TODO: make the hopping matrix taking care of the +/- sign
hopping_shape = (lattice.N_sites, lattice.N_sites)

hopping = np.zeros(hopping_shape,dtype=np.complex128)

for pair in lattice.get_nth_neighbors(n=1):
    i,j = pair.i,pair.j
    if j <= i:
        continue

    if np.dot(pair.r_relative,x_hat) > 0 or np.dot(pair.r_relative,-y_hat) > 0:
        hopping[i,j] = -tdp
    elif np.dot(pair.r_relative,-x_hat) > 0 or np.dot(pair.r_relative,y_hat) > 0:
        hopping[i,j] = tdp
    else:
        raise ValueError(f"Unexpected relative distance for nearest neighbors. {pair.r_relative}")

for pair in lattice.get_nth_neighbors(n=2):
    i,j = pair.i,pair.j
    if j <= i:
        continue

    if np.dot(pair.r_relative,r_x_minus_y_hat) != 0.0:
        hopping[i,j] = tpp
    elif np.dot(pair.r_relative,r_xy_hat) != 0.0:
        hopping[i,j] = -tpp
    else:
        raise ValueError(f"Unexpected relative distance for next-nearest neighbors. {pair.r_relative}")

# add in the j<=i part of the hopping matrix
hopping = hopping + hopping.T.conj()


# add in on-site energies
U_emery = np.zeros((lattice.N_sites),dtype=np.complex128)
for i in range(lattice.N_sites):
    if i % 3 == 0:
        hopping[i,i] = epsilon_d
        U_emery[i] = Ud
    elif i % 3 == 1 or i % 3 == 2:
        hopping[i,i] = epsilon_p
        U_emery[i] = Up
from afqmctools.hamiltonian.model.director import HamiltonianDirector
import afqmctools.utils.io as io

hamiltonian_params = {
    'hamiltonian' : dict(
        t = 0.0,
        U = U_emery,
        custom_one_body = hopping,
    )
}

hamiltonian = HamiltonianDirector(
    lattice=lattice,
    source=hamiltonian_params
).build()

io.write_model_hamiltonian(
    hamiltonian=hamiltonian,
    fname=scratch_dir/"afqmc.h5"
)

8.3. Run Hartree-Fock (HF)#

(sometimes jax doesn’t initialize properly on the first run; just try running again!)

import jax
jax.config.update('jax_platform_name', 'cpu')

from autohf import lattice_hf, AutoHFHamiltonian

from afqmctools.inputs.from_autohf import autohf_to_afqmc

nelec = (8,8)

hf_settings = dict(
    ansatz = 'SD',
    numSteps = 400,
    nelec = nelec,
    numTrials = 20,
    seed = 42,
    measure_spin = True,
    noncollinear = True
)

results = lattice_hf(
    hamiltonian=AutoHFHamiltonian(hamiltonian),
    settings=hf_settings
)

autohf_to_afqmc(
    results,
    output_fname = scratch_dir/"ghf_wfn.h5"
)

8.4. Analyze the HF results#

Before proceeding, it’s important to check that the HF solution is correct. First, check that AutoHF actually converged.

We will also want to check the charge and spin density on the lattice. You can you the lattice model visualizer vis.plot_lattice(lattice,title="Lieb Lattice",density=...) to plot an arbitrary density on the lattice.

import jax.numpy as jnp

# now look at the density
autohf_results = results[0]
autohf_funcs = results[1]
makeRDMs = autohf_funcs['makeRDMs']
state = autohf_results['state']
rdm = makeRDMs(state)

L = lattice.L
# noncollinear
M = lattice.N_sites
new_shape = (L[0],L[1],lattice.nb)
rho_up = rdm[:M,:M].diagonal().reshape(new_shape)
rho_down = rdm[M:,M:].diagonal().reshape(new_shape)

rho_charge = rho_up + rho_down
rho_spin = 0.5*(rho_up - rho_down)

print(f"Integrated charge: {rho_charge.real.sum()}")
print(f" mean charge density: {rho_charge.real.mean()} +/- {np.std(rho_charge.real)}")
print(f" mean spin density: {rho_spin.real.mean()} +/- {np.std(rho_spin.real)}")

Cu_inds =  jnp.array([i*lattice.nb for i in range(lattice.N_sites//lattice.nb)])
Opx_inds = jnp.array([i*lattice.nb+1 for i in range(lattice.N_sites//lattice.nb)])
Opy_inds = jnp.array([i*lattice.nb+2 for i in range(lattice.N_sites//lattice.nb)])

Cu_density_up = rdm[:M,:M].diagonal()[Cu_inds]
Opx_density_up = rdm[:M,:M].diagonal()[Opx_inds]
Opy_density_up = rdm[:M,:M].diagonal()[Opy_inds]

Cu_density_down  = rdm[M:,M:].diagonal()[Cu_inds]
Opx_density_down = rdm[M:,M:].diagonal()[Opx_inds]
Opy_density_down = rdm[M:,M:].diagonal()[Opy_inds]

Cu_density = Cu_density_up + Cu_density_down
Opx_density = Opx_density_up + Opx_density_down
Opy_density = Opy_density_up + Opy_density_down

print(f"For Cu sublattice: mean charge density: {np.mean(Cu_density.real)} +/- {np.std(Cu_density.real)}")
print(f"For O px sublattice: mean charge density: {np.mean(Opx_density.real)} +/- {np.std(Opx_density.real)}")
print(f"For O py sublattice: mean charge density: {np.mean(Opy_density.real)} +/- {np.std(Opy_density.real)}")


vis.plot_lattice(
    lattice,
    density=rho_charge.real,
    density_label="charge density",
    cmap = 'bwr',
    show_labels=False
)

vis.plot_lattice(
    lattice,
    density=rho_spin.real,
    density_label="spin density",
    cmap = 'bwr',
    show_labels=False
)

8.5. 🚧 Old Scratch work below!! 🚧#

import numpy as np

from afqmctools.hamiltonian.model.director import HamiltonianDirector
import afqmctools.utils.io as io
from afqmctools.wavefunction.free_electron import free_electron
from afqmctools.inputs.from_hdf import write_json
import afqmctools.utils.visualize as vis


# define the problem
nelec = (5,5)
nbands = 3

#t_dd = 1.0
t_pd = t = 1.0
t_pp = tprime = 0.5

lattice_params = dict(
    L1 = 4,
    L2 = 4,
    boundary1 = 'PBC',
    boundary2 = 'PBC'
)

lattice = get_lattice(
    params=lattice_params
)

nearest_neighbors = lattice.get_nth_neighbors(n=1)
sites = lattice.sites

nbasis = nbands*lattice.N_sites

hopping_intrasite = np.array([
    [0.0,   t_pd,  t_pd],
    [t_pd,   0.0,  t_pp],
    [t_pd,  t_pp,   0.0]
])

hopping_neatest_x =  np.array([
    [0.0,   t_pd,  0.0],
    [t_pd,   0.0,  t_pp],
    [0.0,  t_pp,   0.0]
])

hopping_neatest_y =  np.array([
    [0.0,    0.0,  t_pd],
    [0.0,   0.0,  t_pp],
    [t_pd,  t_pp,   0.0]
])

hopping = np.zeros((nbasis,nbasis))

# add in intrasite
for site in sites:
    site_slice = slice(site.index*nbands,(site.index+1)*nbands)
    hopping[ site_slice, site_slice ] = hopping_intrasite

# add in nearest-neighbor
for pair in nearest_neighbors:
    i_slice = slice(pair.i*nbands, (pair.i+1)*nbands)
    j_slice = slice(pair.j*nbands, (pair.j+1)*nbands)

    # hopping from Cu to O px/py bands is direction-dependent!
    if (pair.r_relative[0] == 1.0 or pair.r_relative[0] == -1.0):
        hopping[i_slice,j_slice] = hopping_neatest_x
    elif (pair.r_relative[1] == 1.0 or pair.r_relative[1] == -1.0):
        hopping[i_slice,j_slice] = hopping_neatest_y
    else:
        raise ValueError(f"Invalid relative position {pair.r_relative} for a square lattice")
import matplotlib.pyplot as plt

plt.matshow(hopping)
plt.colorbar()
hamiltonian_params = {
    'hamiltonian' : dict(
        nbands = nbands,
        t = hopping,
        U = 8.0
    )
}

hamiltonian = HamiltonianDirector(
    lattice=lattice,
    source=hamiltonian_params
).build()

io.write_model_hamiltion(
    hamiltonian=hamiltonian,
    fname=scratch_dir/"lieb.h5",
    nelec=nelec
)

8.5.1. Tilted Lieb Lattice : Annealing#

"""
This is a benchmarking script, so we'll keep things basic
"""
import numpy as np
import h5py as h5

import afqmctools.systems.lattice as lat
import afqmctools.hamiltonian.model.director as ham
import afqmctools.utils.io as io
import afqmctools.utils.visualize as vis
from afqmctools.wavefunction.free_electron import free_electron
from afqmctools.wavefunction.common import modified_gram_schmidt

from afqmctools.hamiltonian.model.builder import HamiltonianBuilder
from afqmctools.hamiltonian.model.ham_class import HamiltonianComponent,SpinSymm

import faulthandler; faulthandler.enable()

import scipy.sparse as sps

import jax.numpy as jnp
import jax

import jax
# Set the default device to CPU
#jax.config.update('jax_platform_name', 'cpu')

print("JAX devices : ",jax.devices())


import matplotlib.pyplot as plt
import pdb

#### Some input params ######
recompute = True
show_lattice_before = True

annealing_steps = 10
annealing_amplitude = 0.01

L = (4,4)
Cu_per_cell = 2
h = 1/8

#expected HF energy
Ehf_ref = -313.792

#############################

# this is Nelec = (h+1) * Ncell where Ncell is the number of primitive cells
# BUT, be careful. the tilted lattice has two primitive cells per unit cell!
nelec_total = (h+1) * L[0] * L[1] * Cu_per_cell

assert nelec_total % 2 == 0, "must have even number of electrons"

nelec = (int(nelec_total//2),int(nelec_total//2))

print(f" ==== Emery model run parameters ==== ")
print(f"   - hole doping, h = {h}")
print(f"   - nelectrons = {nelec}")


def make_emery(lattice, show_mats=False):
    """
    """

    # 2. define / build the model
    tdp = 1.2
    tpp = 0.7
    epsilon_d = -7.6
    epsilon_p = -3.2
    Ud = 8.4
    Up = 2.0

    x_hat = np.array([1.0,0.0])
    y_hat = np.array([0.0,1.0])

    r_x_minus_y_hat = (1/np.sqrt(2))*np.array([1.0,-1.0])  # \hat{r}_{x-y} = \hat{x} - \hat{y}
    r_xy_hat = (1/np.sqrt(2))*np.array([1.0,1.0])         # \hat{r}_{x+y} = \hat{x} + \hat{y}

    lattice = lattice

    hopping_shape = (lattice.N_sites, lattice.N_sites)

    hopping = np.zeros(hopping_shape)

    for pair in lattice.get_nth_neighbors(n=1):
        i,j = pair.i,pair.j
        if j <= i:
            continue

        # NOTE: hard code the hopping phase between sublattices for now.
        tdp_pq = tdp*np.array(
            [[0.0,1.0,1.0],
            [-1.0,0.0,1.0],
            [-1.0,1.0,0.0]]
        ) # starting on sublattice from row, ending on sublattece from column
        p = pair.i % 3 # sublattice of i
        q = pair.j % 3 # sublattice of j

        # TODO: save origin and destination sublattice to NeighborPair object
        if np.dot(pair.r_relative,x_hat) > 0 or np.dot(pair.r_relative,-y_hat) > 0:
            hopping[i,j] = -tdp_pq[p,q]
        elif np.dot(pair.r_relative,-x_hat) > 0 or np.dot(pair.r_relative,y_hat) > 0:
            hopping[i,j] = tdp_pq[p,q]
        else:
            raise ValueError(f"Unexpected relative distance for nearest neighbors. {pair.r_relative}")

    for pair in lattice.get_nth_neighbors(n=2):
        i,j = pair.i,pair.j
        if j <= i:
            continue

        if np.dot(pair.r_relative,r_x_minus_y_hat) != 0.0:
            hopping[i,j] = tpp
        elif np.dot(pair.r_relative,r_xy_hat) != 0.0:
            hopping[i,j] = -tpp
        else:
            raise ValueError(f"Unexpected relative distance for next-nearest neighbors. {pair.r_relative}")


    # add in the j<=i part of the hopping matrix
    hopping = hopping + hopping.T.conj()


    # add in on-site energies
    U = np.zeros((lattice.N_sites,lattice.N_sites))
    for i in range(lattice.N_sites):
        if i % 3 == 0:
            hopping[i,i] = epsilon_d
            U[i,i] = Ud
        elif i % 3 == 1 or i % 3 == 2:
            hopping[i,i] = epsilon_p
            U[i,i] = Up

    if show_mats:
        plt.matshow(hopping.real)
        plt.title("Hopping matrix")
        plt.colorbar()
        plt.grid(True)
        plt.show()

    if show_mats:
        plt.matshow(U)
        plt.title("U matrix")
        plt.colorbar()
        plt.grid(True)
        plt.show()

    hopping = np.block(
        [[hopping,np.zeros(hopping_shape)]
        ,[np.zeros(hopping_shape),hopping]]
    )

    # 3. save the Hamiltonian
    builder = HamiltonianBuilder(lattice=lattice)
    hopping = sps.csr_matrix(hopping)
    custom_one_body = HamiltonianComponent(
        csr_array=hopping,
        model_type='one_body',
        spin_symm=SpinSymm.NONCOLLINEAR
    )

    # manually add the custom term
    builder.hamiltonian["tij"] = custom_one_body
    custom_hubbard = HamiltonianComponent(
        csr_array=U,
        model_type='hubbard_u',
        hst_type='discrete_spin'
    )
    builder.hamiltonian["Uij"] = custom_hubbard
    builder.finalize()

    hamiltonian = builder.hamiltonian

    io.write_model_hamiltonian(
        hamiltonian=hamiltonian,
        fname="afqmc.h5",
        nelec=nelec
    )

    return hamiltonian

# 1. define the lattice
basis = [ np.array(delta) for delta in [(0,0),(0.5,0),(0,0.5),(1.0,0),(1.5,0),(1.0,0.5)] ]

lattice = lat.CustomLattice(
    L=(4,4),
    a1=np.array((1.0,-1.0)),
    a2=np.array((1.0,1.0)),
    basis=basis,
    axis1_boundary=lat.PBCBoundary,
    axis2_boundary=lat.PBCBoundary,
)

if show_lattice_before:
    vis.plot_lattice(
        lattice,
        show_labels=True,
        show_lattice_vecs=False,
        title="Lieb lattice"
    )

hamiltonian = make_emery(lattice,show_mats=True)


# for initial guess - we need another H with a small twist!
THETA_X = 1/np.sqrt(592560607) # 592560607 is prime
THETA_Y = 1/np.sqrt(47603)     # 47603 is prime
twist = np.array((THETA_X,THETA_Y))

lattice_wt_twist = lat.CustomLattice(
    L=(4,4),
    a1=np.array((1.0,-1.0)),
    a2=np.array((1.0,1.0)),
    basis=basis,
    twist=twist
)

hamiltonian2 = make_emery(lattice_wt_twist)

wfn,spin_symm = free_electron(
    source=hamiltonian2,
    nelec=nelec,
    lattice=lattice,
    spin_symm=SpinSymm.NONCOLLINEAR
)

# add noise to the initial guess
wfn = (wfn[0],wfn[1] + annealing_amplitude*np.random.randn(*wfn[1].shape))

# 4. run autohf
from autohf import lattice_hf, AutoHFHamiltonian
best_E_final = None
for a in range(annealing_steps):
    print(f"Annealing step {a}")

    slater_det = wfn[1][0]

    hf_settings = dict(
        ansatz = 'SD',#'SD_Rot',
        steps = 500, #numSteps = 500,
        #output = f"ghf_wfn_step_{a}.h5",
        nelec = nelec,
        batch_size = 20, #numTrials = 20,
        #seed = 42,
        measure_spin = False,
        noncollinear =  True,
        gpu = False,
        state0_scale = 1.0e-3,
    )

    results, functions = lattice_hf(
        hamiltonian=AutoHFHamiltonian(hamiltonian),
        lattice=lattice,
        settings=hf_settings,
        jaxoptargs=dict(tol=1e-8),
        initial_guess=slater_det,
    )



    # now look at the density
    makeRDMs = functions['makeRDMs']
    state = results['state']
    rdm = makeRDMs(state)

    with h5.File("dev_checkpoint.h5","w") as f:
        f.create_dataset(f"rdm_{a}",data=rdm)

    # get a wavefunction for the next annealing step
    orbitals = results['orbitals'][0]

    slater_det = np.array(orbitals[:,:sum(nelec)])
    # orthonormalize!
    slater_det = modified_gram_schmidt(slater_det)
    slater_det = slater_det + annealing_amplitude*np.random.randn(*slater_det.shape)

    wfn = (wfn[0],[slater_det])

    L = lattice.L

    # noncollinear
    M = lattice.N_sites
    new_shape = (L[0],L[1],lattice.nb)
    rho_up = rdm[:M,:M].diagonal().reshape(new_shape)
    rho_down = rdm[M:,M:].diagonal().reshape(new_shape)

    rho_charge = rho_up + rho_down
    rho_spin = 0.5*(rho_up - rho_down)



    if best_E_final is None:
        best_E_final = results['E']
        best_rho_charge = rho_charge
        best_rho_spin = rho_spin
        best_rdm = rdm

    if results['E'] < best_E_final:
        print("Found a better energy: ",results['E'])
        best_E_final = results['E']
        best_rho_charge = rho_charge
        best_rho_spin = rho_spin
        best_rdm = rdm

    if recompute:
        with h5.File("dev_checkpoint.h5","a") as f:
            f.create_dataset("E",data=best_E_final)
            f.create_dataset(f"best_rho_charge",data=best_rho_charge)
            f.create_dataset(f"best_rho_spin",data=best_rho_spin)


print(f"\nBest Energy = {best_E_final} t")
print(f"            = {best_E_final / (L[0] * L[1] * Cu_per_cell)} t / Cu site\n\n")

rdm = best_rdm
rho_charge = best_rho_charge
rho_spin = best_rho_spin

print(f"Integrated charge: {rho_charge.real.sum()}")

print(f" mean charge density: {rho_charge.real.mean()} +/- {np.std(rho_charge.real)}")
print(f" mean spin density: {rho_spin.real.mean()} +/- {np.std(rho_spin.real)}")

Cu_inds =  jnp.array([i*lattice.nb for i in range(lattice.N_sites//lattice.nb)])
Opx_inds = jnp.array([i*lattice.nb+1 for i in range(lattice.N_sites//lattice.nb)])
Opy_inds = jnp.array([i*lattice.nb+2 for i in range(lattice.N_sites//lattice.nb)])

Cu_density_up = rdm[:M,:M].diagonal()[Cu_inds]
Opx_density_up = rdm[:M,:M].diagonal()[Opx_inds]
Opy_density_up = rdm[:M,:M].diagonal()[Opy_inds]

Cu_density_down  = rdm[M:,M:].diagonal()[Cu_inds]
Opx_density_down = rdm[M:,M:].diagonal()[Opx_inds]
Opy_density_down = rdm[M:,M:].diagonal()[Opy_inds]

Cu_density = Cu_density_up + Cu_density_down
Opx_density = Opx_density_up + Opx_density_down
Opy_density = Opy_density_up + Opy_density_down

print(f"For Cu sublattice: mean charge density: {np.mean(Cu_density.real)} +/- {np.std(Cu_density.real)}")
print(f"For O px sublattice: mean charge density: {np.mean(Opx_density.real)} +/- {np.std(Opx_density.real)}")
print(f"For O py sublattice: mean charge density: {np.mean(Opy_density.real)} +/- {np.std(Opy_density.real)}")


if True:
    vis.plot_lattice(
        lattice,
        density=rho_charge.real,
        density_label="charge density",
        cmap = 'bwr',
        show_labels=False,
        save="charge.png"
    )

    vis.plot_lattice(
        lattice,
        density=rho_spin.real,
        density_label="spin density",
        cmap = 'bwr',
        show_labels=False,
        save="spin.png"
    )