From 5042d96a53084b303e0f369fb12c09743ee433a2 Mon Sep 17 00:00:00 2001
From: Joseph Weston <joseph.weston08@gmail.com>
Date: Wed, 5 Apr 2017 15:15:20 +0200
Subject: [PATCH] create a wrapper class for builders constructed by
 'wraparound'

This adds some extra attributed on finalization, which allow the
momentum parameter names and symmetry of the original Builder to
be accessed from the finalized wrapped-around builder.
---
 kwant/wraparound.py | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/kwant/wraparound.py b/kwant/wraparound.py
index f3417ac6..2146430d 100644
--- a/kwant/wraparound.py
+++ b/kwant/wraparound.py
@@ -11,7 +11,8 @@ import inspect
 import cmath
 import tinyarray as ta
 
-from .builder import Builder, herm_conj, HermConjOfFunc
+from . import builder
+from .builder import herm_conj, HermConjOfFunc
 from .lattice import TranslationalSymmetry
 from ._common import get_parameters
 
@@ -66,6 +67,19 @@ def _modify_signature(func, parameter_names, takes_kwargs):
     return func
 
 
+## This wrapper is needed so that finalized systems that
+## have been wrapped can be queried for their symmetry, which
+## is needed for Brillouin zone calculations (plotting).
+
+class WrappedBuilder(builder.Builder):
+
+    def finalized(self):
+        ret = super().finalized()
+        ret._momentum_names = self._momentum_names
+        ret._wrapped_symmetry = self._wrapped_symmetry
+        return ret
+
+
 def wraparound(builder, keep=None, *, coordinate_names=('x', 'y', 'z')):
     """Replace translational symmetries by momentum parameters.
 
@@ -220,13 +234,17 @@ def wraparound(builder, keep=None, *, coordinate_names=('x', 'y', 'z')):
                zip(coordinate_names, builder.symmetry.periods)]
 
     if keep is None:
-        ret = Builder()
+        ret = WrappedBuilder()
         sym = builder.symmetry
     else:
         periods = list(builder.symmetry.periods)
-        ret = Builder(TranslationalSymmetry(periods.pop(keep)))
+        ret = WrappedBuilder(TranslationalSymmetry(periods.pop(keep)))
         sym = TranslationalSymmetry(*periods)
         momenta.pop(keep)
+    # Store the names of the momentum parameters and the symmetry of the
+    # old Builder (this will be needed for band structure plotting)
+    ret._momentum_names = momenta
+    ret._wrapped_symmetry = builder.symmetry
 
     # Wrapped around system retains conservation law and chiral symmetry.
     # We use 'bind_site' to add the momenta arguments if required.
-- 
GitLab