Documentation

Construct the documentation for the release of the package.

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):
-    """Define the cost function for fixed point iteration.
+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 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 : 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 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._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)
+    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, mf_guess, nk=100, optimizer=scipy.optimize.anderson, optimizer_kwargs={}
-):
-    """Solve the mean-field self-consistent equation.
+    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 : Model
-        The model object.
-    mf_guess : numpy.array
-        The initial guess for the mean-field tight-binding model.
-    nk : int, optional
-        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 {}.
+    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
     -------
-    result : numpy.array
-        The mean-field tight-binding model.
+    :
+        Mean-field correction solution in the tight-binding dictionary format.
     """
-    shape = Model._size
+    shape = model._ndof
     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._ndof)})