adjust docstrings and ensure consistent naming schemes

pymf/kwant_helper/utils.py

 import inspect
 from copy import copy
 from itertools import product
+from typing import Callable
+from pymf.tb.tb import tb_type
 
 import kwant
 import numpy as np
 from scipy.sparse import coo_array
 
 
-def builder_to_tb(builder, params={}, return_data=False):
-    """Construct a tight-binding model dictionary from a `kwant.Builder`.
+def builder_to_tb(
+    builder: kwant.Builder, params: dict = {}, return_data: bool = False
+) -> tb_type:
+    """Construct a tight-binding dictionary from a `kwant.Builder` system.
 
     Parameters
     ----------
-    builder : `kwant.Builder`
-        Either builder for non-interacting system or interacting Hamiltonian.
-    params : dict
-        Dictionary of parameters to evaluate the Hamiltonian.
-    return_data : bool
+    builder :
+       system to convert to tight-binding dictionary.
+    params :
+        Dictionary of parameters to evaluate the builder on.
+    return_data :
         Returns dictionary with sites and number of orbitals per site.
 
     Returns
     -------
-    h_0 : dict
-        Tight-binding model of non-interacting systems.
-    data : dict
+    h_0 :
+        Tight-binding dictionary that corresponds to the builder.
+    data :
         Data with sites and number of orbitals. Only if `return_data=True`.
     """
     builder = copy(builder)
@@ -124,27 +128,35 @@ def builder_to_tb(builder, params={}, return_data=False):
         return h_0
 
 
-def build_interacting_syst(builder, lattice, func_onsite, func_hop, max_neighbor=1):
-    """Construct an auxiliary `kwant` system to build Hamiltonian matrix.
+def build_interacting_syst(
+    builder: kwant.Builder,
+    lattice: kwant.lattice,
+    func_onsite: Callable,
+    func_hop: Callable,
+    max_neighbor: int = 1,
+) -> kwant.Builder:
+    """
+    Construct an auxiliary `kwant` system that encodes the interactions.
 
     Parameters
     ----------
-    builder : `kwant.Builder`
+    builder :
         Non-interacting `kwant` system.
-    func_onsite : function
-        Onsite function.
-    func_hop : function
-        Hopping function.
-    max_neighbor : int
-        Maximal nearest-neighbor order.
+    lattice :
+        Lattice of the system.
+    func_onsite :
+        Onsite interactions function.
+    func_hop :
+        Hopping/inter unit cell interactions function.
+    max_neighbor :
+        The maximal number of neighbouring unit cells (along a lattice vector)
+        connected by interaction. Interaction goes to zero after this distance.
 
     Returns
     -------
-    int_builder : `kwant.Builder`
-        Dummy `kwant.Builder` to compute interaction matrix.
+    int_builder :
+        Auxiliary `kwant.Builder` that encodes the interactions of the system.
     """
-    # lattice_info = list(builder.sites())[0][0]
-    # lattice = kwant.lattice.general(lattice_info.prim_vecs, norbs=lattice_info.norbs)
     int_builder = kwant.Builder(kwant.TranslationalSymmetry(*builder.symmetry.periods))
     int_builder[builder.sites()] = func_onsite
     for neighbors in range(max_neighbor):

pymf/mf.py

 import numpy as np
 from scipy.fftpack import ifftn
 
-from pymf.tb.tb import add_tb
+from pymf.tb.tb import add_tb, tb_type
 from pymf.tb.transforms import ifftn_to_tb, tb_to_khamvector
 
 
-def construct_density_matrix_kgrid(kham, filling):
+def construct_density_matrix_kgrid(
+    kham: np.ndarray, filling: float
+) -> (np.ndarray, float):
     """Calculate density matrix on a k-space grid.
 
     Parameters
     ----------
-    kham : npndarray
-         Hamiltonian in k-space of shape (len(dim), norbs, norbs)
-    filling : float
+    kham :
+        Hamiltonian from which to construct the density matrix.
+        The hamiltonian is sampled on a grid of k-points and has shape (nk, nk, ..., ndof, ndof),
+        where ndof is number of internal degrees of freedom.
+    filling :
         Number of particles in a unit cell.
+        Used to determine the Fermi level.
 
     Returns
     -------
-     np.ndarray, float
-         Density matrix in k-space and Fermi energy.
+    :
+        Density matrix on a k-space grid with shape (nk, nk, ..., ndof, ndof) and Fermi energy.
     """
     vals, vecs = np.linalg.eigh(kham)
     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, fermi
+    density_matrix_krid = occ_vecs @ np.moveaxis(occ_vecs, -1, -2).conj()
+    return density_matrix_krid, fermi
 
 
-def construct_density_matrix(h, filling, nk):
+def construct_density_matrix(h: tb_type, filling: float, nk: int) -> (tb_type, float):
     """Compute the density matrix in real-space tight-binding format.
 
     Parameters
     ----------
-    h : dict
-        Tight-binding model.
-    filling : float
-        Filling of the system.
-    nk : int
-        Number of k-points in the grid.
+    h :
+        Hamiltonian tight-binding dictionary from which to construct the density matrix.
+    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
     -------
-    (dict, float)
-        Density matrix in real-space tight-binding format and Fermi energy.
+    :
+        Density matrix tight-binding dictionary and Fermi energy.
     """
     ndim = len(list(h)[0])
     if ndim > 0:
         kham = tb_to_khamvector(h, nk=nk)
-        rho_grid, fermi = construct_density_matrix_kgrid(kham, filling)
+        density_matrix_krid, fermi = construct_density_matrix_kgrid(kham, filling)
         return (
-            ifftn_to_tb(ifftn(rho_grid, axes=np.arange(ndim))),
+            ifftn_to_tb(ifftn(density_matrix_krid, axes=np.arange(ndim))),
             fermi,
         )
     else:
-        rho, fermi = construct_density_matrix_kgrid(h[()], filling)
-        return {(): rho}, fermi
+        density_matrix, fermi = construct_density_matrix_kgrid(h[()], filling)
+        return {(): density_matrix}, fermi
 
 
-def meanfield(density_matrix_tb, h_int):
-    """Compute the mean-field in k-space.
+def meanfield(density_matrix: tb_type, h_int: tb_type) -> tb_type:
+    """Compute the mean-field correction from the density matrix.
 
     Parameters
     ----------
-    density_matrix : dict
-        Density matrix in real-space tight-binding format.
-    h_int : dict
-        Interaction tb model.
-
+    density_matrix :
+        Density matrix tight-binding dictionary.
+    h_int :
+        Interaction hermitian Hamiltonian tight-binding dictionary.
+        The interaction must be of density-density type, h_int[R][i, j] * c_i^dagger(R) c_j^dagger(0) c_j(0) c_i(R).
+        For example in 1D system with ndof internal degrees of freedom,
+        h_int[(2,)] = U * np.ones((ndof, ndof)) is a Coulomb repulsion interaction
+        with strength U between unit cells separated by 2 lattice vectors, where
+        the interaction is the same between all internal degrees of freedom.
     Returns
     -------
-    dict
-        Mean-field tb model.
+    :
+        Mean-field correction tight-binding dictionary.
     """
-    n = len(list(density_matrix_tb)[0])
+    n = len(list(density_matrix)[0])
     local_key = tuple(np.zeros((n,), dtype=int))
 
     direct = {
@@ -82,7 +93,7 @@ def meanfield(density_matrix_tb, h_int):
             np.array(
                 [
                     np.diag(
-                        np.einsum("pp,pn->n", density_matrix_tb[local_key], h_int[vec])
+                        np.einsum("pp,pn->n", density_matrix[local_key], h_int[vec])
                     )
                     for vec in frozenset(h_int)
                 ]
@@ -92,26 +103,26 @@ def meanfield(density_matrix_tb, h_int):
     }
 
     exchange = {
-        vec: -1 * h_int.get(vec, 0) * density_matrix_tb[vec]  # / (2 * np.pi)#**2
-        for vec in frozenset(h_int)
+        vec: -1 * h_int.get(vec, 0) * density_matrix[vec] for vec in frozenset(h_int)
     }
     return add_tb(direct, exchange)
 
 
-def fermi_on_grid(vals, filling):
+def fermi_on_grid(vals: np.ndarray, filling: float) -> float:
     """Compute the Fermi energy on a grid of k-points.
 
     Parameters
     ----------
-    vals : ndarray
-        Eigenvalues of the hamiltonian in k-space of shape (len(dim), norbs, norbs)
-    filling : int
-         Number of particles in a unit cell.
-
+    vals :
+        Eigenvalues of a hamiltonian sampled on a k-point grid with shape (nk, nk, ..., ndof, ndof),
+        where ndof is number of internal degrees of freedom.
+    filling :
+        Number of particles in a unit cell.
+        Used to determine the Fermi level.
     Returns
     -------
-    fermi_energy : float
-         Fermi energy
+    :
+        Fermi energy
     """
     norbs = vals.shape[-1]
     vals_flat = np.sort(vals.flatten())

pymf/model.py

@@ -12,44 +12,46 @@ def _check_hermiticity(h):
         op_vector = tuple(-1 * np.array(vector))
         op_vector = tuple(-1 * np.array(vector))
         if not np.allclose(h[vector], h[op_vector].conj().T):
-            raise ValueError("Hamiltonian is not Hermitian.")
+            raise ValueError("Tight-binding dictionary must be hermitian.")
 
 
 def _tb_type_check(tb):
     for count, key in enumerate(tb):
         if not isinstance(tb[key], np.ndarray):
-            raise ValueError("Inputted dictionary values are not np.ndarray's")
+            raise ValueError(
+                "Values of the tight-binding dictionary must be numpy arrays"
+            )
         shape = tb[key].shape
         if count == 0:
             size = shape[0]
         if not len(shape) == 2:
-            raise ValueError("Inputted dictionary values are not square matrices")
+            raise ValueError(
+                "Values of the tight-binding dictionary must be square matrices"
+            )
         if not size == shape[0]:
-            raise ValueError("Inputted dictionary elements shapes are not consistent")
+            raise ValueError(
+                "Values of the tight-binding dictionary must have consistent shape"
+            )
 
 
 class Model:
     """
-    Data class which defines the mean-field tight-binding problem
-    and computes the mean-field Hamiltonian.
+    Data class which defines the interacting tight-binding problem.
 
     Parameters
     ----------
     h_0 :
-        Non-interacting Hamiltonian.
+        Non-interacting hermitian Hamiltonian tight-binding dictionary.
     h_int :
-        Interaction Hamiltonian.
+        Interaction hermitian Hamiltonian tight-binding dictionary.
+        The interaction must be of density-density type, h_int[R][i, j] * c_i^dagger(R) c_j^dagger(0) c_j(0) c_i(R).
+        For example in 1D system with ndof internal degrees of freedom,
+        h_int[(2,)] = U * np.ones((ndof, ndof)) is a Coulomb repulsion interaction
+        with strength U between unit cells separated by 2 lattice vectors, where
+        the interaction is the same between all internal degrees of freedom.
     filling :
-        Filling of the system.
-
-    Attributes
-    ----------
-    h_0 :
-        Non-interacting Hamiltonian.
-    h_int :
-        Interaction Hamiltonian.
-    filling :
-        Filling of the system.
+        Number of particles in a unit cell.
+        Used to determine the Fermi level.
     """
 
     def __init__(self, h_0: tb_type, h_int: tb_type, filling: float) -> None:
@@ -71,24 +73,24 @@ class Model:
         _check_hermiticity(h_0)
         _check_hermiticity(h_int)
 
-    def mfield(self, mf_tb: tb_type, nk: int = 200) -> tb_type:
-        """Compute single mean field iteration.
+    def mfield(self, mf: tb_type, nk: int = 200) -> tb_type:
+        """Computes a new mean-field correction from a given one.
 
         Parameters
         ----------
-        mf_tb :
-            Mean-field tight-binding model.
+        mf :
+            Initial mean-field correction tight-binding dictionary.
         nk :
-            Number of k-points in the grid along a single direction.
+            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
         -------
         :
-            New mean-field tight-binding model.
+            new mean-field correction tight-binding dictionary.
         """
         rho, fermi_energy = construct_density_matrix(
-            add_tb(self.h_0, mf_tb), self.filling, nk
+            add_tb(self.h_0, mf), self.filling, nk
         )
         return add_tb(
             meanfield(rho, self.h_int),

pymf/observables.py

 import numpy as np
+
 from pymf.tb.tb import tb_type
 
 
@@ -8,9 +9,9 @@ def expectation_value(density_matrix: tb_type, observable: tb_type) -> complex:
     Parameters
     ----------
     density_matrix :
-        Density matrix in tight-binding format.
+        Density matrix tight-binding dictionary.
     observable :
-        Observable in tight-binding format.
+        Observable tight-binding dictionary.
 
     Returns
     -------

pymf/params/param_transforms.py

 import numpy as np
+
 from pymf.tb.tb import tb_type
 
 
 def tb_to_flat(tb: tb_type) -> np.ndarray:
-    """Convert a hermitian tight-binding dictionary to flat complex matrix.
+    """Parametrise a hermitian tight-binding dictionary by a flat complex vector.
 
     Parameters
     ----------
     tb :
-        Hermitian tigh-binding model
+        Hermitian tigh-binding dictionary
 
     Returns
     -------
-    flat :
-        1D complex array that parametrises the tb model.
+    :
+        1D complex array that parametrises the tight-binding dictionary.
     """
     if len(list(tb)[0]) == 0:
         matrix = np.array(list(tb.values()))
@@ -25,7 +26,7 @@ def tb_to_flat(tb: tb_type) -> np.ndarray:
 
 
 def flat_to_tb(
-    flat: np.ndarray,
+    tb_param_complex: np.ndarray,
     shape: tuple[int, int],
     tb_keys: list[tuple[None] | tuple[int, ...]],
 ) -> tb_type:
@@ -35,28 +36,28 @@ def flat_to_tb(
 
     Parameters
     ----------
-    flat :
+    tb_param_complex :
         1d complex array that parametrises the tb model.
     shape :
         Tuple (n, n) where n is the number of internal degrees of freedom
         (e.g. orbitals, spin, sublattice) within the tight-binding model.
     tb_keys :
-        A list of the keys within the tight-binding model (all the hoppings).
+        List of keys of the tight-binding dictionary.
 
     Returns
     -------
     tb :
-        tight-binding model
+        tight-binding dictionary
     """
     if len(tb_keys[0]) == 0:
         matrix = np.zeros((shape[-1], shape[-2]), dtype=complex)
-        matrix[np.triu_indices(shape[-1])] = flat
+        matrix[np.triu_indices(shape[-1])] = tb_param_complex
         matrix += matrix.conj().T
         matrix[np.diag_indices(shape[-1])] /= 2
         return {(): matrix}
     matrix = np.zeros(shape, dtype=complex)
     N = len(tb_keys) // 2 + 1
-    matrix[:N] = flat.reshape(N, *shape[1:])
+    matrix[:N] = tb_param_complex.reshape(N, *shape[1:])
     matrix[N:] = np.moveaxis(matrix[-(N + 1) :: -1], -1, -2).conj()
 
     tb_keys = np.array(list(tb_keys))

pymf/params/rparams.py

+import numpy as np
+
 from pymf.params.param_transforms import (
     complex_to_real,
     flat_to_tb,
     real_to_complex,
     tb_to_flat,
 )
-import numpy as np
 from pymf.tb.tb import tb_type
 
 
 def tb_to_rparams(tb: tb_type) -> np.ndarray:
-    """Convert a mean-field tight-binding model to a set of real parameters.
+    """Parametrise a hermitian tight-binding dictionary by a real vector.
 
     Parameters
     ----------
     tb :
-        Mean-field tight-binding model.
+        tight-binding dictionary.
 
     Returns
     -------
     :
-        1D real vector that parametrises the tb model.
+        1D real vector that parametrises the tight-binding dictionary.
     """
     return complex_to_real(tb_to_flat(tb))
 
 
 def rparams_to_tb(
-    r_params: np.ndarray, key_list: list[tuple[None] | tuple[int, ...]], size: int
+    tb_params: np.ndarray, tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int
 ) -> tb_type:
-    """Extract mean-field tight-binding model from a set of real parameters.
+    """Extract a hermitian tight-binding dictionary from a real vector parametrisation.
 
     Parameters
     ----------
-    r_params :
-        Real parameters.
-    key_list :
-        List of the keys within the tight-binding model (all the hoppings).
-    size :
-        Number of internal degrees of freedom (e.g. orbitals, spin, sublattice) within
-        the tight-binding model.
+    tb_params :
+        1D real array that parametrises the tight-binding dictionary.
+    tb_keys :
+        List of keys of the tight-binding dictionary.
+    ndof :
+        Number internal degrees of freedom within the unit cell.
 
     Returns
     -------
     :
-        Mean-field tight-binding model.
+        Tight-biding dictionary.
     """
-    flat_matrix = real_to_complex(r_params)
-    return flat_to_tb(flat_matrix, (len(key_list), size, size), key_list)
+    flat_matrix = real_to_complex(tb_params)
+    return flat_to_tb(flat_matrix, (len(tb_keys), ndof, ndof), tb_keys)

pymf/solvers.py

@@ -9,52 +9,52 @@ from pymf.model import Model
 from pymf.tb.utils import calculate_fermi_energy
 
 
-def cost(mf_param: np.ndarray, model: Model, nk: Optional[int] = 100) -> np.ndarray:
-    """Define the cost function for fixed point iteration.
+def cost(mf_param: np.ndarray, model: Model, nk: int = 100) -> np.ndarray:
+    """Defines the cost function for root solver.
 
-    The cost function is the difference between the input mean-field real space
-    parametrisation and a new mean-field.
+    The cost function is the difference between the computed and inputted mean-field.
 
     Parameters
     ----------
     mf_param :
-        1D real array that parametrises the mean-field tight-binding correction.
+        1D real array that parametrises the mean-field correction.
     Model :
-        Object which defines the mean-field problem to solve.
+        Interacting tight-binding problem definition.
     nk :
-        The number of k-points to use in the grid.
+        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
     -------
     :
-        1D real array which contains the difference between the input mean-field
-        and the new mean-field.
+        1D real array that is the difference between the computed and inputted mean-field
+        parametrisations
     """
     shape = model._size
-    mf_tb = rparams_to_tb(mf_param, list(model.h_int), shape)
-    mf_tb_new = model.mfield(mf_tb, nk=nk)
-    mf_params_new = tb_to_rparams(mf_tb_new)
+    mf = rparams_to_tb(mf_param, list(model.h_int), shape)
+    mf_new = model.mfield(mf, nk=nk)
+    mf_params_new = tb_to_rparams(mf_new)
     return mf_params_new - mf_param
 
 
 def solver(
     model: Model,
     mf_guess: np.ndarray,
-    nk: Optional[int] = 100,
+    nk: int = 100,
     optimizer: Optional[Callable] = scipy.optimize.anderson,
-    optimizer_kwargs: Optional[dict[str, str]] = {},
+    optimizer_kwargs: Optional[dict[str, str]] = {"M": 0},
 ) -> tb_type:
-    """Solve the mean-field self-consistent equation.
+    """Solve for the mean-field correction through self-consistent root finding.
 
     Parameters
     ----------
     model :
-        The model object.
+        Interacting tight-binding problem definition.
     mf_guess :
-        The initial guess for the mean-field tight-binding model.
+        The initial guess for the mean-field correction in the tight-binding dictionary format.
     nk :
-        The number of k-points to use in the grid. The default is 100. In the
-        0-dimensional case, this parameter is ignored.
+        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.
     optimizer :
         The solver used to solve the fixed point iteration.
     optimizer_kwargs :
@@ -63,7 +63,7 @@ def solver(
     Returns
     -------
     :
-        The mean-field tight-binding model.
+        Mean-field correction solution in the tight-binding dictionary format.
     """
     shape = model._size
     mf_params = tb_to_rparams(mf_guess)

pymf/tb/tb.py

@@ -4,37 +4,37 @@ tb_type = dict[tuple[None] | tuple[int, ...], np.ndarray]
 
 
 def add_tb(tb1: tb_type, tb2: tb_type) -> tb_type:
-    """Add up two tight-binding models together.
+    """Add up two tight-binding dictionaries together.
 
     Parameters
     ----------
     tb1 :
-        Tight-binding model.
+        Tight-binding dictionary.
     tb2 :
-        Tight-binding model.
+        Tight-binding dictionary.
 
     Returns
     -------
     :
-        Sum of the two tight-binding models.
+        Sum of the two tight-binding dictionaries.
     """
     return {k: tb1.get(k, 0) + tb2.get(k, 0) for k in frozenset(tb1) | frozenset(tb2)}
 
 
 def scale_tb(tb: tb_type, scale: float) -> tb_type:
-    """Scale a tight-binding model.
+    """Scale a tight-binding dictionary by a constant.
 
     Parameters
     ----------
-    tb : dict
-        Tight-binding model.
-    scale : float
-        The scaling factor.
+    tb :
+        Tight-binding dictionary.
+    scale :
+        Constant to scale the tight-binding dictionary by.
 
     Returns
     -------
     :
-        Scaled tight-binding model.
+        Scaled tight-binding dictionary.
     """
     return {k: tb.get(k, 0) * scale for k in frozenset(tb)}
 

pymf/tb/transforms.py

@@ -7,21 +7,24 @@ ks_type = Optional[np.ndarray]
 
 
 def tb_to_khamvector(tb: tb_type, nk: int, ks: ks_type = None) -> np.ndarray:
-    """Real-space tight-binding model to hamiltonian on k-space grid.
+    """Evaluate a tight-binding dictionary on a k-space grid.
 
     Parameters
     ----------
     tb :
-        A dictionary with real-space vectors as keys and complex np.arrays as values.
+        Tight-binding dictionary to evaluate on a k-space grid.
     nk :
-        Number of k-points along each direction.
+        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.
     ks :
-        Set of k-points. Repeated for all directions.
+        Set of points along which to evalaute the k-point grid. Repeated for all dimensions.
+        If not provided, a linear grid is used based on nk.
 
     Returns
     -------
     :
-        Hamiltonian evaluated on a k-point grid.
+        Tight-binding dictionary evaluated on a k-space grid.
+        Has shape (nk, nk, ..., ndof, ndof), where ndof is number of internal degrees of freedom.
     """
     ndim = len(list(tb)[0])
     if ks is None:
@@ -41,17 +44,18 @@ def tb_to_khamvector(tb: tb_type, nk: int, ks: ks_type = None) -> np.ndarray:
 
 
 def ifftn_to_tb(ifft_array: np.ndarray) -> tb_type:
-    """Convert an array from ifftn to a tight-binding model format.
+    """
+    Convert the result of `scipy.fft.ifftn` to a tight-binding dictionary.
 
     Parameters
     ----------
     ifft_array :
-        An array obtained from ifftn.
-
+        Result of `scipy.fft.ifftn` to convert to a tight-binding dictionary.
+        The input to `scipy.fft.ifftn` should be from `tb_to_khamvector`.
     Returns
     -------
     :
-        A dictionary with real-space vectors as keys and complex np.arrays as values.
+        Tight-binding dictionary.
     """
     size = ifft_array.shape[:-2]
 
@@ -62,21 +66,20 @@ def ifftn_to_tb(ifft_array: np.ndarray) -> tb_type:
 
 def kham_to_tb(
     kham: np.ndarray,
-    hopping_vecs: list[tuple[None] | tuple[int, ...]],
+    tb_keys: list[tuple[None] | tuple[int, ...]],
     ks: ks_type = None,
 ) -> tb_type:
-    """Extract hopping matrices from Bloch Hamiltonian.
+    """Convert a Hamiltonian evaluated on a k-grid to a tight-binding dictionary.
 
     Parameters
     ----------
     kham :
-        Bloch Hamiltonian matrix kham[k_x, ..., k_n, i, j]
-    hopping_vecs :
-        List of hopping vectors, will be the keys to the tb.
+        Hamiltonian sampled on a grid of k-points with shape (nk, nk, ..., ndof, ndof),
+        where ndof is number of internal degrees of freedom.
+    tb_keys :
+        List of keys for which to compute the tight-binding dictionary.
     ks :
-        Set of k-points. Repeated for all directions.
-        If system is finite, this option is ignored.
-
+        I have no clue why we need this, so this goes on the chopping board.
     Returns
     -------
     :
@@ -94,16 +97,16 @@ def kham_to_tb(
         hopps = (
             np.einsum(
                 "ij,jkl->ikl",
-                np.exp(1j * np.einsum("ij,jk->ik", hopping_vecs, k_grid)),
+                np.exp(1j * np.einsum("ij,jk->ik", tb_keys, k_grid)),
                 kham,
             )
             * (dk / (2 * np.pi)) ** ndim
         )
 
-        h_0 = {}
-        for i, vector in enumerate(hopping_vecs):
-            h_0[tuple(vector)] = hopps[i]
+        h = {}
+        for i, vector in enumerate(tb_keys):
+            h[tuple(vector)] = hopps[i]
 
-        return h_0
+        return h
     else:
         return {(): kham}

pymf/tb/utils.py

@@ -7,25 +7,25 @@ from pymf.tb.transforms import tb_to_khamvector
 
 
 def generate_guess(
-    vectors: list[tuple[None] | tuple[int, ...]], ndof: int, scale: float = 1
+    tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int, scale: float = 1
 ) -> tb_type:
-    """Generate guess for a tight-binding model.
+    """Generate guess tight-binding dictionary.
 
     Parameters
     ----------
-    vectors :
-        List of hopping vectors.
+    tb_keys :
+       List of tight-binding dictionary keys the guess contains.
     ndof :
-        Number internal degrees of freedom (e.g. orbitals, spin, sublattice),
+        Number internal degrees of freedom within the unit cell.
     scale :
-        Scale of the random guess. Default is 1.
+        Scale of the random guess.
     Returns
     -------
     :
-        Guess in the form of a tight-binding model.
+        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)
@@ -41,25 +41,43 @@ def generate_guess(
     return guess
 
 
-def generate_vectors(cutoff: int, dim: int) -> list[tuple[None] | tuple[int, ...]]:
-    """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 :
-        Maximum distance along each direction.
+        Maximum distance along each dimension to generate tight-bindign dictionary keys for.
     dim :
-        Dimension of the vectors.
+        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: tb_type, filling: float, nk: int = 100):
-    """Calculate the Fermi energy for a given filling."""
+    """
+    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_khamvector(tb, nk, ks=None)
     vals = np.linalg.eigvalsh(kham)
     return fermi_on_grid(vals, filling)

pymf/tests/test_zero_hint.py

@@ -21,7 +21,7 @@ def test_zero_hint(seed):
     dim = np.random.randint(0, 3)
     ndof = np.random.randint(2, 10)
     filling = np.random.randint(1, ndof)
-    random_hopping_vecs = utils.generate_vectors(cutoff, dim)
+    random_hopping_vecs = utils.generate_tb_keys(cutoff, dim)
 
     zero_key = tuple([0] * dim)
     h_0_random = utils.generate_guess(random_hopping_vecs, ndof, scale=1)