Bundle density matrix and fermi energy calc

codes/mf.py

@@ -4,19 +4,19 @@ from codes.tb.transforms import tb_to_khamvector, ifftn_to_tb
 from scipy.fftpack import ifftn
 
 
-def density_matrix_kgrid(kham, fermi_energy):
+def density_matrix_kgrid(kham, filling):
     """
      Parameters
      ----------
      kham : npndarray
          Hamiltonian in k-space of shape (len(dim), norbs, norbs)
-     fermi_energy : float
-         Fermi level
+     filling : float
+        Number of particles in a unit cell.
 
      Returns
      -------
-     np.ndarray
-         Density matrix in k-space.
+     np.ndarray, flaot
+         Density matrix in k-space and Fermi energy.
 
     Notes
     -----
@@ -24,11 +24,12 @@ def density_matrix_kgrid(kham, fermi_energy):
 
     """
     vals, vecs = np.linalg.eigh(kham)
-    unocc_vals = vals > fermi_energy
+    fermi = fermi_on_grid(vals, filling)
+    unocc_vals = vals > fermi
     occ_vecs = vecs
     np.moveaxis(occ_vecs, -1, -2)[unocc_vals, :] = 0
     rho_krid = occ_vecs @ np.moveaxis(occ_vecs, -1, -2).conj()
-    return rho_krid
+    return rho_krid, fermi
 
 def density_matrix(h, filling, nk):
     """
@@ -51,9 +52,9 @@ def density_matrix(h, filling, nk):
     ndim = len(list(h)[0])
     if ndim > 0:
         kham = tb_to_khamvector(h, nk=nk)
-        fermi = fermi_on_grid(kham, filling)
+        rho_grid, fermi = density_matrix_kgrid(kham, filling)
         return (
-            ifftn_to_tb(ifftn(density_matrix_kgrid(kham, fermi), axes=np.arange(ndim))),
+            ifftn_to_tb(ifftn(rho_grid, axes=np.arange(ndim))),
             fermi,
         )
     else:
@@ -100,14 +101,14 @@ def meanfield(density_matrix_tb, h_int):
     return add_tb(direct, exchange)
 
 
-def fermi_on_grid(kham, filling):
+def fermi_on_grid(vals, filling):
     """
      Compute the Fermi energy on a grid of k-points.
 
      Parameters
      ----------
-     kham : ndarray
-        Hamiltonian in k-space of shape (len(dim), norbs, norbs)
+     vals : ndarray
+        Eigenvalues of the hamiltonian in k-space of shape (len(dim), norbs, norbs)
      filling : int
          Number of particles in a unit cell.
      Returns
@@ -115,9 +116,6 @@ def fermi_on_grid(kham, filling):
     fermi_energy : float
          Fermi energy
     """
-
-    vals = np.linalg.eigvalsh(kham)
-
     norbs = vals.shape[-1]
     vals_flat = np.sort(vals.flatten())
     ne = len(vals_flat)

codes/solvers.py

 from codes.params.rparams import tb_to_rparams, rparams_to_tb
-from codes.mf import fermi_on_grid
-from codes.tb.transforms import tb_to_khamvector
+from codes.tb.utils import calculate_fermi_energy
 from codes.tb.tb import add_tb
 import scipy
 from functools import partial
@@ -59,5 +58,5 @@ def solver(
     result = rparams_to_tb(
         optimizer(f, mf_params, **optimizer_kwargs), list(Model.h_int), shape
     )
-    fermi = fermi_on_grid(tb_to_khamvector(add_tb(Model.h_0, result), nk=nk), Model.filling)
+    fermi = calculate_fermi_energy(add_tb(Model.h_0, result), Model.filling, nk=nk)
     return add_tb(result, {Model._local_key: -fermi * np.eye(Model._size)})

codes/tb/utils.py

@@ -89,4 +89,5 @@ def calculate_fermi_energy(tb, filling, nk=100):
     Calculate the Fermi energy for a given filling.
     """
     kham = tb_to_khamvector(tb, nk, ks=None)
-    return fermi_on_grid(kham, filling)
+    vals = np.linalg.eigvalsh(kham)
+    return fermi_on_grid(vals, filling)