add automatic calculation of n-th nearest neighbors

doc/source/images/ab_ring.py.diff

@@ -10,7 +10,7 @@
  def make_system(a=1, t=1.0, W=10, r1=10, r2=20):
 @@ -37,12 +38,13 @@
      sys[lat.shape(ring, (0, r1 + 1))] = 4 * t
-     sys[lat.nearest] = -t
+     sys[lat.neighbors()] = -t
  
 -    # In order to introduce a flux through the ring, we introduce a phase on
 -    # the hoppings on the line cut through one of the arms.  Since we want to
@@ -40,14 +40,14 @@
 +        rsq = x**2 + y**2
 +        return ( r1**2 < rsq < r2**2)
 +    sys[lat.shape(ring, (0, 11))] = 4 * t
-+    sys[lat.nearest] = -t
++    sys[lat.neighbors()] = -t
 +    sym_lead0 = kwant.TranslationalSymmetry((-a, 0))
 +    lead0 = kwant.Builder(sym_lead0)
 +    def lead_shape(pos):
 +        (x, y) = pos
 +        return (-1 < x < 1) and ( 0.5 * W < y < 1.5 * W )
 +    lead0[lat.shape(lead_shape, (0, W))] = 4 * t
-+    lead0[lat.nearest] = -t
++    lead0[lat.neighbors()] = -t
 +    lead1 = lead0.reversed()
 +    sys.attach_lead(lead0)
 +    sys.attach_lead(lead1)
@@ -62,14 +62,14 @@
 +        rsq = x**2 + y**2
 +        return ( r1**2 < rsq < r2**2)
 +    sys[lat.shape(ring, (0, 11))] = 4 * t
-+    sys[lat.nearest] = -t
++    sys[lat.neighbors()] = -t
 +    sym_lead0 = kwant.TranslationalSymmetry((-a, 0))
 +    lead0 = kwant.Builder(sym_lead0)
 +    def lead_shape(pos):
 +        (x, y) = pos
 +        return (-1 < x < 1) and ( -W/2 < y < W/2  )
 +    lead0[lat.shape(lead_shape, (0, 0))] = 4 * t
-+    lead0[lat.nearest] = -t
++    lead0[lat.neighbors()] = -t
 +    lead1 = lead0.reversed()
 +    sys.attach_lead(lead0)
 +    sys.attach_lead(lead1, lat(0, 0))

doc/source/images/graphene.py.diff

@@ -8,7 +8,7 @@
  
  
  # Define the graphene lattice
-@@ -102,22 +103,40 @@
+@@ -100,22 +101,40 @@
          smatrix = kwant.solve(sys, energy)
          data.append(smatrix.transmission(0, 1))
  
@@ -57,7 +57,7 @@
  
  
  def main():
-@@ -130,17 +149,22 @@
+@@ -128,17 +147,22 @@
          return 0 if site.family == a else 1
  
      # Plot the closed system without leads.

doc/source/tutorial/ab_ring.py

@@ -38,7 +38,7 @@ def make_system(a=1, t=1.0, W=10, r1=10, r2=20):
     # and add the corresponding lattice points using the `shape`-function
 #HIDDEN_BEGIN_lcak
     sys[lat.shape(ring, (0, r1 + 1))] = 4 * t
-    sys[lat.nearest] = -t
+    sys[lat.neighbors()] = -t
 #HIDDEN_END_lcak
 
     # In order to introduce a flux through the ring, we introduce a phase on
@@ -76,7 +76,7 @@ def make_system(a=1, t=1.0, W=10, r1=10, r2=20):
         return (-W / 2 < y < W / 2)
 
     lead[lat.shape(lead_shape, (0, 0))] = 4 * t
-    lead[lat.nearest] = -t
+    lead[lat.neighbors()] = -t
 #HIDDEN_END_qwgr
 
     #### Attach the leads and return the system. ####

doc/source/tutorial/quantum_well.py

@@ -35,13 +35,13 @@ def make_system(a=1, t=1.0, W=10, L=30, L_well=10):
         return 4 * t + potential(site, pot)
 
     sys[(lat(x, y) for x in range(L) for y in range(W))] = onsite
-    sys[lat.nearest] = -t
+    sys[lat.neighbors()] = -t
 #HIDDEN_END_coid
 
     #### Define and attach the leads. ####
     lead = kwant.Builder(kwant.TranslationalSymmetry((-a, 0)))
     lead[(lat(0, j) for j in xrange(W))] = 4 * t
-    lead[lat.nearest] = -t
+    lead[lat.neighbors()] = -t
     sys.attach_lead(lead)
     sys.attach_lead(lead.reversed())
 

doc/source/tutorial/quantum_wire_revisited.py

@@ -29,7 +29,7 @@ def make_system(a=1, t=1.0, W=10, L=30):
     sys[(lat(x, y) for x in range(L) for y in range(W))] = 4 * t
 #HIDDEN_END_vvjt
 #HIDDEN_BEGIN_nooi
-    sys[lat.nearest] = -t
+    sys[lat.neighbors()] = -t
 #HIDDEN_END_nooi
 
     #### Define and attach the leads. ####
@@ -37,7 +37,7 @@ def make_system(a=1, t=1.0, W=10, L=30):
 #HIDDEN_BEGIN_iepx
     lead = kwant.Builder(kwant.TranslationalSymmetry((-a, 0)))
     lead[(lat(0, j) for j in xrange(W))] = 4 * t
-    lead[lat.nearest] = -t
+    lead[lat.neighbors()] = -t
 #HIDDEN_END_iepx
 
     # Attach the left lead and its reversed copy.

doc/source/tutorial/superconductor_transport.py

@@ -36,8 +36,8 @@ def make_system(a=1, W=10, L=10, barrier=1.5, barrierpos=(3, 4),
          for y in range(W))] = mu - 4 * t - barrier
 
     # hoppings for both electrons and holes
-    sys[lat_e.nearest] = -t
-    sys[lat_h.nearest] = t
+    sys[lat_e.neighbors()] = -t
+    sys[lat_h.neighbors()] = t
 
     # Superconducting order parameter enters as hopping between
     # electrons and holes
@@ -53,12 +53,12 @@ def make_system(a=1, W=10, L=10, barrier=1.5, barrierpos=(3, 4),
     # left electron lead
     lead0 = kwant.Builder(sym_left)
     lead0[(lat_e(0, j) for j in xrange(W))] = 4 * t - mu
-    lead0[lat_e.nearest] = -t
+    lead0[lat_e.neighbors()] = -t
 
     # left hole lead
     lead1 = kwant.Builder(sym_left)
     lead1[(lat_h(0, j) for j in xrange(W))] = mu - 4 * t
-    lead1[lat_h.nearest] = t
+    lead1[lat_h.neighbors()] = t
 #HIDDEN_END_ttth
 
     # Then the lead to the right

doc/source/tutorial/tutorial1.rst

@@ -308,16 +308,16 @@ feature of kwant:
 In regular lattices, hoppings form large groups such that hoppings within a
 group can be transformed into one another by lattice translations. In order to
 allow to easily manipulate such hoppings, an object
-`~kwant.builder.HoppingKind` is provided. When given a `~kwant.builder.Builder` as
-an argument, `~kwant.builder.HoppingKind` yields all the hoppings of a
+`~kwant.builder.HoppingKind` is provided. When given a `~kwant.builder.Builder`
+as an argument, `~kwant.builder.HoppingKind` yields all the hoppings of a
 certain kind that can be added to this builder without adding new sites. When
 `~kwant.builder.HoppingKind` is given to `~kwant.builder.Builder` as a key, it
 means that something is done to all the possible hoppings of this kind. A list
 of `~kwant.builder.HoppingKind` objects corresponding to nearest neighbors in
-pre-defined lattices in kwant (that is `~kwant.lattice.chain`,
-`~kwant.lattice.square`, and `~kwant.lattice.honeycomb`) is stored in
-``lat.nearest``. ``sys[lat.nearest] = -t`` then sets all of those hopping
-matrix elements at once. More detailed example of using
+lattices in kwant is obtained using ``lat.neighbors()``. ``sys[lat.neighbors()]
+= -t`` then sets all of those hopping matrix elements at once. In order to set
+values for all the nth-nearest neighbors at once, one can similarly use
+``sys[lat.neighbors(n)] = -t``. More detailed example of using
 `~kwant.builder.HoppingKind` directly will be provided in
 :ref:`tutorial_spinorbit`.
 

doc/source/tutorial/tutorial2.rst

@@ -59,9 +59,9 @@ we can simply write:
 Note that the Zeeman energy adds to the onsite term, whereas the Rashba
 spin-orbit term adds to the hoppings (due to the derivative operator).
 Furthermore, the hoppings in x and y-direction have a different matrix
-structure. We now cannot use ``lat.nearest`` to add all the hoppings at once,
-since we now have to distinguish x and y-direction. Because of that, we have to
-explicitly specify the hoppings in the form expected by
+structure. We now cannot use ``lat.neighbors()`` to add all the hoppings at
+once, since we now have to distinguish x and y-direction. Because of that, we
+have to explicitly specify the hoppings in the form expected by
 `~kwant.builder.HoppingKind`:
 
 - A tuple with relative lattice indices.  For example, `(1, 0)` means
@@ -248,7 +248,7 @@ provided by the lattice:
 
 Here, ``lat.shape`` takes as a second parameter a (real-space) point that is
 inside the desired shape. The hoppings can still be added using
-``lat.nearest`` as before.
+``lat.neighbors()`` as before.
 
 Up to now, the system contains constant hoppings and onsite energies,
 and we still need to include the phase shift due to the magnetic flux.

doc/source/whatsnew/0.3.rst

@@ -14,7 +14,7 @@ obsolete.  Where previously one would have had ::
 
 now it suffices to write ::
 
-    sys[lat.nearest] = t
+    sys[lat.neighbors()] = t
 
 This is possible because `~kwant.builder.Builder` now accepts *functions* as
 keys in addition to `~kwant.builder.Site` objects and tuples of them
@@ -125,9 +125,13 @@ Lattice and shape improvements
 ------------------------------
 `~kwant.lattice.Monoatomic.closest` now returns an exact, and not approximately
 closest point. A new method `~kwant.lattice.Monoatomic.n_closest` was added,
-which returns n closest lattice points. Likewise
-`~kwant.lattice.Polyatomic.shape` has acquired an improved flood-fill
+which returns n closest lattice points.
+
+Likewise `~kwant.lattice.Polyatomic.shape` has acquired an improved flood-fill
 algorithm, making it work better on narrow ribbon (which were sometimes buggy
 before with non-square lattices). Additionally, it was made symmetry-aware, so
 if a shape is used for a lead, no conditions with regard to coordnate parallel
 to the lead period are required.
+
+Finally, lattices now have a method `~kwant.lattice.Polyatomic.neighbors`,
+which calculates all the n-th shortest possible hoppings on this lattice.

kwant/lattice.py

@@ -11,6 +11,7 @@ from __future__ import division
 __all__ = ['TranslationalSymmetry', 'general', 'Polyatomic', 'Monatomic']
 
 from math import sqrt
+from itertools import product
 import numpy as np
 import tinyarray as ta
 from . import builder
@@ -104,6 +105,88 @@ class Polyatomic(object):
         """
         return Shape(self, function, start)
 
+    def neighbors(self, n=1, eps=1e-8):
+        """
+        Return n-th nearest neighbor hoppings.
+
+        Parameters
+        ----------
+        n : integer
+            Order of the hoppings to return.
+        eps : float
+            A cutoff for when to consider lengths to be approximately equal.
+
+        Returns
+        -------
+        hoppings : list of kwant.builder.HopplingKind objects
+            A list n-th nearest neighbor hoppings.
+
+        Notes
+        -----
+        The hoppings are ordered lexicographically according to sublattice from
+        which they originate, sublattice on which they end, and their lattice
+        coordinates. Out of the two equivalent hoppings (a hopping and its
+        reverse) only the lexicographically larger one is returned.
+        """
+        # This algorithm is not designed to be fast and can be improved,
+        # however there is no real need.
+        sls = self.sublattices
+        nvec = len(self.prim_vecs)
+        sublat_pairs = [(i, j) for (i, j) in product(sls, sls)
+                        if sls.index(j) >= sls.index(i)]
+        def first_nonnegative(tag):
+            for i in tag:
+                if i < 0:
+                    return False
+                elif i > 0:
+                    return True
+                else:
+                    continue
+            return True
+
+        # Find the correct number of neighbors to calculate on each lattice.
+        cutoff = n + 2
+        while True:
+            max_dist = []
+            sites = []
+            for i, j in sublat_pairs:
+                origin = j(*ta.zeros(nvec)).pos
+                tags = i.n_closest(origin, n=cutoff**nvec)
+
+                ij_dist = [np.linalg.norm(i(*tag).pos - origin)
+                              for tag in tags]
+                sites.append((tags, (i, j), ij_dist))
+            max_dist = [i[2][-1] for i in sites]
+            distances = np.r_[tuple((i[2] for i in sites))]
+            distances = np.sort(distances)
+            group_boundaries = np.argwhere(np.diff(distances) > eps)
+            if len(group_boundaries) < n:
+                cutoff += 1
+                continue
+            n_dist = distances[group_boundaries[n]]
+            if np.all(max_dist > n_dist):
+                break
+            cutoff += 1
+
+        # We now have all the required sites, we need to find n-th.
+        result = []
+        for group in sites:
+            tags, distance = group[0], group[2]
+            i, j = group[1]
+            tags = np.array([tag for tag, dist in zip(tags, distance)
+                             if abs(dist - n_dist) < eps])
+            if len(tags):
+                # Sort the tags.
+                tags = tags[np.lexsort(tags.T[::-1])][::-1]
+                # Throw away equivalent hoppings if
+                # two sublattices are the same.
+                if i == j and len(tags) > 1:
+                    tags = tags[: len(tags) // 2]
+                for tag in tags:
+                    result.append(builder.HoppingKind(tag, j, i))
+        return result
+
+
     def vec(self, int_vec):
         """
         Return the coordinates of a Bravais lattice vector in real space.
@@ -489,15 +572,12 @@ class Shape(object):
 def chain(a=1, name=''):
     """Create a one-dimensional lattice."""
     lat = Monatomic(((a,),), name=name)
-    lat.nearest = [builder.HoppingKind((1,), lat, lat)]
     return lat
 
 
 def square(a=1, name=''):
     """Create a square lattice."""
     lat = Monatomic(((a, 0), (0, a)), name=name)
-    lat.nearest = [builder.HoppingKind((1, 0), lat, lat),
-                   builder.HoppingKind((0, 1), lat, lat)]
     return lat
 
 
@@ -506,7 +586,4 @@ def honeycomb(a=1, name=''):
     lat = Polyatomic(((a, 0), (0.5 * a, 0.5 * a * sqrt(3))),
                             ((0, 0), (0, a / sqrt(3))), name=name)
     lat.a, lat.b = lat.sublattices
-    lat.nearest = [builder.HoppingKind((0, 0), lat.a, lat.b),
-                   builder.HoppingKind((0, 1), lat.a, lat.b),
-                   builder.HoppingKind((-1, 1), lat.a, lat.b)]
     return lat

kwant/solvers/tests/_test_sparse.py

@@ -382,7 +382,7 @@ def test_wavefunc_ldos_consistency(wave_function, ldos):
             h = np.random.rand(n, n) + 1j * np.random.rand(n, n)
             h += h.conjugate().transpose()
             b[site] = h
-        for hopping_kind in square.nearest:
+        for hopping_kind in square.neighbors():
             for hop in hopping_kind(b):
                 b[hop] = 10 * np.random.rand(n, n) + 1j * np.random.rand(n, n)
     sys.attach_lead(left_lead)

kwant/tests/test_lattice.py

@@ -43,6 +43,15 @@ def test_general():
     assert_raises(ValueError, lat, 0, 1)
 
 
+def test_neighbors():
+    lat = lattice.honeycomb()
+    num_nth_nearest = [len(lat.neighbors(n)) for n in range(5)]
+    assert num_nth_nearest == [2, 3, 6, 3, 6]
+    lat = lattice.square()
+    num_nth_nearest = [len(lat.neighbors(n)) for n in range(5)]
+    assert num_nth_nearest == [1, 2, 2, 2, 4]
+
+
 def test_shape():
     def in_circle(pos):
         return pos[0] ** 2 + pos[1] ** 2 < 3

kwant/tests/test_plotter.py

@@ -28,7 +28,7 @@ def sys_2d(W=3, r1=3, r2=8):
         return r1 ** 2 < rsq < r2 ** 2
 
     sys[lat.shape(ring, (0, r1 + 1))] = 4 * t
-    sys[lat.nearest] = -t
+    sys[lat.neighbors()] = -t
     sym_lead0 = kwant.TranslationalSymmetry(lat.vec((-1, 0)))
     lead0 = kwant.Builder(sym_lead0)
     lead2 = kwant.Builder(sym_lead0)
@@ -38,7 +38,7 @@ def sys_2d(W=3, r1=3, r2=8):
     lead0[lat.shape(lead_shape, (0, 0))] = 4 * t
     lead2[lat.shape(lead_shape, (0, 0))] = 4 * t
     sys.attach_lead(lead2)
-    lead0[lat.nearest] = - t
+    lead0[lat.neighbors()] = - t
     lead1 = lead0.reversed()
     sys.attach_lead(lead0)
     sys.attach_lead(lead1)
@@ -47,8 +47,6 @@ def sys_2d(W=3, r1=3, r2=8):
 
 def sys_3d(W=3, r1=2, r2=4, a=1, t=1.0):
     lat = kwant.lattice.general(((a, 0, 0), (0, a, 0), (0, 0, a)))
-    lat.nearest = [kwant.builder.HoppingKind(*kind) for kind in
-                   [((1, 0, 0), lat), ((0, 1, 0), lat), ((0, 0, 1), lat)]]
     sys = kwant.Builder()
 
     def ring(pos):
@@ -56,14 +54,14 @@ def sys_3d(W=3, r1=2, r2=4, a=1, t=1.0):
         rsq = x ** 2 + y ** 2
         return (r1 ** 2 < rsq < r2 ** 2) and abs(z) < 2
     sys[lat.shape(ring, (0, -r2 + 1, 0))] = 4 * t
-    sys[lat.nearest] = - t
+    sys[lat.neighbors()] = - t
     sym_lead0 = kwant.TranslationalSymmetry(lat.vec((-1, 0, 0)))
     lead0 = kwant.Builder(sym_lead0)
 
     lead_shape = lambda pos: (-W / 2 < pos[1] < W / 2) and abs(pos[2]) < 2
 
     lead0[lat.shape(lead_shape, (0, 0, 0))] = 4 * t
-    lead0[lat.nearest] = - t
+    lead0[lat.neighbors()] = - t
     lead1 = lead0.reversed()
     sys.attach_lead(lead0)
     sys.attach_lead(lead1)