Coverage for kwant/wraparound.py : 97%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2011-2018 Kwant authors.
2#
3# This file is part of Kwant. It is subject to the license terms in the file
4# LICENSE.rst found in the top-level directory of this distribution and at
5# https://kwant-project.org/license. A list of Kwant authors can be found in
6# the file AUTHORS.rst at the top-level directory of this distribution and at
7# https://kwant-project.org/authors.
9import collections
10import inspect
11import cmath
12import warnings
14import tinyarray as ta
15import numpy as np
16import scipy.linalg
17import scipy.spatial
19from . import builder, system, plotter
20from .linalg import lll
21from .builder import herm_conj, HermConjOfFunc
22from .lattice import TranslationalSymmetry
23from ._common import get_parameters, memoize
26__all__ = ['wraparound', 'plot_2d_bands']
29def _set_signature(func, params):
30 """Set the signature of 'func'.
32 Parameters
33 ----------
34 func : callable
35 params: sequence of str
36 Parameter names to put in the signature. These will be added as
37 'POSITIONAL_ONLY' type parameters.
38 """
39 params = [inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY)
40 for name in params]
41 func.__signature__ = inspect.Signature(params)
43@memoize
44def _callable_herm_conj(val):
45 """Keep the same id for every 'val'."""
46 return HermConjOfFunc(val)
48## This wrapper is needed so that finalized systems that
49## have been wrapped can be queried for their symmetry, which
50## is needed for Brillouin zone calculations (plotting).
52class WrappedBuilder(builder.Builder):
54 def finalized(self):
55 ret = super().finalized()
56 ret._momentum_names = self._momentum_names
57 ret._wrapped_symmetry = self._wrapped_symmetry
58 return ret
61def wraparound(builder, keep=None, *, coordinate_names='xyz'):
62 """Replace translational symmetries by momentum parameters.
64 A new Builder instance is returned. By default, each symmetry is replaced
65 by one scalar momentum parameter that is appended to the already existing
66 arguments of the system. Optionally, one symmetry may be kept by using the
67 `keep` argument. The momentum parameters will have names like 'k_n' where
68 the 'n' are specified by 'coordinate_names'.
70 Parameters
71 ----------
72 builder : `~kwant.builder.Builder`
73 keep : int, optional
74 Which (if any) translational symmetry to keep.
75 coordinate_names : sequence of strings, default: 'xyz'
76 The names of the coordinates along the symmetry
77 directions of 'builder'.
79 Notes
80 -----
81 Wraparound is stop-gap functionality until Kwant 2.x which will include
82 support for higher-dimension translational symmetry in the low-level system
83 format. It will be deprecated in the 2.0 release of Kwant.
84 """
86 @memoize
87 def bind_site(val):
88 def f(*args):
89 a, *args = args
90 return val(a, *args[:mnp])
92 assert callable(val)
93 _set_signature(f, get_parameters(val) + momenta)
94 return f
96 @memoize
97 def bind_hopping_as_site(elem, val):
98 def f(*args):
99 a, *args = args
100 phase = cmath.exp(1j * ta.dot(elem, args[mnp:]))
101 v = val(a, sym.act(elem, a), *args[:mnp]) if callable(val) else val
102 pv = phase * v
103 return pv + herm_conj(pv)
105 params = ('_site0',)
106 if callable(val):
107 params += get_parameters(val)[2:] # cut off both site parameters
108 _set_signature(f, params + momenta)
109 return f
111 @memoize
112 def bind_hopping(elem, val):
113 def f(*args):
114 a, b, *args = args
115 phase = cmath.exp(1j * ta.dot(elem, args[mnp:]))
116 v = val(a, sym.act(elem, b), *args[:mnp]) if callable(val) else val
117 return phase * v
119 params = ('_site0', '_site1')
120 if callable(val):
121 params += get_parameters(val)[2:] # cut off site parameters
122 _set_signature(f, params + momenta)
123 return f
125 @memoize
126 def bind_sum(num_sites, *vals):
127 """Construct joint signature for all 'vals'."""
129 def f(*in_args):
130 acc = 0
131 for val, selection in val_selection_pairs:
132 if selection: # Otherwise: reuse previous out_args.
133 out_args = tuple(in_args[i] for i in selection)
134 if callable(val):
135 acc = acc + val(*out_args)
136 else:
137 acc = acc + val
138 return acc
140 params = collections.OrderedDict()
142 # Add the leading one or two 'site' parameters.
143 site_params = ['_site{}'.format(i) for i in range(num_sites)]
144 for name in site_params:
145 params[name] = inspect.Parameter(
146 name, inspect.Parameter.POSITIONAL_ONLY)
148 # Add all the other parameters (except for the momenta). Setup the
149 # 'selections'.
150 selections = []
151 for val in vals:
152 if not callable(val):
153 selections.append(())
154 continue
155 val_params = get_parameters(val)[num_sites:]
156 assert val_params[mnp:] == momenta
157 val_params = val_params[:mnp]
158 selections.append((*site_params, *val_params))
159 for p in val_params:
160 # Skip parameters that exist in previously added functions.
161 if p in params:
162 continue
163 params[p] = inspect.Parameter(
164 p, inspect.Parameter.POSITIONAL_ONLY)
166 # Sort values such that ones with the same arguments are bunched.
167 # Prepare 'val_selection_pairs' that is used in the function 'f' above.
168 params_keys = list(params.keys())
169 val_selection_pairs = []
170 prev_selection = None
171 argsort = sorted(range(len(selections)), key=selections.__getitem__)
172 momenta_sel = tuple(range(mnp, 0, 1))
173 for i in argsort:
174 selection = selections[i]
175 if selection and selection != prev_selection:
176 prev_selection = selection = tuple(
177 params_keys.index(s) for s in selection) + momenta_sel
178 else:
179 selection = ()
180 val_selection_pairs.append((vals[i], selection))
182 # Finally, add the momenta.
183 for k in momenta:
184 params[k] = inspect.Parameter(
185 k, inspect.Parameter.POSITIONAL_ONLY)
187 f.__signature__ = inspect.Signature(params.values())
188 return f
190 try:
191 momenta = ['k_{}'.format(coordinate_names[i])
192 for i in range(len(builder.symmetry.periods))]
193 except IndexError:
194 raise ValueError("All symmetry directions must have a name specified "
195 "in coordinate_names")
197 if keep is None:
198 ret = WrappedBuilder()
199 sym = builder.symmetry
200 else:
201 periods = list(builder.symmetry.periods)
202 ret = WrappedBuilder(TranslationalSymmetry(periods.pop(keep)))
203 sym = TranslationalSymmetry(*periods)
204 momenta.pop(keep)
205 momenta = tuple(momenta)
206 mnp = -len(momenta) # Used by the bound functions above.
208 # Store the names of the momentum parameters and the symmetry of the
209 # old Builder (this will be needed for band structure plotting)
210 ret._momentum_names = momenta
211 ret._wrapped_symmetry = builder.symmetry
213 # Wrapped around system retains conservation law and chiral symmetry.
214 # We use 'bind_site' to add the momenta arguments if required.
215 cons = builder.conservation_law
216 ret.conservation_law = bind_site(cons) if callable(cons) else cons
217 chiral = builder.chiral
218 ret.chiral = bind_site(chiral) if callable(chiral) else chiral
220 if builder.particle_hole is not None or builder.time_reversal is not None:
221 warnings.warn('`particle_hole` and `time_reversal` symmetries are set '
222 'on the input builder. However they are ignored for the '
223 'wrapped system, since Kwant lacks a way to express the '
224 'existence (or not) of a symmetry at k != 0.',
225 RuntimeWarning, stacklevel=2)
226 ret.particle_hole = None
227 ret.time_reversal = None
229 ret.vectorize = builder.vectorize
231 sites = {}
232 hops = collections.defaultdict(list)
234 # Store lists of values, so that multiple values can be assigned to the
235 # same site or hopping.
236 for site, val in builder.site_value_pairs():
237 # Every 'site' is in the FD of the original symmetry.
238 # Move the sites to the FD of the remaining symmetry, this guarantees that
239 # every site in the new system is an image of an original FD site translated
240 # purely by the remaining symmetry.
241 sites[ret.symmetry.to_fd(site)] = [val] # a list to append wrapped hoppings
243 for hop, val in builder.hopping_value_pairs():
244 a, b = hop
245 # 'a' is in the FD of original symmetry.
246 # Get translation from FD of original symmetry to 'b',
247 # this is different from 'b_dom = sym.which(b)'.
248 b_dom = builder.symmetry.which(b)
249 # Throw away part that is in the remaining translation direction, so we get
250 # an element of 'sym' which is being wrapped
251 b_dom = ta.array([t for i, t in enumerate(b_dom) if i != keep])
252 # Pull back using the remainder, which is purely in the wrapped directions.
253 # This guarantees that 'b_wa' is an image of an original FD site translated
254 # purely by the remaining symmetry.
255 b_wa = sym.act(-b_dom, b)
256 # Move the hopping to the FD of the remaining symmetry
257 a, b_wa = ret.symmetry.to_fd(a, b_wa)
259 if a == b_wa:
260 # The hopping gets wrapped-around into an onsite Hamiltonian.
261 # Since site `a` already exists in the system, we can simply append.
262 sites[a].append(bind_hopping_as_site(b_dom, val))
263 else:
264 # The hopping remains a hopping.
265 if any(b_dom):
266 # The hopping got wrapped-around.
267 val = bind_hopping(b_dom, val)
269 # Make sure that there is only one entry for each hopping
270 # pointing in one direction, modulo the remaining translations.
271 b_wa_r, a_r = ret.symmetry.to_fd(b_wa, a)
272 if (b_wa_r, a_r) in hops:
273 assert (a, b_wa) not in hops
274 if callable(val):
275 val = _callable_herm_conj(val)
276 else:
277 val = herm_conj(val)
279 hops[b_wa_r, a_r].append((val, b_dom))
280 else:
281 hops[a, b_wa].append((val, b_dom))
283 # Copy stuff into result builder, converting lists of more than one element
284 # into summing functions.
285 for site, vals in sites.items():
286 if len(vals) == 1:
287 # no need to bind onsites without extra wrapped hoppings
288 ret[site] = vals[0]
289 else:
290 val = vals[0]
291 vals[0] = bind_site(val) if callable(val) else val
292 ret[site] = bind_sum(1, *vals)
294 for hop, vals_doms in hops.items():
295 if len(vals_doms) == 1:
296 # no need to bind hoppings that are not already bound
297 val, b_dom = vals_doms[0]
298 ret[hop] = val
299 else:
300 new_vals = [bind_hopping(b_dom, val) if callable(val)
301 and not any(b_dom) # skip hoppings already bound
302 else val for val, b_dom in vals_doms]
303 ret[hop] = bind_sum(2, *new_vals)
304 return ret
307def plot_2d_bands(syst, k_x=31, k_y=31, params=None,
308 mask_brillouin_zone=False, extend_bbox=0, file=None,
309 show=True, dpi=None, fig_size=None, ax=None):
310 """Plot 2D band structure of a wrapped around system.
312 This function is primarily useful for systems that have translational
313 symmetry vectors that are non-orthogonal (e.g. graphene). This function
314 will properly plot the band structure in an orthonormal basis in k-space,
315 as opposed to in the basis of reciprocal lattice vectors (which would
316 produce a "skewed" Brillouin zone).
318 If your system has orthogonal lattice vectors, you are probably better
319 off using `kwant.plotter.spectrum`.
321 Parameters
322 ----------
323 syst : `kwant.system.FiniteSystem`
324 A 2D system that was finalized from a Builder produced by
325 `kwant.wraparound.wraparound`. Note that this *must* be a finite
326 system; so `kwant.wraparound.wraparound` should have been called with
327 ``keep=None``.
328 k_x, k_y : int or sequence of float, default: 31
329 Either a number of sampling points, or a sequence of points at which
330 the band structure is to be evaluated, in units of inverse length.
331 params : dict, optional
332 Dictionary of parameter names and their values, not including the
333 momentum parameters.
334 mask_brillouin_zone : bool, default: False
335 If True, then the band structure will only be plotted over the first
336 Brillouin zone. By default the band structure is plotted over a
337 rectangular bounding box that contains the Brillouin zone.
338 extend_bbox : float, default: 0
339 Amount by which to extend the region over which the band structure is
340 plotted, expressed as a proportion of the Brillouin zone bounding box
341 length. i.e. ``extend_bbox=0.1`` will extend the region by 10% (in all
342 directions).
343 file : string or file object, optional
344 The output file. If None, output will be shown instead.
345 show : bool, default: False
346 Whether ``matplotlib.pyplot.show()`` is to be called, and the output is
347 to be shown immediately. Defaults to `True`.
348 dpi : float, optional
349 Number of pixels per inch. If not set the ``matplotlib`` default is
350 used.
351 fig_size : tuple, optional
352 Figure size `(width, height)` in inches. If not set, the default
353 ``matplotlib`` value is used.
354 ax : ``matplotlib.axes.Axes`` instance, optional
355 If `ax` is not `None`, no new figure is created, but the plot is done
356 within the existing Axes `ax`. in this case, `file`, `show`, `dpi`
357 and `fig_size` are ignored.
359 Returns
360 -------
361 fig : matplotlib figure
362 A figure with the output if `ax` is not set, else None.
364 Notes
365 -----
366 This function produces plots where the units of momentum are inverse
367 length. This is contrary to `kwant.plotter.bands`, where the units
368 of momentum are inverse lattice constant.
370 If the lattice vectors for the symmetry of ``syst`` are not orthogonal,
371 then part of the plotted band structure will be outside the first Brillouin
372 zone (inside the bounding box of the brillouin zone). Setting
373 ``mask_brillouin_zone=True`` will cause the plot to be truncated outside of
374 the first Brillouin zone.
376 See Also
377 --------
378 kwant.plotter.spectrum
379 """
380 if not hasattr(syst, '_wrapped_symmetry'):
381 raise TypeError("Expecting a system that was produced by "
382 "'kwant.wraparound.wraparound'.")
383 if isinstance(syst, system.InfiniteSystem):
384 msg = ("All symmetry directions must be wrapped around: specify "
385 "'keep=None' when calling 'kwant.wraparound.wraparound'.")
386 raise TypeError(msg)
387 if isinstance(syst, builder.Builder): 387 ↛ 388line 387 didn't jump to line 388, because the condition on line 387 was never true
388 msg = ("Expecting a finalized system: remember to finalize your "
389 "system with 'syst.finalized()'.")
390 raise TypeError(msg)
392 params = params or {}
393 lat_ndim, space_ndim = syst._wrapped_symmetry.periods.shape
395 if lat_ndim != 2:
396 raise ValueError("Expected a system with a 2D translational symmetry.")
397 if space_ndim != lat_ndim: 397 ↛ 398line 397 didn't jump to line 398, because the condition on line 397 was never true
398 raise ValueError("Lattice dimension must equal realspace dimension.")
400 # columns of B are lattice vectors
401 B = np.array(syst._wrapped_symmetry.periods).T
402 # columns of A are reciprocal lattice vectors
403 A = np.linalg.pinv(B).T
405 ## calculate the bounding box for the 1st Brillouin zone
407 # Get lattice points that neighbor the origin, in basis of lattice vectors
408 reduced_vecs, transf = lll.lll(A.T)
409 neighbors = ta.dot(lll.voronoi(reduced_vecs), transf)
410 # Add the origin to these points.
411 klat_points = np.concatenate(([[0] * lat_ndim], neighbors))
412 # Transform to cartesian coordinates and rescale.
413 # Will be used in 'outside_bz' function, later on.
414 klat_points = 2 * np.pi * np.dot(klat_points, A.T)
415 # Calculate the Voronoi cell vertices
416 vor = scipy.spatial.Voronoi(klat_points)
417 around_origin = vor.point_region[0]
418 bz_vertices = vor.vertices[vor.regions[around_origin]]
419 # extract bounding box
420 k_max = np.max(np.abs(bz_vertices), axis=0)
422 ## build grid along each axis, if needed
423 ks = []
424 for k, km in zip((k_x, k_y), k_max):
425 k = np.array(k)
426 if not k.shape:
427 if extend_bbox:
428 km += km * extend_bbox
429 k = np.linspace(-km, km, k)
430 ks.append(k)
432 # TODO: It is very inefficient to call 'momentum_to_lattice' once for
433 # each point (for trivial Hamiltonians 60% of the time is spent
434 # doing this). We should instead transform the whole grid in one call.
436 def momentum_to_lattice(k):
437 k, residuals = scipy.linalg.lstsq(A, k)[:2]
438 if np.any(abs(residuals) > 1e-7): 438 ↛ 439line 438 didn't jump to line 439, because the condition on line 438 was never true
439 raise RuntimeError("Requested momentum doesn't correspond"
440 " to any lattice momentum.")
441 return k
443 def ham(k_x, k_y=None, **params):
444 # transform into the basis of reciprocal lattice vectors
445 k = momentum_to_lattice([k_x] if k_y is None else [k_x, k_y])
446 p = dict(zip(syst._momentum_names, k), **params)
447 return syst.hamiltonian_submatrix(params=p, sparse=False)
449 def outside_bz(k_x, k_y, **_):
450 dm = scipy.spatial.distance_matrix(klat_points, [[k_x, k_y]])
451 return np.argmin(dm) != 0 # is origin no closest 'klat_point' to 'k'?
453 fig = plotter.spectrum(ham,
454 x=('k_x', ks[0]),
455 y=('k_y', ks[1]) if lat_ndim == 2 else None,
456 params=params,
457 mask=(outside_bz if mask_brillouin_zone else None),
458 file=file, show=show, dpi=dpi,
459 fig_size=fig_size, ax=ax)
460 return fig