add kgrid_to_tb fun; bump down the default nk values

pymf/__init__.py

@@ -14,7 +14,7 @@ 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
+from .tb.transforms import tb_to_kgrid, kgrid_to_tb
 from .tb.utils import generate_guess, calculate_fermi_energy
 
 
@@ -29,6 +29,7 @@ __all__ = [
     "construct_density_matrix",
     "meanfield",
     "tb_to_kgrid",
+    "kgrid_to_tb",
     "__version__",
     "__version_tuple__",
 ]

pymf/mf.py

 import numpy as np
-from scipy.fftpack import ifftn
 from typing import Tuple
 
 from pymf.tb.tb import add_tb, _tb_type
-from pymf.tb.transforms import ifftn_to_tb, tb_to_kgrid
+from pymf.tb.transforms import tb_to_kgrid, kgrid_to_tb
 
 
 def construct_density_matrix_kgrid(
@@ -61,7 +60,7 @@ def construct_density_matrix(
         kham = tb_to_kgrid(h, nk=nk)
         density_matrix_krid, fermi = construct_density_matrix_kgrid(kham, filling)
         return (
-            ifftn_to_tb(ifftn(density_matrix_krid, axes=np.arange(ndim))),
+            kgrid_to_tb(density_matrix_krid),
             fermi,
         )
     else:

pymf/model.py

@@ -76,7 +76,7 @@ class Model:
         _check_hermiticity(h_0)
         _check_hermiticity(h_int)
 
-    def mfield(self, mf: _tb_type, nk: int = 200) -> _tb_type:
+    def mfield(self, mf: _tb_type, nk: int = 20) -> _tb_type:
         """Computes a new mean-field correction from a given one.
 
         Parameters

pymf/solvers.py

@@ -9,7 +9,7 @@ from pymf.model import Model
 from pymf.tb.utils import calculate_fermi_energy
 
 
-def cost(mf_param: np.ndarray, model: Model, nk: int = 100) -> np.ndarray:
+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.
@@ -40,7 +40,7 @@ def cost(mf_param: np.ndarray, model: Model, nk: int = 100) -> np.ndarray:
 def solver(
     model: Model,
     mf_guess: np.ndarray,
-    nk: int = 100,
+    nk: int = 20,
     optimizer: Optional[Callable] = scipy.optimize.anderson,
     optimizer_kwargs: Optional[dict[str, str]] = {"M": 0},
 ) -> _tb_type:

pymf/tb/transforms.py

 import itertools
 import numpy as np
+from scipy.fftpack import ifftn
 
 from pymf.tb.tb import _tb_type
 
@@ -37,6 +38,25 @@ def tb_to_kgrid(tb: _tb_type, nk: int) -> np.ndarray:
     return np.sum(tb_array * k_dependency, axis=0)
 
 
+def kgrid_to_tb(kgrid_array: np.ndarray) -> _tb_type:
+    """
+    Convert a k-space grid array to a tight-binding dictionary.
+
+    Parameters
+    ----------
+    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
+    -------
+    :
+        Tight-binding dictionary.
+    """
+    ndim = len(kgrid_array.shape) - 2
+    return ifftn_to_tb(ifftn(kgrid_array, axes=np.arange(ndim)))
+
+
 def ifftn_to_tb(ifft_array: np.ndarray) -> _tb_type:
     """
     Convert the result of `scipy.fft.ifftn` to a tight-binding dictionary.