create solvers and interface modules

Co-authored-by: Johanna Zijderveld

codes/hf.py

@@ -2,7 +2,7 @@ from scipy.ndimage import convolve
 import numpy as np
 import codes.utils as utils
 
-def density_matrix(vals, vecs, E_F, dim):
+def density_matrix(vals, vecs, E_F):
     """
     Returns the mean field F_ij(k) = <psi_i(k)|psi_j(k)> for all k-points and
     eigenvectors below the Fermi level.
@@ -22,6 +22,7 @@ def density_matrix(vals, vecs, E_F, dim):
         Density matrix rho=rho[kx, ky, ..., i, j] where i,j are cell indices.
     """
     norbs = vals.shape[-1]
+    dim = len(vals.shape) - 1
     nk = vals.shape[0]
 
     if dim > 0:
@@ -78,7 +79,7 @@ def convolution(M1, M2):
     return V_output
 
 
-def compute_mf(rho, H_int, dim):
+def compute_mf(rho, H_int):
     """
     Compute mean-field correction at self-consistent loop.
 
@@ -96,6 +97,7 @@ def compute_mf(rho, H_int, dim):
     """
     
     nk = rho.shape[0]
+    dim = len(rho.shape) - 2
     
     if dim > 0:
         H0_int = H_int[*([0]*dim)]
@@ -124,7 +126,7 @@ def total_energy(h, rho):
     total_energy : float
         System total energy computed as tr[h@rho].
     """
-    return np.sum(np.trace(h @ rho, axis1=-1, axis2=-2)).real
+    return np.sum(np.trace(h @ rho, axis1=-1, axis2=-2)).real / np.prod(rho.shape[:-2])
 
 def updated_matrices(mf_k, model):
     """
@@ -151,9 +153,8 @@ def updated_matrices(mf_k, model):
     vals, vecs = np.linalg.eigh(hamiltonians)
     vecs = np.linalg.qr(vecs)[0]
     E_F = utils.get_fermi_energy(vals, model.filling)
-    rho = density_matrix(vals=vals, vecs=vecs, E_F=E_F, dim=model.dim)
+    rho = density_matrix(vals=vals, vecs=vecs, E_F=E_F)
     return rho, compute_mf(
         rho=rho,
-        H_int=model.H_int,
-        dim=model.dim) - E_F * np.eye(model.hamiltonians_0.shape[-1])
+        H_int=model.H_int) - E_F * np.eye(model.hamiltonians_0.shape[-1])