add tests and fixes for an edge case in wraparound

When the fundamental domains of the constituent symmetries
of a wrapped around system do not coincide, it was previously
possible for wrapping around to fail when 'keep' was provided,
even if wrapping around succeeded with 'keep=None'.

kwant/tests/test_wraparound.py

@@ -250,3 +250,65 @@ def test_plot_2d_bands():
         plot_2d_bands(syst, extend_bbox=1.2, k_x=11, k_y=11, file=out)
         # test mask Brillouin zone
         plot_2d_bands(syst, mask_brillouin_zone=True, k_x=11, k_y=11, file=out)
+
+
+def test_fd_mismatch():
+    # The fundamental domains of the two 1D symmetries that make up T are
+    # incompatible. This produced a bug where certain systems could be wrapped
+    # around in all directions, but could not be wrapped around when 'keep' is
+    # provided.
+    sqrt3 = np.sqrt(3)
+    lat = kwant.lattice.general([(sqrt3, 0), (-sqrt3/2, 1.5)])
+    T = kwant.TranslationalSymmetry((sqrt3, 0), (0, 3))
+
+    syst1 = kwant.Builder(T)
+    syst1[lat(1, 1)] = syst1[lat(0, 1)] = 1
+    syst1[lat(1, 1), lat(0, 1)] = 1
+
+    syst2 = kwant.Builder(T)
+    syst2[lat(0, 0)] = syst2[lat(0, -1)] = 1
+    syst2[lat(0, 0), lat(0, -1)] = 1
+
+    # combine the previous two
+    syst3 = kwant.Builder(T)
+    syst3 += syst1
+    syst3 += syst2
+
+    for syst in (syst1, syst2, syst3):
+        wraparound(syst)
+        wraparound(syst, keep=0)
+        wraparound(syst, keep=1)
+
+    ## Test that spectrum of non-trivial system (including above cases)
+    ## is the same, regardless of the way in which it is wrapped around
+    lat = kwant.lattice.general([(sqrt3, 0), (-sqrt3/2, 1.5)],
+                                [(sqrt3 / 2, 0.5), (0, 1)])
+    a, b = lat.sublattices
+    T = kwant.TranslationalSymmetry((3 * sqrt3, 0), (0, 3))
+    syst = kwant.Builder(T)
+    syst[(l(i, j) for i in range(-1, 2) for j in range(-1, 2)
+         for l in (a, b))] = 0
+    syst[kwant.HoppingKind((0, 0), b, a)] = -1j
+    syst[kwant.HoppingKind((1, 0), b, a)] = 2j
+    syst[kwant.HoppingKind((0, 1), a, b)] = -2j
+
+    def spectrum(syst, keep):
+        syst = wraparound(syst, keep=keep).finalized()
+        if keep is None:
+            def _(*args):
+                return np.linalg.eigvalsh(syst.hamiltonian_submatrix(args=args))
+        else:
+            def _(*args):
+                args = list(args)
+                assert len(args) == 2
+                kext = args.pop(keep)
+                kint = args[0]
+                B = kwant.physics.Bands(syst, args=(kint,))
+                return B(kext)
+        return _
+
+    spectra = [spectrum(syst, keep) for keep in (None, 0, 1)]
+    # Check that spectra are the same at k=0. Checking finite k
+    # is more tricky, as we would also need to transform the k vectors
+    E_k = np.array([spec(0, 0) for spec in spectra]).transpose()
+    assert all(np.allclose(E, E[0]) for E in E_k)

kwant/wraparound.py

@@ -269,14 +269,22 @@ def wraparound(builder, keep=None, *, coordinate_names=('x', 'y', 'z')):
     mnp = -len(sym.periods)      # Used by the bound functions above.
 
     # Store lists of values, so that multiple values can be assigned to the
-    # same site or hopping.
+    # same site or hopping. We map to the FD of 'sym', as the hopping-processing
+    # code assumes the sites are in the FD.
     for site, val in builder.site_value_pairs():
-        sites[site] = [bind_site(val) if callable(val) else val]
+        sites[sym.to_fd(site)] = [bind_site(val) if callable(val) else val]
 
     for hop, val in builder.hopping_value_pairs():
-        a, b = hop
+        # Map hopping to FD of 'sym', as the code afterwards assumes that 'a'
+        # is in this domain.
+        a, b = sym.to_fd(*hop)
         b_dom = sym.which(b)
         b_wa = sym.act(-b_dom, b)
+        # Now map 'b_wa' to another domain of 'sym', so that the hopping
+        # is compatible with 'ret.symmetry'. This is necessary when the
+        # fundamental domains of the symmetries do not coincide.
+        w = ret.symmetry.which(b_wa)
+        b_wa = ret.symmetry.act(w, sym.to_fd(ret.symmetry.act(-w, b_wa)))
 
         if a == b_wa:
             # The hopping gets wrapped-around into an onsite Hamiltonian.