-- two p orbitals on a honeycomb lattice
dAB = 0.2
tnn = 1.1
HTB = NewTightBinding()
HTB.Name = "dichalcogenide tight binding"
HTB.Cell = {{sqrt(3),0,0},
            {sqrt(3/4),3/2,0},
            {0,0,1}}
HTB.Atoms = { {"A", {0,0,0},       {{"p", {"0"}}}},
              {"B", {sqrt(3),1,0}, {{"p", {"0"}}}}}
HTB.Hopping = {{"A.p","A.p",{         0,   0,0},{{-dAB/2}}},
               {"B.p","B.p",{         0,   0,0},{{ dAB/2}}},
               {"A.p","B.p",{         0,   1,0},{{ tnn  }}},
               {"B.p","A.p",{         0,  -1,0},{{ tnn  }}},
               {"A.p","B.p",{ sqrt(3/4),-1/2,0},{{ tnn  }}},
               {"B.p","A.p",{-sqrt(3/4), 1/2,0},{{ tnn  }}},
               {"A.p","B.p",{-sqrt(3/4),-1/2,0},{{ tnn  }}},
               {"B.p","A.p",{ sqrt(3/4), 1/2,0},{{ tnn  }}}}

kpath = { {"M",{0,2*pi/3,0}}, 100, {"G",{0,0,0}}, 100, {"K",{4*pi/(3*sqrt(3)),0,0}}, 100, {"M",{2*pi/(sqrt(3)),0,0}} }

-- the character of an eigenstate: its weight on atom A
function OnA(psi)
  return Complex.Re(psi[1]*Conjugate(psi[1]))
end

-- without a self-energy: band energies and characters, one row per k-point
E, CA = TightBinding.CalculateBands(HTB, kpath, OnA)
print("without a self-energy, at Gamma (k-point 101):")
print(string.format("  energies %8.4f %8.4f   character A %6.4f %6.4f", E[101][1], E[101][2], CA[101][1], CA[101][2]))

-- a self-energy on atom A only, with two poles at -2 and 2
SigmaA = ResponseFunction.New({{0, -2, 2}, {0.5, 0.5}, mu=0, type="ListOfPoles"})
G, GA = TightBinding.CalculateBands(HTB, kpath, {SigmaA, 0}, OnA)
print("with a self-energy on atom A, at Gamma:")
t, tA = ResponseFunction.ToTable(G[101]), ResponseFunction.ToTable(GA[101])
for i = 1, #t[2] do
  print(string.format("  pole %8.4f   weight %6.4f   on A %6.4f", t[1][i+1], t[2][i], tA[2][i]))
end

-- the plots, the bands in blue (on A) to red (on B)
Frame = {"Frame",{{"Ymin",-4},{"Ymax",4},{"dYTick",1},{"YLabel","Energy"}}}
pl0 = Graphics.PlotBands(kpath, E, CA, {Frame})
pl1 = Graphics.PlotBands(kpath, G, GA, {Frame, {"GammaL",0.05}})
for name, pl in pairs({BandsWithout=pl0, BandsWithSigma=pl1}) do
  file = io.open(name..".svg", "w")
  file:write(Graphics.ToSVG(pl, {{"RelativeSize",true}}))
  file:close()
end