diff --git a/pymf/model.py b/pymf/model.py
index 6a61b2d4cda86bc770691c23b970c238f6ca3928..b41813a151b29f332d15c686e0446b5c755f5896 100644
--- a/pymf/model.py
+++ b/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(
diff --git a/pymf/observables.py b/pymf/observables.py
index bf770349e11798bbe4d86dae5f3f076cf86043f1..6267da2e8366f9e3f00ca66e2a7a775e5d9e60be 100644
--- a/pymf/observables.py
+++ b/pymf/observables.py
@@ -1,19 +1,20 @@
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(
diff --git a/pymf/params/param_transforms.py b/pymf/params/param_transforms.py
index 98b68e2485597b9479ab30755e76b37312e15968..5ab2a79a9285867fb20c3b29525b22b70d42e59a 100644
--- a/pymf/params/param_transforms.py
+++ b/pymf/params/param_transforms.py
@@ -1,18 +1,19 @@
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 :]
diff --git a/pymf/params/rparams.py b/pymf/params/rparams.py
index bb78ed086ef08366f3d6c91f90ea0f91d838793b..8e9c70e0c5075ff382afb8be11ee80aa33e4ad09 100644
--- a/pymf/params/rparams.py
+++ b/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)
diff --git a/pymf/solvers.py b/pymf/solvers.py
index 9099f39749878dfac9374f19291593d6262e9814..b0a51d75858cf60dcc61f82c719cc15275cbeed5 100644
--- a/pymf/solvers.py
+++ b/pymf/solvers.py
@@ -1,14 +1,15 @@
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)})
diff --git a/pymf/tb/tb.py b/pymf/tb/tb.py
index a059323b98b434dd41b3a8b413442fa21c95b88d..11da27b24ceb67b46f229f0db5a395ead5b2340f 100644
--- a/pymf/tb/tb.py
+++ b/pymf/tb/tb.py
@@ -1,25 +1,27 @@
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)
diff --git a/pymf/tb/transforms.py b/pymf/tb/transforms.py
index 4befee2240393754d58623f0095745e3de6d7efb..f886b88fc4c9b689f1d7a52d4bb76166abea95d7 100644
--- a/pymf/tb/transforms.py
+++ b/pymf/tb/transforms.py
@@ -1,24 +1,27 @@
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:
diff --git a/pymf/tb/utils.py b/pymf/tb/utils.py
index cc20753f704bf160e6b0eec752ae1ae50adebc46..a940e93cf5f42ab86bb710a176ffe5c1efbf0fe8 100644
--- a/pymf/tb/utils.py
+++ b/pymf/tb/utils.py
@@ -1,26 +1,27 @@
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)