Interface refactoring

Find a design which is more intuitive and minimal.

codes/tb/transforms.py

 import numpy as np
-from .utils import quad_vecNDim
 from scipy.fftpack import ifftn
-from itertools import product
 import itertools as it
 
 
@@ -29,6 +27,28 @@ def tb2kfunc(h_0):
     return bloch_ham
 
 
+def ifftn2tb(ifftArray):
+    """
+    Converts an array from ifftn to a tight-binding model format.
+
+    Parameters
+    ----------
+    ifftArray : ndarray
+        An array obtained from ifftn.
+
+    Returns
+    -------
+    dict
+        A dictionary with real-space vectors as keys and complex np.arrays as values.
+    """
+
+    size = ifftArray.shape[:-2]
+
+    keys = [np.arange(-size[0] // 2 + 1, size[0] // 2) for i in range(len(size))]
+    keys = it.product(*keys)
+    return {tuple(k): ifftArray[tuple(k)] for k in keys}
+
+
 def kfunc2kham(nk, hk, dim, return_ks=False, hermitian=True):
     """
     Evaluates Hamiltonian on a k-point grid.
@@ -61,7 +81,7 @@ def kfunc2kham(nk, hk, dim, return_ks=False, hermitian=True):
     k_pts = np.tile(ks, dim).reshape(dim, nk)
 
     ham = []
-    for k in product(*k_pts):
+    for k in it.product(*k_pts):
         ham.append(hk(k))
     ham = np.array(ham)
     if hermitian:
@@ -81,7 +101,7 @@ def tb2kham(h_0, nK=200, ndim=1):
     return kham
 
 
-def kfunc2tbFFT(kfunc, nSamples, ndim=1):
+def kfunc2tb(kfunc, nSamples, ndim=1):
     """
     Applies FFT on a k-space function to obtain a real-space components.
 
@@ -95,8 +115,8 @@ def kfunc2tbFFT(kfunc, nSamples, ndim=1):
     Returns
     -------
 
-    ndarray
-        An array with real-space components of kfuncs
+    dict
+        A dictionary with real-space vectors as keys and complex np.arrays as values.
     """
 
     ks = np.linspace(
@@ -112,40 +132,8 @@ def kfunc2tbFFT(kfunc, nSamples, ndim=1):
         kfuncOnGrid = np.array([[kfunc((kx, ky)) for ky in ks] for kx in ks])
     if ndim > 2:
         raise NotImplementedError("n > 2 not implemented")
-    return ifftn(kfuncOnGrid, axes=np.arange(ndim))
-
+    ifftnArray = ifftn(kfuncOnGrid, axes=np.arange(ndim))
+    return ifftn2tb(ifftnArray)
 
 def kdens2tbFFT(kdens, ndim=1):
-
-    return ifftn(kdens, axes=np.arange(ndim))
-
-
-def kfunc2tbQuad(kfunc, ndim=1):
-    """
-    Inverse Fourier transforms a k-space function to a real-space function using a
-    ndim quadrature integration.
-
-    Parameters
-    ----------
-    kfunc : function
-        A function that takes a k-space vector and returns a np.array.
-    ndim : int
-        Dimension of the k-space
-
-    Returns
-    -------
-    function
-        A function that takes a real-space integer vector and returns a np.array.
-    """
-
-    def tbfunc(vector):
-        def integrand(k):
-            return (
-                kfunc(k)
-                * np.exp(1j * np.dot(k, np.array(vector, dtype=float)))
-                / (2 * np.pi)
-            )
-
-        return quad_vecNDim(integrand, -np.pi, np.pi, ndim=ndim)[0]
-
-    return tbfunc
+    return ifftn(kdens, axes=np.arange(ndim))
\ No newline at end of file