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,17 +97,24 @@ 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)]
         local_density = np.diag(np.average(rho, axis=tuple([i for i in range(dim)])))
         exchange_mf = convolution(rho, H_int) * nk ** (-dim)
         direct_mf = np.diag(np.einsum("i,ij->j", local_density, H0_int))
+        dc_direct = local_density.T @ H0_int @ local_density
+        dc_exchange = np.einsum('kij, kji', exchange_mf, rho) * nk ** (-dim)
+        dc_energy = 0.5*(-dc_exchange + dc_direct) * np.eye(direct_mf.shape[-1])
     else:
         local_density = np.diag(rho)
         exchange_mf = rho * H_int
         direct_mf = np.diag(np.einsum("i,ij->j", local_density, H_int))
-    return direct_mf - exchange_mf
+        dc_energy_direct = np.diag(np.einsum("ij, i, j->", H_int, local_density, local_density))
+        dc_energy_cross = np.diag(np.einsum("ij, ij, ji->", H_int, rho, rho))
+        dc_energy = 2 * dc_energy_direct - dc_energy_cross
+    return direct_mf - exchange_mf# - dc_energy * np.eye(direct_mf.shape[-1])
 
 def total_energy(h, rho):
     """
@@ -124,7 +132,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 +159,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])