Documentation

Construct the documentation for the release of the package.

pymf/tb/utils.py

 from itertools import product
-
 import numpy as np
 
-from pymf.mf import fermi_on_grid
-from pymf.tb.transforms import tb_to_khamvector
+from pymf.tb.tb import _tb_type
+from pymf.mf import fermi_on_kgrid
+from pymf.tb.transforms import tb_to_kgrid
 
 
-def generate_guess(vectors, ndof, scale=1):
-    """Generate guess for a tight-binding model.
+def generate_guess(
+    tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int, scale: float = 1
+) -> _tb_type:
+    """Generate hermitian guess tight-binding dictionary.
 
     Parameters
     ----------
-    vectors : list
-        List of hopping vectors.
-    ndof : int
-        Number internal degrees of freedom (orbitals),
-    scale : float
-        The scale of the guess. Maximum absolute value of each element of the guess.
-
+    tb_keys :
+       List of hopping vectors (tight-binding dictionary keys) the guess contains.
+    ndof :
+        Number internal degrees of freedom within the unit cell.
+    scale :
+        Scale of the random guess.
     Returns
     -------
-    guess : tb dictionary
-        Guess in the form of a tight-binding model.
+    :
+        Hermitian guess tight-binding dictionary.
     """
     guess = {}
-    for vector in vectors:
+    for vector in tb_keys:
         if vector not in guess.keys():
             amplitude = scale * np.random.rand(ndof, ndof)
             phase = 2 * np.pi * np.random.rand(ndof, ndof)
@@ -40,25 +41,44 @@ def generate_guess(vectors, ndof, scale=1):
     return guess
 
 
-def generate_vectors(cutoff, dim):
-    """Generate hopping vectors up to a cutoff.
+def generate_tb_keys(cutoff: int, dim: int) -> list[tuple[None] | tuple[int, ...]]:
+    """Generate tight-binding dictionary keys up to a cutoff.
 
     Parameters
     ----------
-    cutoff : int
-        Maximum distance along each direction.
-    dim : int
-        Dimension of the vectors.
+    cutoff :
+        Maximum distance along each dimension to generate tight-bindign dictionary keys for.
+    dim :
+        Dimension of the tight-binding dictionary.
 
     Returns
     -------
-    List of hopping vectors.
+    :
+        List of generated tight-binding dictionary keys up to a cutoff.
     """
     return [*product(*([[*range(-cutoff, cutoff + 1)]] * dim))]
 
 
-def calculate_fermi_energy(tb, filling, nk=100):
-    """Calculate the Fermi energy for a given filling."""
-    kham = tb_to_khamvector(tb, nk, ks=None)
+def calculate_fermi_energy(tb: _tb_type, filling: float, nk: int = 100):
+    """
+    Calculate the Fermi energy of a given tight-binding dictionary.
+
+    Parameters
+    ----------
+    tb :
+        Tight-binding dictionary.
+    filling :
+        Number of particles in a unit cell.
+        Used to determine the Fermi level.
+    nk :
+        Number of k-points in a grid to sample the Brillouin zone along each dimension.
+        If the system is 0-dimensional (finite), this parameter is ignored.
+
+    Returns
+    -------
+    :
+        Fermi energy.
+    """
+    kham = tb_to_kgrid(tb, nk)
     vals = np.linalg.eigvalsh(kham)
-    return fermi_on_grid(vals, filling)
+    return fermi_on_kgrid(vals, filling)