From 18c7fd24e3ca39921f739cfda627a645d5fc6f44 Mon Sep 17 00:00:00 2001
From: Anton Akhmerov <anton.akhmerov@gmail.com>
Date: Wed, 7 Dec 2016 22:35:55 +0100
Subject: [PATCH] allow site pickling
---
kwant/builder.py | 5 ++++-
kwant/tests/test_builder.py | 8 +++++++-
2 files changed, 11 insertions(+), 2 deletions(-)
diff --git a/kwant/builder.py b/kwant/builder.py
index f3b08e6a..6859cdf6 100644
--- a/kwant/builder.py
+++ b/kwant/builder.py
@@ -1,4 +1,4 @@
-# Copyright 2011-2015 Kwant authors.
+# Copyright 2011-2016 Kwant authors.
#
# This file is part of Kwant. It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
@@ -77,6 +77,9 @@ class Site(tuple):
sf = self.family
return '<Site {0} of {1}>'.format(self.tag, sf.name if sf.name else sf)
+ def __getnewargs__(self):
+ return (self.family, self.tag, True)
+
@property
def pos(self):
"""Real space position of the site.
diff --git a/kwant/tests/test_builder.py b/kwant/tests/test_builder.py
index 91359908..65dcf150 100644
--- a/kwant/tests/test_builder.py
+++ b/kwant/tests/test_builder.py
@@ -1,4 +1,4 @@
-# Copyright 2011-2015 Kwant authors.
+# Copyright 2011-2016 Kwant authors.
#
# This file is part of Kwant. It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
@@ -7,6 +7,7 @@
# http://kwant-project.org/authors.
import warnings
+import pickle
from random import Random
import itertools as it
from pytest import raises
@@ -795,3 +796,8 @@ def test_ModesLead_and_SelfEnergyLead():
fsyst = syst.finalized()
ts2 = [kwant.greens_function(fsyst, e).transmission(1, 0) for e in energies]
assert_almost_equal(ts2, ts)
+
+
+def test_site_pickle():
+ site = kwant.lattice.square()(0, 0)
+ assert pickle.loads(pickle.dumps(site)) == site
--
GitLab