From 08cbaaadbff8021d0a583c3b817b77c91b4fcefb Mon Sep 17 00:00:00 2001
From: Joerg Behrmann <behrmann@physik.fu-berlin.de>
Date: Tue, 24 May 2016 14:26:31 +0200
Subject: [PATCH] define ordering for `SiteFamily` and add a test for this
---
kwant/builder.py | 8 ++++++++
kwant/tests/test_builder.py | 17 +++++++++++++++++
2 files changed, 25 insertions(+)
diff --git a/kwant/builder.py b/kwant/builder.py
index 40047f27..b0749521 100644
--- a/kwant/builder.py
+++ b/kwant/builder.py
@@ -12,6 +12,7 @@ __all__ = ['Builder', 'Site', 'SiteFamily', 'SimpleSiteFamily', 'Symmetry',
import abc
import warnings
import operator
+from functools import total_ordering
from itertools import islice, chain
import tinyarray as ta
import numpy as np
@@ -82,6 +83,7 @@ class Site(tuple):
return self.family.pos(self.tag)
+@total_ordering
class SiteFamily(metaclass=abc.ABCMeta):
"""Abstract base class for site families.
@@ -154,6 +156,12 @@ class SiteFamily(metaclass=abc.ABCMeta):
except AttributeError:
return True
+ def __lt__(self, other):
+ try:
+ return self.canonical_repr < other.canonical_repr
+ except AttributeError:
+ return False
+
@abc.abstractmethod
def normalize_tag(self, tag):
"""Return a normalized version of the tag.
diff --git a/kwant/tests/test_builder.py b/kwant/tests/test_builder.py
index 71ef0e73..f270cf3d 100644
--- a/kwant/tests/test_builder.py
+++ b/kwant/tests/test_builder.py
@@ -111,6 +111,23 @@ def test_site_families():
assert_not_equal(fam, 'a')
+def test_site_families_sorting():
+ fam1 = builder.SimpleSiteFamily('fam1')
+ fam2 = builder.SimpleSiteFamily('fam2')
+
+ rng = Random(123)
+
+ tags = [(rng.randint(0,10), rng.randint(0,10)) for i in range(10)]
+ tags_sorted = [t for t in sorted(tags)]
+
+ sites = [fam1(*t) for t in tags] + [fam2(*t) for t in tags]
+ rng.shuffle(sites)
+ sites_sorted = ([fam1(*t) for t in tags_sorted] +
+ [fam2(*t) for t in tags_sorted])
+
+ assert_equal(list(sorted(sites)), sites_sorted)
+
+
class VerySimpleSymmetry(builder.Symmetry):
def __init__(self, period):
self.period = period
--
GitLab