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