fix docstrings and use type hints

pymf/model.py

@@ -4,7 +4,7 @@ from pymf.mf import (
     construct_density_matrix,
     meanfield,
 )
-from pymf.tb.tb import add_tb
+from pymf.tb.tb import add_tb, tb_type
 
 
 def _check_hermiticity(h):
@@ -29,7 +29,30 @@ def _tb_type_check(tb):
 
 
 class Model:
-    def __init__(self, h_0, h_int, filling):
+    """
+    Data class which defines the mean-field tight-binding problem
+    and computes the mean-field Hamiltonian.
+
+    Parameters
+    ----------
+    h_0 :
+        Non-interacting Hamiltonian.
+    h_int :
+        Interaction Hamiltonian.
+    filling :
+        Filling of the system.
+
+    Attributes
+    ----------
+    h_0 :
+        Non-interacting Hamiltonian.
+    h_int :
+        Interaction Hamiltonian.
+    filling :
+        Filling of the system.
+    """
+
+    def __init__(self, h_0: tb_type, h_int: tb_type, filling: float) -> None:
         _tb_type_check(h_0)
         self.h_0 = h_0
         _tb_type_check(h_int)
@@ -48,19 +71,20 @@ class Model:
         _check_hermiticity(h_0)
         _check_hermiticity(h_int)
 
-    def mfield(self, mf_tb, nk=200):  # method or standalone?
+    def mfield(self, mf_tb: tb_type, nk: int = 200) -> tb_type:
         """Compute single mean field iteration.
 
         Parameters
         ----------
-        mf_tb : dict
+        mf_tb :
             Mean-field tight-binding model.
-        nk : int
-            Number of k-points in the grid.
+        nk :
+            Number of k-points in the grid along a single direction.
+            If the system is 0-dimensional (finite), this parameter is ignored.
 
         Returns
         -------
-        dict
+        :
             New mean-field tight-binding model.
         """
         rho, fermi_energy = construct_density_matrix(

pymf/observables.py

 import numpy as np
+from pymf.tb.tb import tb_type
 
 
-def expectation_value(density_matrix, observable):
+def expectation_value(density_matrix: tb_type, observable: tb_type) -> complex:
     """Compute the expectation value of an observable with respect to a density matrix.
 
     Parameters
     ----------
-    density_matrix : dict
+    density_matrix :
         Density matrix in tight-binding format.
-    observable : dict
+    observable :
         Observable in tight-binding format.
 
     Returns
     -------
-    complex
+    :
         Expectation value.
     """
     return np.sum(

pymf/params/param_transforms.py

 import numpy as np
+from pymf.tb.tb import tb_type
 
 
-def tb_to_flat(tb):
+def tb_to_flat(tb: tb_type) -> np.ndarray:
     """Convert a hermitian tight-binding dictionary to flat complex matrix.
 
     Parameters
     ----------
-    tb : dict with nd-array elements
-        Hermitian tigh-binding dictionary
+    tb :
+        Hermitian tigh-binding model
 
     Returns
     -------
-    flat : complex 1d numpy array
-        Flattened tight-binding dictionary
+    flat :
+        1D complex array that parametrises the tb model.
     """
     if len(list(tb)[0]) == 0:
         matrix = np.array(list(tb.values()))
@@ -23,24 +24,29 @@ def tb_to_flat(tb):
     return sorted_vals[:N].flatten()
 
 
-def flat_to_tb(flat, shape, tb_keys):
+def flat_to_tb(
+    flat: np.ndarray,
+    shape: tuple[int, int],
+    tb_keys: list[tuple[None] | tuple[int, ...]],
+) -> tb_type:
     """Reverse operation to `tb_to_flat`.
 
     It takes a flat complex 1d array and return the tight-binding dictionary.
 
     Parameters
     ----------
-    flat : dict with nd-array elements
-        Hermitian tigh-binding dictionary
-    shape : tuple
-        shape of the tb elements
-    tb_keys : iterable
-        original tb key elements
+    flat :
+        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).
 
     Returns
     -------
-    tb : dict
-        tight-binding dictionary
+    tb :
+        tight-binding model
     """
     if len(tb_keys[0]) == 0:
         matrix = np.zeros((shape[-1], shape[-2]), dtype=complex)
@@ -59,17 +65,22 @@ def flat_to_tb(flat, shape, tb_keys):
     return tb
 
 
-def complex_to_real(z):
-    """Split real and imaginary parts of a complex array.
+def complex_to_real(z: np.ndarray) -> np.ndarray:
+    """Split and concatenate real and imaginary parts of a complex array.
 
     Parameters
     ----------
-    z : array
+    z :
         Complex array.
+
+    Returns
+    -------
+    :
+        Real array that concatenates the real and imaginary parts of the input array.
     """
     return np.concatenate((np.real(z), np.imag(z)))
 
 
-def real_to_complex(z):
+def real_to_complex(z: np.ndarray) -> np.ndarray:
     """Undo `complex_to_real`."""
     return z[: len(z) // 2] + 1j * z[len(z) // 2 :]

pymf/params/rparams.py

@@ -4,40 +4,44 @@ from pymf.params.param_transforms import (
     real_to_complex,
     tb_to_flat,
 )
+import numpy as np
+from pymf.tb.tb import tb_type
 
 
-def tb_to_rparams(tb):
+def tb_to_rparams(tb: tb_type) -> np.ndarray:
     """Convert a mean-field tight-binding model to a set of real parameters.
 
     Parameters
     ----------
-    tb : dict
+    tb :
         Mean-field tight-binding model.
 
     Returns
     -------
-    dict
-        Real parameters.
+    :
+        1D real vector that parametrises the tb model.
     """
-    return complex_to_real(tb_to_flat(tb))  # placeholder for now
+    return complex_to_real(tb_to_flat(tb))
 
 
-def rparams_to_tb(r_params, key_list, size):
+def rparams_to_tb(
+    r_params: np.ndarray, key_list: list[tuple[None] | tuple[int, ...]], size: int
+) -> tb_type:
     """Extract mean-field tight-binding model from a set of real parameters.
 
     Parameters
     ----------
-    r_params : dict
+    r_params :
         Real parameters.
-    key_list : list
-        List of the keys of the mean-field tight-binding model, meaning all the
-        hoppings.
-    size : tuple
-        Shape of the mean-field tight-binding model.
+    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.
 
     Returns
     -------
-    dict
+    :
         Mean-field tight-binding model.
     """
     flat_matrix = real_to_complex(r_params)

pymf/solvers.py

 from functools import partial
-
 import numpy as np
 import scipy
+from typing import Optional, Callable
 
 from pymf.params.rparams import rparams_to_tb, tb_to_rparams
-from pymf.tb.tb import add_tb
+from pymf.tb.tb import add_tb, tb_type
+from pymf.model import Model
 from pymf.tb.utils import calculate_fermi_energy
 
 
-def cost(mf_param, Model, nk=100):
+def cost(mf_param: np.ndarray, model: Model, nk: Optional[int] = 100) -> np.ndarray:
     """Define the cost function for fixed point iteration.
 
     The cost function is the difference between the input mean-field real space
@@ -16,50 +17,59 @@ def cost(mf_param, Model, nk=100):
 
     Parameters
     ----------
-    mf_param : numpy.array
-        The mean-field real space parametrisation.
-    Model : Model
-        The model object.
-    nk : int, optional
-        The number of k-points to use in the grid. The default is 100.
+    mf_param :
+        1D real array that parametrises the mean-field tight-binding correction.
+    Model :
+        Object which defines the mean-field problem to solve.
+    nk :
+        The number of k-points to use in the grid.
+
+    Returns
+    -------
+    :
+        1D real array which contains the difference between the input mean-field
+        and the new mean-field.
     """
-    shape = Model._size
-    mf_tb = rparams_to_tb(mf_param, list(Model.h_int), shape)
-    mf_tb_new = Model.mfield(mf_tb, nk=nk)
+    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)
     return mf_params_new - mf_param
 
 
 def solver(
-    Model, mf_guess, nk=100, optimizer=scipy.optimize.anderson, optimizer_kwargs={}
-):
+    model: Model,
+    mf_guess: np.ndarray,
+    nk: Optional[int] = 100,
+    optimizer: Optional[Callable] = scipy.optimize.anderson,
+    optimizer_kwargs: Optional[dict[str, str]] = {},
+) -> tb_type:
     """Solve the mean-field self-consistent equation.
 
     Parameters
     ----------
-    Model : Model
+    model :
         The model object.
-    mf_guess : numpy.array
+    mf_guess :
         The initial guess for the mean-field tight-binding model.
-    nk : int, optional
+    nk :
         The number of k-points to use in the grid. The default is 100. In the
         0-dimensional case, this parameter is ignored.
-    optimizer : scipy.optimize, optional
-        The optimizer to use to solve for fixed-points. The default is
-        scipy.optimize.anderson.
-    optimizer_kwargs : dict, optional
-        The keyword arguments to pass to the optimizer. The default is {}.
+    optimizer :
+        The solver used to solve the fixed point iteration.
+    optimizer_kwargs :
+        The keyword arguments to pass to the optimizer.
 
     Returns
     -------
-    result : numpy.array
+    :
         The mean-field tight-binding model.
     """
-    shape = Model._size
+    shape = model._size
     mf_params = tb_to_rparams(mf_guess)
-    f = partial(cost, Model=Model, nk=nk)
+    f = partial(cost, model=model, nk=nk)
     result = rparams_to_tb(
-        optimizer(f, mf_params, **optimizer_kwargs), list(Model.h_int), shape
+        optimizer(f, mf_params, **optimizer_kwargs), list(model.h_int), shape
     )
-    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)})
+    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)})

pymf/tb/tb.py

 import numpy as np
 
+tb_type = dict[tuple[None] | tuple[int, ...], np.ndarray]
 
-def add_tb(tb1, tb2):
+
+def add_tb(tb1: tb_type, tb2: tb_type) -> tb_type:
     """Add up two tight-binding models together.
 
     Parameters
     ----------
-    tb1 : dict
+    tb1 :
         Tight-binding model.
-    tb2 : dict
+    tb2 :
         Tight-binding model.
 
     Returns
     -------
-    dict
+    :
         Sum of the two tight-binding models.
     """
     return {k: tb1.get(k, 0) + tb2.get(k, 0) for k in frozenset(tb1) | frozenset(tb2)}
 
 
-def scale_tb(tb, scale):
+def scale_tb(tb: tb_type, scale: float) -> tb_type:
     """Scale a tight-binding model.
 
     Parameters
@@ -31,13 +33,13 @@ def scale_tb(tb, scale):
 
     Returns
     -------
-    dict
+    :
         Scaled tight-binding model.
     """
     return {k: tb.get(k, 0) * scale for k in frozenset(tb)}
 
 
-def compare_dicts(dict1, dict2, atol=1e-10):
+def compare_dicts(dict1: dict, dict2: dict, atol: float = 1e-10) -> None:
     """Compare two dictionaries."""
     for key in frozenset(dict1) | frozenset(dict2):
         assert np.allclose(dict1[key], dict2[key], atol=atol)

pymf/tb/transforms.py

 import itertools
 import numpy as np
+from typing import Optional
+from pymf.tb.tb import tb_type
 
+ks_type = Optional[np.ndarray]
 
-def tb_to_khamvector(tb, nk, ks=None):
+
+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.
 
     Parameters
     ----------
-    tb : dict
+    tb :
         A dictionary with real-space vectors as keys and complex np.arrays as values.
-    nk : int
+    nk :
         Number of k-points along each direction.
-    ks : 1D-array
+    ks :
         Set of k-points. Repeated for all directions.
 
     Returns
     -------
-    ndarray
+    :
         Hamiltonian evaluated on a k-point grid.
-
     """
     ndim = len(list(tb)[0])
     if ks is None:
@@ -37,17 +40,17 @@ def tb_to_khamvector(tb, nk, ks=None):
     return np.sum(tb_array * k_dependency, axis=0)
 
 
-def ifftn_to_tb(ifft_array):
+def ifftn_to_tb(ifft_array: np.ndarray) -> tb_type:
     """Convert an array from ifftn to a tight-binding model format.
 
     Parameters
     ----------
-    ifft_array : ndarray
+    ifft_array :
         An array obtained from ifftn.
 
     Returns
     -------
-    dict
+    :
         A dictionary with real-space vectors as keys and complex np.arrays as values.
     """
     size = ifft_array.shape[:-2]
@@ -57,22 +60,26 @@ def ifftn_to_tb(ifft_array):
     return {tuple(k): ifft_array[tuple(k)] for k in keys}
 
 
-def kham_to_tb(kham, hopping_vecs, ks=None):
+def kham_to_tb(
+    kham: np.ndarray,
+    hopping_vecs: list[tuple[None] | tuple[int, ...]],
+    ks: ks_type = None,
+) -> tb_type:
     """Extract hopping matrices from Bloch Hamiltonian.
 
     Parameters
     ----------
-    kham : nd-array
+    kham :
         Bloch Hamiltonian matrix kham[k_x, ..., k_n, i, j]
-    hopping_vecs : list
+    hopping_vecs :
         List of hopping vectors, will be the keys to the tb.
-    ks : 1D-array
-        Set of k-points. Repeated for all directions. If the system is finite,
-        ks=None`.
+    ks :
+        Set of k-points. Repeated for all directions.
+        If system is finite, this option is ignored.
 
     Returns
     -------
-    scf_model : dict
+    :
         Tight-binding model of Hartree-Fock solution.
     """
     if ks is not None:

pymf/tb/utils.py

 from itertools import product
-
 import numpy as np
 
+from pymf.tb.tb import tb_type
 from pymf.mf import fermi_on_grid
 from pymf.tb.transforms import tb_to_khamvector
 
 
-def generate_guess(vectors, ndof, scale=1):
+def generate_guess(
+    vectors: list[tuple[None] | tuple[int, ...]], ndof: int, scale: float = 1
+) -> tb_type:
     """Generate guess for a tight-binding model.
 
     Parameters
     ----------
-    vectors : list
+    vectors :
         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.
-
+    ndof :
+        Number internal degrees of freedom (e.g. orbitals, spin, sublattice),
+    scale :
+        Scale of the random guess. Default is 1.
     Returns
     -------
-    guess : tb dictionary
+    :
         Guess in the form of a tight-binding model.
     """
     guess = {}
@@ -40,14 +41,14 @@ def generate_guess(vectors, ndof, scale=1):
     return guess
 
 
-def generate_vectors(cutoff, dim):
+def generate_vectors(cutoff: int, dim: int) -> list[tuple[None] | tuple[int, ...]]:
     """Generate hopping vectors up to a cutoff.
 
     Parameters
     ----------
-    cutoff : int
+    cutoff :
         Maximum distance along each direction.
-    dim : int
+    dim :
         Dimension of the vectors.
 
     Returns
@@ -57,7 +58,7 @@ def generate_vectors(cutoff, dim):
     return [*product(*([[*range(-cutoff, cutoff + 1)]] * dim))]
 
 
-def calculate_fermi_energy(tb, filling, nk=100):
+def calculate_fermi_energy(tb: tb_type, filling: float, nk: int = 100):
     """Calculate the Fermi energy for a given filling."""
     kham = tb_to_khamvector(tb, nk, ks=None)
     vals = np.linalg.eigvalsh(kham)