add supercell test with callables

meanfi/tests/test_kwant.py

@@ -2,6 +2,8 @@ import numpy as np
 import pytest
 import kwant
 
+import itertools as it
+
 from meanfi.kwant_helper.utils import builder_to_tb, tb_to_builder
 from meanfi.tb.utils import generate_tb_keys, guess_tb
 from meanfi.tb.tb import compare_dicts
@@ -40,5 +42,82 @@ def test_kwant_conversion(seed):
         random_tb, list(dummy_tb.sites()), dummy_tb.symmetry.periods
     )
     # convert builder back to tb and compare
-    random_builder_tb = builder_to_tb(random_builder)
+    random_builder_tb = builder_to_tb(random_builder, return_data=True)
     compare_dicts(random_tb, random_builder_tb)
+
+
+@pytest.mark.parametrize("seed", range(repeat_number))
+def test_kwant_supercell(seed):
+    np.random.seed(seed)
+    ndim = np.random.randint(1, 3)
+    cutoff = np.random.randint(1, 3)
+    sites_in_cell = np.random.randint(1, 4)
+    ndof_per_site = [np.random.randint(1, 4) for site in range(sites_in_cell)]
+    keyList = generate_tb_keys(cutoff, ndim)
+    n_cells = np.random.randint(1, 4)
+
+    vecs = np.random.rand(ndim, ndim)
+
+    # set a dummy lattice to read sites from
+    lattice = kwant.lattice.general(
+        vecs,
+        basis=np.random.rand(sites_in_cell, ndim) @ vecs,
+        norbs=ndof_per_site,
+    )
+
+    def random_matrix_kwant_digest(n, m, k):
+        matrix = np.zeros((n, m))
+        for i in zip(it.product(range(n), range(m))):
+            matrix[*i[0]] = kwant.digest.uniform(str(n * m * np.prod(i[0]) + k))
+        return matrix
+
+    def onsite(site, alpha, beta):
+        n = site.family.norbs
+        amplitude = alpha * random_matrix_kwant_digest(n, n, 0)
+        phase = 1j * 2 * np.pi * beta * random_matrix_kwant_digest(n, n, 1)
+        onsite_matrix = amplitude * phase
+        onsite_matrix += onsite_matrix.conj().T
+        return onsite_matrix
+
+    def hopping(site1, site2, gamma, delta):
+        n1 = site1.family.norbs
+        n2 = site2.family.norbs
+        amplitude = gamma * random_matrix_kwant_digest(n1, n2, 0)
+        phase = 1j * 2 * np.pi * delta * random_matrix_kwant_digest(n1, n2, 1)
+        hopping_matrix = amplitude * phase
+        return hopping_matrix
+
+    random_builder = kwant.Builder(
+        kwant.TranslationalSymmetry(*n_cells * lattice.prim_vecs)
+    )
+    for i, sublattice in enumerate(lattice.sublattices):
+        random_builder[lattice.shape(lambda pos: True, tuple(ndim * [0]))] = onsite
+    random_builder[lattice.neighbors()] = hopping
+
+    params_num = np.random.rand(4)
+    params = dict(
+        alpha=params_num[0],
+        beta=params_num[1],
+        gamma=params_num[2],
+        delta=params_num[3],
+    )
+
+    random_tb, data = builder_to_tb(random_builder, params=params, return_data=True)
+    random_builder_test = tb_to_builder(random_tb, data["sites"], data["periods"])
+    for site_pair in zip(it.product(data["sites"], data["sites"])):
+        site1, site2 = site_pair[0]
+        if site1 == site2:
+            assert np.isclose(
+                random_builder[site1](site=site1, alpha=params["alpha"], beta=params["beta"]),
+                random_builder_test[site1],
+            ).all()
+        else:
+            try:
+                assert np.isclose(
+                    random_builder[site1, site2](site1, site2, gamma=params["gamma"], delta=params["delta"]),
+                    random_builder_test[site1, site2],
+                ).all()
+            except KeyError:
+                continue
+            except:
+                raise
\ No newline at end of file