Draft: Resolve "bring back tb_to_hk or allow a user specified k-grid in tb_to_kgrid"

meanfi/__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, kgrid_to_tb
+from .tb.transforms import tb_to_kgrid, kgrid_to_tb, tb_to_kfunc
 from .tb.utils import guess_tb, fermi_energy
 
 
@@ -29,6 +29,7 @@ __all__ = [
     "density_matrix",
     "meanfield",
     "tb_to_kgrid",
+    "tb_to_kfunc",
     "kgrid_to_tb",
     "__version__",
     "__version_tuple__",

meanfi/tb/transforms.py

 import itertools
 import numpy as np
+from typing import Callable
 from scipy.fftpack import ifftn
 
 from meanfi.tb.tb import _tb_type
@@ -76,3 +77,31 @@ def ifftn_to_tb(ifft_array: np.ndarray) -> _tb_type:
     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}
+
+
+def tb_to_kfunc(tb: _tb_type) -> Callable:
+    """
+    Fourier transforms a real-space tight-binding model to a k-space function.
+
+    Parameters
+    ----------
+    tb :
+        Tight-binding dictionary.
+
+    Returns
+    -------
+    :
+        A function that takes a k-space vector and returns a complex np.array.
+
+    Notes
+    -----
+    Function doesn't work for zero dimensions.
+    """
+
+    def kfunc(k):
+        ham = 0
+        for vector in tb.keys():
+            ham += tb[vector] * np.exp(-1j * np.dot(k, np.array(vector, dtype=float)))
+        return ham
+
+    return kfunc

meanfi/tests/test_tb.py

@@ -6,7 +6,8 @@ import pytest
 from scipy.fftpack import ifftn
 
 from meanfi.tb.tb import compare_dicts
-from meanfi.tb.transforms import ifftn_to_tb, tb_to_kgrid
+from meanfi.tb.utils import guess_tb, generate_tb_keys
+from meanfi.tb.transforms import ifftn_to_tb, tb_to_kgrid, tb_to_kfunc
 
 repeat_number = 10
 
@@ -27,3 +28,29 @@ def test_fourier(seed):
     kham = tb_to_kgrid(h_0, nk=nk)
     tb_new = ifftn_to_tb(ifftn(kham, axes=np.arange(ndim)))
     compare_dicts(h_0, tb_new)
+
+
+@pytest.mark.parametrize("seed", range(repeat_number))
+def test_kfunc(seed):
+    np.random.seed(seed)
+    cutoff = np.random.randint(1, 4)
+    dim = np.random.randint(1, 3)
+    ndof = np.random.randint(2, 10)
+    nk = np.random.randint(3, 10)
+    random_hopping_vecs = generate_tb_keys(cutoff, dim)
+    random_tb = guess_tb(random_hopping_vecs, ndof, scale=1)
+
+    kfunc = tb_to_kfunc(random_tb)
+    kham = tb_to_kgrid(random_tb, nk=nk)
+
+    # evaluate kfunc on same grid as kham
+    ks = np.linspace(-np.pi, np.pi, nk, endpoint=False)
+    ks = np.concatenate((ks[nk // 2 :], ks[: nk // 2]), axis=0)
+    k_pts = np.tile(ks, dim).reshape(dim, nk)
+
+    ham_kfunc = []
+    for k in it.product(*k_pts):
+        ham_kfunc.append(kfunc(k))
+    ham_kfunc = np.array(ham_kfunc).reshape(*kham.shape)
+
+    assert np.allclose(kham, ham_kfunc)