import numpy as np
import sisl

from quant_met.routines import self_consistency_loop

a = 1.0  # Lattice constant
t = -1.0  # Nearest-neighbor hopping
bond = a  # Bond length

# Create an atom object with appropriate cutoff range
atom = sisl.Atom(1, R=bond + 0.01)

# Generate 2D square lattice geometry
geom = sisl.Geometry(
    [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
    [atom, atom],
    [[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, 10.0]],  # Unit cell (2D in 3D space)
)

hamiltonian = sisl.Hamiltonian(geom)
search_radius = [0.1 * bond, bond + 0.01]  # Search radius for neighbors

for ia in geom:
    idx_a = geom.close(ia, R=search_radius)
    hamiltonian[ia, idx_a[0]] = 0.0  # On-site energy
    for i in idx_a[1:]:
        hamiltonian[ia, i] = t  # Nearest-neighbor hopping

hamiltonian.finalize()

hamiltonian.write("hamiltonian.HSX")

[[6.28318531 0.         0.        ]
 [0.         6.28318531 0.        ]
 [0.         0.         0.62831853]]
[6.28318531 0.         0.        ]

# n_k is the number of k-points along each reciprocal direction
n_k = 10
k_grid_obj = sisl.MonkhorstPack(hamiltonian.geometry, [n_k, n_k, 1])  # 2D grid

# Extract the k-points as a NumPy array
k_grid = k_grid_obj.k

beta = 1.0 / 0.2  # inverse temperature (T = 0.01)
hubbard_int = 1.0  # On-site interaction strength (Hubbard U)
epsilon = 1e-5  # convergence threshold

hubbard_int_orbital_basis = np.array(
    [hubbard_int] * hamiltonian.no,
    dtype=np.float64,
)  # no = number of orbitals

final_gap = self_consistency_loop(
    hamiltonian=hamiltonian,
    kgrid=k_grid_obj,
    beta=beta,
    hubbard_int_orbital_basis=hubbard_int_orbital_basis,
    epsilon=epsilon,
)

[0.3552019 +0.000000e+00j 0.35520021-8.992404e-21j]