optimize site creation

kwant/builder.py

@@ -54,7 +54,9 @@ class Site(tuple):
     group = property(operator.itemgetter(0))
     tag = property(operator.itemgetter(1))
 
-    def __new__(cls, group, tag):
+    def __new__(cls, group, tag, _i_know_what_i_do=False):
+        if _i_know_what_i_do:
+            return tuple.__new__(cls, (group, tag))
         try:
             tag = group.normalize_tag(tag)
         except (TypeError, ValueError):
@@ -885,7 +887,7 @@ class Builder(object):
         for site0 in self.H:
             if site0.group is not group_a:
                 continue
-            site1 = site0.shifted(d, group_b)
+            site1 = Site(group_b, site0.tag + d, True)
             if symtofd(site1) in H: # if site1 in self
                 yield site0, site1
 

kwant/lattice.py

@@ -96,6 +96,8 @@ class PolyatomicLattice(object):
         sites : sequence of `Site` objects
             all the sites that belong to the lattice and fit inside the shape.
         """
+        Site = builder.Site
+
         dim = len(start)
         if dim != self.prim_vecs.shape[1]:
             raise ValueError('Dimensionality of start position does not match'
@@ -125,7 +127,7 @@ class PolyatomicLattice(object):
                 for sl in sls:
                     if not function(vec + sl.offset):
                         continue
-                    yield sl(*tag)
+                    yield Site(sl, tag, True)
                     any_hits = True
                 if not any_hits:
                     continue
@@ -313,7 +315,7 @@ class TranslationalSymmetry(builder.Symmetry):
         try:
             det_x_inv_m_part, m_part, det_m = self.site_group_data[a.group]
         except KeyError:
-            self.add_site_group(gr)
+            self.add_site_group(a.group)
             return self.act(element, a, b)
         try:
             delta = ta.dot(m_part, element)
@@ -321,9 +323,10 @@ class TranslationalSymmetry(builder.Symmetry):
             msg = 'Expecting a {0}-tuple group element, but got `{1}` instead.'
             raise ValueError(msg.format(self.num_directions, element))
         if b is None:
-            return a.shifted(delta)
+            return builder.Site(a.group, a.tag + delta, True)
         else:
-            return a.shifted(delta), b.shifted(delta)
+            return builder.Site(a.group, a.tag + delta, True), \
+                builder.Site(b.group, b.tag + delta, True)
 
     def to_fd(self, a, b=None):
         return self.act(-self.which(a), a, b)