pymf/__init__.py → meanfi/__init__.py

@@ -6,10 +6,30 @@ except ImportError:
     __version__ = "unknown"
     __version_tuple__ = (0, 0, "unknown", "unknown")
 
-from .mf import construct_density_matrix
+from .mf import (
+    density_matrix,
+    meanfield,
+)
+from .solvers import solver
+from .model import Model
+from .observables import expectation_value
+from .tb.tb import add_tb, scale_tb
+from .tb.transforms import tb_to_kgrid, kgrid_to_tb
+from .tb.utils import guess_tb, fermi_energy
+
 
 __all__ = [
-    "construct_density_matrix",
+    "solver",
+    "Model",
+    "expectation_value",
+    "add_tb",
+    "scale_tb",
+    "guess_tb",
+    "fermi_energy",
+    "density_matrix",
+    "meanfield",
+    "tb_to_kgrid",
+    "kgrid_to_tb",
     "__version__",
     "__version_tuple__",
 ]

pymf/kwant_helper/kwant_examples.py → meanfi/kwant_helper/kwant_examples.py

 import kwant
 import numpy as np
 
-from pymf.kwant_helper.utils import build_interacting_syst
+from meanfi.kwant_helper.utils import build_interacting_syst
 
 s0 = np.identity(2)
 sz = np.diag([1, -1])

pymf/kwant_helper/utils.py → meanfi/kwant_helper/utils.py

 import inspect
 from copy import copy
 from itertools import product
+from typing import Callable
 
-import kwant
 import numpy as np
 from scipy.sparse import coo_array
+import kwant
+import kwant.lattice
+import kwant.builder
 
+from meanfi.tb.tb import _tb_type
 
-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.Builder, params: dict = {}, return_data: bool = False
+) -> _tb_type:
+    """Construct a tight-binding dictionary from a `kwant.builder.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
+    :
+        Tight-binding dictionary that corresponds to the builder.
+    :
         Data with sites and number of orbitals. Only if `return_data=True`.
     """
     builder = copy(builder)
@@ -45,8 +52,8 @@ def builder_to_tb(builder, params={}, return_data=False):
         col = copy(row)
         row, col = np.array([*product(row, col)]).T
         try:
+            _params = {}
             for arg in inspect.getfullargspec(val).args:
-                _params = {}
                 if arg in params:
                     _params[arg] = params[arg]
             val = val(site, **_params)
@@ -73,8 +80,8 @@ def builder_to_tb(builder, params={}, return_data=False):
         ]
         row, col = np.array([*product(row, col)]).T
         try:
+            _params = {}
             for arg in inspect.getfullargspec(val).args:
-                _params = {}
                 if arg in params:
                     _params[arg] = params[arg]
             val = val(a, b, **_params)
@@ -124,28 +131,38 @@ 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.Builder,
+    lattice: kwant.lattice.Polyatomic,
+    func_onsite: Callable,
+    func_hop: Callable,
+    max_neighbor: int = 1,
+) -> kwant.builder.Builder:
+    """
+    Construct an auxiliary `kwant` system that encodes the interactions.
 
     Parameters
     ----------
-    builder : `kwant.Builder`
-        Non-interacting `kwant` system.
-    func_onsite : function
-        Onsite function.
-    func_hop : function
-        Hopping function.
-    max_neighbor : int
-        Maximal nearest-neighbor order.
+    builder :
+        Non-interacting `kwant.builder.Builder` system.
+    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.
+    :
+        Auxiliary `kwant.builder.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 = kwant.builder.Builder(
+        kwant.lattice.TranslationalSymmetry(*builder.symmetry.periods)
+    )
     int_builder[builder.sites()] = func_onsite
     for neighbors in range(max_neighbor):
         int_builder[lattice.neighbors(neighbors + 1)] = func_hop

pymf/mf.py → meanfi/mf.py

 import numpy as np
-from scipy.fftpack import ifftn
+from typing import Tuple
 
-from pymf.tb.tb import add_tb
-from pymf.tb.transforms import ifftn_to_tb, tb_to_khamvector
+from meanfi.tb.tb import add_tb, _tb_type
+from meanfi.tb.transforms import tb_to_kgrid, kgrid_to_tb
 
 
-def construct_density_matrix_kgrid(kham, filling):
+def density_matrix_kgrid(kham: np.ndarray, filling: float) -> Tuple[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)
+    fermi = fermi_on_kgrid(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):
-    """Compute the density matrix in real-space tight-binding format.
+def density_matrix(h: _tb_type, filling: float, nk: int) -> Tuple[_tb_type, float]:
+    """Compute the real-space density matrix tight-binding dictionary.
 
     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)
+        kham = tb_to_kgrid(h, nk=nk)
+        _density_matrix_krid, fermi = density_matrix_kgrid(kham, filling)
         return (
-            ifftn_to_tb(ifftn(rho_grid, axes=np.arange(ndim))),
+            kgrid_to_tb(_density_matrix_krid),
             fermi,
         )
     else:
-        rho, fermi = construct_density_matrix_kgrid(h[()], filling)
-        return {(): rho}, fermi
+        _density_matrix, fermi = 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.
     Returns
     -------
-    dict
-        Mean-field tb model.
+    :
+        Mean-field correction tight-binding dictionary.
+
+    Notes
+    -----
+
+    The interaction h_int must be of density-density type.
+    For example, h_int[(1,)][i, j] = V means a repulsive interaction
+    of strength V between two particles with internal degrees of freedom i and j
+    separated by 1 lattice vector.
     """
-    n = len(list(density_matrix_tb)[0])
+    n = len(list(density_matrix)[0])
     local_key = tuple(np.zeros((n,), dtype=int))
 
     direct = {
@@ -82,7 +94,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 +104,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_kgrid(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 → meanfi/model.py

 import numpy as np
 
-from pymf.mf import (
-    construct_density_matrix,
+from meanfi.mf import (
+    density_matrix,
     meanfield,
 )
-from pymf.tb.tb import add_tb
+from meanfi.tb.tb import add_tb, _tb_type
 
 
 def _check_hermiticity(h):
@@ -12,24 +12,52 @@ 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:
-    def __init__(self, h_0, h_int, filling):
+    """
+    Data class which defines the interacting tight-binding problem.
+
+    Parameters
+    ----------
+    h_0 :
+        Non-interacting hermitian Hamiltonian tight-binding dictionary.
+    h_int :
+        Interaction hermitian Hamiltonian tight-binding dictionary.
+    filling :
+        Number of particles in a unit cell.
+        Used to determine the Fermi level.
+
+    Notes
+    -----
+
+    The interaction h_int must be of density-density type.
+    For example, h_int[(1,)][i, j] = V means a repulsive interaction
+    of strength V between two particles with internal degrees of freedom i and j
+    separated by 1 lattice vector.
+    """
+
+    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)
@@ -42,31 +70,30 @@ class Model:
 
         _first_key = list(h_0)[0]
         self._ndim = len(_first_key)
-        self._size = h_0[_first_key].shape[0]
+        self._ndof = h_0[_first_key].shape[0]
         self._local_key = tuple(np.zeros((self._ndim,), dtype=int))
 
         _check_hermiticity(h_0)
         _check_hermiticity(h_int)
 
-    def mfield(self, mf_tb, nk=200):  # method or standalone?
-        """Compute single mean field iteration.
+    def mfield(self, mf: _tb_type, nk: int = 20) -> _tb_type:
+        """Computes a new mean-field correction from a given one.
 
         Parameters
         ----------
-        mf_tb : dict
-            Mean-field tight-binding model.
-        nk : int
-            Number of k-points in the grid.
+        mf :
+            Initial mean-field correction tight-binding dictionary.
+        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
-            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
-        )
+        rho, fermi_energy = density_matrix(add_tb(self.h_0, mf), self.filling, nk)
         return add_tb(
             meanfield(rho, self.h_int),
-            {self._local_key: -fermi_energy * np.eye(self._size)},
+            {self._local_key: -fermi_energy * np.eye(self._ndof)},
         )

pymf/observables.py → meanfi/observables.py

 import numpy as np
 
+from meanfi.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 in tight-binding format.
-    observable : dict
-        Observable in tight-binding format.
+    density_matrix :
+        Density matrix tight-binding dictionary.
+    observable :
+        Observable tight-binding dictionary.
 
     Returns
     -------
-    complex
+    :
         Expectation value.
     """
     return np.sum(

pymf/params/param_transforms.py → meanfi/params/param_transforms.py

 import numpy as np
 
+from meanfi.tb.tb import _tb_type
 
-def tb_to_flat(tb):
-    """Convert a hermitian tight-binding dictionary to flat complex matrix.
+
+def tb_to_flat(tb: _tb_type) -> np.ndarray:
+    """Parametrise a hermitian tight-binding dictionary by a flat complex vector.
 
     Parameters
     ----------
-    tb : dict with nd-array elements
+    tb :
         Hermitian tigh-binding dictionary
 
     Returns
     -------
-    flat : complex 1d numpy array
-        Flattened tight-binding dictionary
+    :
+        1D complex array that parametrises the tight-binding dictionary.
     """
     if len(list(tb)[0]) == 0:
         matrix = np.array(list(tb.values()))
@@ -23,34 +25,39 @@ def tb_to_flat(tb):
     return sorted_vals[:N].flatten()
 
 
-def flat_to_tb(flat, shape, tb_keys):
+def flat_to_tb(
+    tb_param_complex: np.ndarray,
+    ndof: 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
+    tb_param_complex :
+        1d complex array that parametrises the tb model.
+    ndof :
+        Number internal degrees of freedom within the unit cell.
+    tb_keys :
+        List of keys of the tight-binding dictionary.
 
     Returns
     -------
-    tb : dict
+    tb :
         tight-binding dictionary
     """
+    shape = (len(tb_keys), ndof, ndof)
     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))
@@ -59,17 +66,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 :]

meanfi/params/rparams.py

+import numpy as np
+
+from meanfi.params.param_transforms import (
+    complex_to_real,
+    flat_to_tb,
+    real_to_complex,
+    tb_to_flat,
+)
+from meanfi.tb.tb import _tb_type
+
+
+def tb_to_rparams(tb: _tb_type) -> np.ndarray:
+    """Parametrise a hermitian tight-binding dictionary by a real vector.
+
+    Parameters
+    ----------
+    tb :
+        tight-binding dictionary.
+
+    Returns
+    -------
+    :
+        1D real vector that parametrises the tight-binding dictionary.
+    """
+    return complex_to_real(tb_to_flat(tb))
+
+
+def rparams_to_tb(
+    tb_params: np.ndarray, tb_keys: list[tuple[None] | tuple[int, ...]], ndof: int
+) -> _tb_type:
+    """Extract a hermitian tight-binding dictionary from a real vector parametrisation.
+
+    Parameters
+    ----------
+    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
+    -------
+    :
+        Tight-biding dictionary.
+    """
+    flat_matrix = real_to_complex(tb_params)
+    return flat_to_tb(flat_matrix, ndof, tb_keys)

meanfi/solvers.py

+from functools import partial
+import numpy as np
+import scipy
+from typing import Optional, Callable
+
+from meanfi.params.rparams import rparams_to_tb, tb_to_rparams
+from meanfi.tb.tb import add_tb, _tb_type
+from meanfi.model import Model
+from meanfi.tb.utils import fermi_energy
+
+
+def cost(mf_param: np.ndarray, model: Model, nk: int = 20) -> np.ndarray:
+    """Defines the cost function for root solver.
+
+    The cost function is the difference between the computed and inputted mean-field.
+
+    Parameters
+    ----------
+    mf_param :
+        1D real array that parametrises the mean-field correction.
+    Model :
+        Interacting tight-binding problem definition.
+    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
+    -------
+    :
+        1D real array that is the difference between the computed and inputted mean-field
+        parametrisations
+    """
+    shape = model._ndof
+    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: int = 20,
+    optimizer: Optional[Callable] = scipy.optimize.anderson,
+    optimizer_kwargs: Optional[dict[str, str]] = {"M": 0},
+) -> _tb_type:
+    """Solve for the mean-field correction through self-consistent root finding.
+
+    Parameters
+    ----------
+    model :
+        Interacting tight-binding problem definition.
+    mf_guess :
+        The initial guess for the mean-field correction in the tight-binding dictionary format.
+    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.
+    optimizer :
+        The solver used to solve the fixed point iteration.
+        Default uses `scipy.optimize.anderson`.
+    optimizer_kwargs :
+        The keyword arguments to pass to the optimizer.
+
+    Returns
+    -------
+    :
+        Mean-field correction solution in the tight-binding dictionary format.
+    """
+    shape = model._ndof
+    mf_params = tb_to_rparams(mf_guess)
+    f = partial(cost, model=model, nk=nk)
+    result = rparams_to_tb(
+        optimizer(f, mf_params, **optimizer_kwargs), list(model.h_int), shape
+    )
+    fermi = fermi_energy(add_tb(model.h_0, result), model.filling, nk=nk)
+    return add_tb(result, {model._local_key: -fermi * np.eye(model._ndof)})

pymf/tb/tb.py → meanfi/tb/tb.py

 import numpy as np
 
+_tb_type = dict[tuple[int, ...], np.ndarray]
 
-def add_tb(tb1, tb2):
-    """Add up two tight-binding models together.
+
+def add_tb(tb1: _tb_type, tb2: _tb_type) -> _tb_type:
+    """Add up two tight-binding dictionaries together.
 
     Parameters
     ----------
-    tb1 : dict
-        Tight-binding model.
-    tb2 : dict
-        Tight-binding model.
+    tb1 :
+        Tight-binding dictionary.
+    tb2 :
+        Tight-binding dictionary.
 
     Returns
     -------
-    dict
-        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, scale):
-    """Scale a tight-binding model.
+def scale_tb(tb: _tb_type, scale: float) -> _tb_type:
+    """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
     -------
-    dict
-        Scaled tight-binding model.
+    :
+        Scaled tight-binding dictionary.
     """
     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 → meanfi/tb/transforms.py

 import itertools
 import numpy as np
+from scipy.fftpack import ifftn
 
+from meanfi.tb.tb import _tb_type
 
-def tb_to_khamvector(tb, nk, ks=None):
-    """Real-space tight-binding model to hamiltonian on k-space grid.
+
+def tb_to_kgrid(tb: _tb_type, nk: int) -> np.ndarray:
+    """Evaluate a tight-binding dictionary on a k-space grid.
 
     Parameters
     ----------
-    tb : dict
-        A dictionary with real-space vectors as keys and complex np.arrays as values.
-    nk : int
-        Number of k-points along each direction.
-    ks : 1D-array
-        Set of k-points. Repeated for all directions.
+    tb :
+        Tight-binding dictionary to evaluate on a k-space grid.
+    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
     -------
-    ndarray
-        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:
-        ks = np.linspace(-np.pi, np.pi, nk, endpoint=False)
-        ks = np.concatenate((ks[nk // 2 :], ks[: nk // 2]), axis=0)  # shift for ifft
+    ks = np.linspace(-np.pi, np.pi, nk, endpoint=False)
+    ks = np.concatenate((ks[nk // 2 :], ks[: nk // 2]), axis=0)  # shift for ifft
     kgrid = np.meshgrid(*([ks] * ndim), indexing="ij")
 
     num_keys = len(list(tb.keys()))
@@ -37,66 +38,41 @@ def tb_to_khamvector(tb, nk, ks=None):
     return np.sum(tb_array * k_dependency, axis=0)
 
 
-def ifftn_to_tb(ifft_array):
-    """Convert an array from ifftn to a tight-binding model format.
+def kgrid_to_tb(kgrid_array: np.ndarray) -> _tb_type:
+    """
+    Convert a k-space grid array to a tight-binding dictionary.
 
     Parameters
     ----------
-    ifft_array : ndarray
-        An array obtained from ifftn.
-
+    kgrid_array :
+        K-space grid array to convert to a tight-binding dictionary.
+        The array should be of shape (nk, nk, ..., ndof, ndof),
+        where ndof is number of internal degrees of freedom.
     Returns
     -------
-    dict
-        A dictionary with real-space vectors as keys and complex np.arrays as values.
+    :
+        Tight-binding dictionary.
     """
-    size = ifft_array.shape[:-2]
-
-    keys = [np.arange(-size[0] // 2 + 1, size[0] // 2) for i in range(len(size))]
-    keys = itertools.product(*keys)
-    return {tuple(k): ifft_array[tuple(k)] for k in keys}
+    ndim = len(kgrid_array.shape) - 2
+    return ifftn_to_tb(ifftn(kgrid_array, axes=np.arange(ndim)))
 
 
-def kham_to_tb(kham, hopping_vecs, ks=None):
-    """Extract hopping matrices from Bloch Hamiltonian.
+def ifftn_to_tb(ifft_array: np.ndarray) -> _tb_type:
+    """
+    Convert the result of `scipy.fft.ifftn` to a tight-binding dictionary.
 
     Parameters
     ----------
-    kham : nd-array
-        Bloch Hamiltonian matrix kham[k_x, ..., k_n, i, j]
-    hopping_vecs : list
-        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`.
-
+    ifft_array :
+        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
     -------
-    scf_model : dict
-        Tight-binding model of Hartree-Fock solution.
+    :
+        Tight-binding dictionary.
     """
-    if ks is not None:
-        ndim = len(kham.shape) - 2
-        dk = np.diff(ks)[0]
-        nk = len(ks)
-        k_pts = np.tile(ks, ndim).reshape(ndim, nk)
-        k_grid = np.array(np.meshgrid(*k_pts))
-        k_grid = k_grid.reshape(k_grid.shape[0], np.prod(k_grid.shape[1:]))
-        kham = kham.reshape(np.prod(kham.shape[:ndim]), *kham.shape[-2:])
-
-        hopps = (
-            np.einsum(
-                "ij,jkl->ikl",
-                np.exp(1j * np.einsum("ij,jk->ik", hopping_vecs, k_grid)),
-                kham,
-            )
-            * (dk / (2 * np.pi)) ** ndim
-        )
-
-        h_0 = {}
-        for i, vector in enumerate(hopping_vecs):
-            h_0[tuple(vector)] = hopps[i]
+    size = ifft_array.shape[:-2]
 
-        return h_0
-    else:
-        return {(): kham}
+    keys = [np.arange(-size[0] // 2 + 1, size[0] // 2) for i in range(len(size))]
+    keys = itertools.product(*keys)
+    return {tuple(k): ifft_array[tuple(k)] for k in keys}

pymf/tb/utils.py → meanfi/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 meanfi.tb.tb import _tb_type
+from meanfi.mf import fermi_on_kgrid
+from meanfi.tb.transforms import tb_to_kgrid
 
 
-def generate_guess(vectors, ndof, scale=1):
-    """Generate guess for a tight-binding model.
+def guess_tb(
+    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 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)

pymf/tests/test_graphene.py → meanfi/tests/test_graphene.py

@@ -2,12 +2,14 @@
 import numpy as np
 import pytest
 
-from pymf.kwant_helper import kwant_examples, utils
-from pymf.model import Model
-from pymf.solvers import solver
-from pymf.tb.tb import add_tb
-from pymf.tb.transforms import tb_to_khamvector
-from pymf.tb.utils import generate_guess
+from meanfi.kwant_helper import kwant_examples, utils
+from meanfi import (
+    Model,
+    solver,
+    tb_to_kgrid,
+    guess_tb,
+    add_tb,
+)
 
 
 def compute_gap(tb, fermi_energy=0, nk=100):
@@ -27,7 +29,7 @@ def compute_gap(tb, fermi_energy=0, nk=100):
      gap : float
      Indirect gap.
     """
-    kham = tb_to_khamvector(tb, nk, ks=None)
+    kham = tb_to_kgrid(tb, nk)
     vals = np.linalg.eigvalsh(kham)
 
     emax = np.max(vals[vals <= fermi_energy])
@@ -35,7 +37,7 @@ def compute_gap(tb, fermi_energy=0, nk=100):
     return np.abs(emin - emax)
 
 
-repeat_number = 10
+repeat_number = 5
 # %%
 graphene_builder, int_builder = kwant_examples.graphene_extended_hubbard()
 h_0 = utils.builder_to_tb(graphene_builder)
@@ -72,12 +74,10 @@ def gap_prediction(U, V):
     nk = 40
 
     h_int = utils.builder_to_tb(int_builder, params)
-    guess = generate_guess(frozenset(h_int), len(list(h_0.values())[0]))
+    guess = guess_tb(frozenset(h_int), len(list(h_0.values())[0]))
     model = Model(h_0, h_int, filling)
 
-    mf_sol = solver(
-        model, guess, nk=nk, optimizer_kwargs={"verbose": True, "M": 0, "f_tol": 1e-8}
-    )
+    mf_sol = solver(model, guess, nk=nk, optimizer_kwargs={"M": 0, "f_tol": 1e-8})
     gap = compute_gap(add_tb(h_0, mf_sol), nk=200)
 
     # Check if the gap is predicted correctly

meanfi/tests/test_hat.py

+# %%
+import numpy as np
+import pytest
+
+from meanfi import (
+    Model,
+    solver,
+    guess_tb,
+    scale_tb,
+    add_tb,
+    expectation_value,
+    density_matrix,
+)
+
+from meanfi.tb.utils import generate_tb_keys
+
+
+# %%
+def total_energy(ham_tb, rho_tb):
+    return np.real(expectation_value(rho_tb, ham_tb))
+
+
+# %%
+U0 = 1
+filling = 2
+nk = 10
+repeat_number = 3
+ndof = 4
+cutoff = 1
+
+
+# %%
+@np.vectorize
+def mf_rescaled(alpha, mf0, h0):
+    hamiltonian = add_tb(h0, scale_tb(mf0, alpha))
+    rho, _ = density_matrix(hamiltonian, filling=filling, nk=nk)
+    hamiltonian = add_tb(h0, scale_tb(mf0, np.sign(alpha)))
+    return total_energy(hamiltonian, rho)
+
+
+@pytest.mark.parametrize("seed", range(repeat_number))
+def test_mexican_hat(seed):
+    np.random.seed(seed)
+    h0s = []
+    h_ints = []
+    for ndim in np.arange(4):
+        keys = generate_tb_keys(cutoff, ndim)
+        h0s.append(guess_tb(keys, ndof))
+        h_int = guess_tb(keys, ndof)
+        h_int[keys[len(keys) // 2]] += U0
+        h_ints.append(h_int)
+
+    for h0, h_int in zip(h0s, h_ints):
+        guess = guess_tb(frozenset(h_int), ndof)
+        _model = Model(h0, h_int, filling=filling)
+        mf_sol_groundstate = solver(
+            _model, mf_guess=guess, nk=nk, optimizer_kwargs={"M": 0}
+        )
+
+        alphas = np.random.uniform(0, 50, 100)
+        alphas = np.where(alphas == 1, 0, alphas)
+        assert np.all(
+            mf_rescaled(alphas, mf0=mf_sol_groundstate, h0=h0)
+            > mf_rescaled(np.array([1]), mf0=mf_sol_groundstate, h0=h0)
+        )