new file mode 100644

@@ -0,0 +1,274 @@

+# Copyright 2011-2017 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

+# http://kwant-project.org/license.  A list of Kwant authors can be found in

+# the file AUTHORS.rst at the top-level directory of this distribution and at

+# http://kwant-project.org/authors.

+

+import numpy as np

+from matplotlib.colors import ListedColormap

+

+

+kr_data = [[ 0.98916316, 0.98474381, 0.99210697],

+           [ 0.98723538, 0.98138853, 0.98740721],

+           [ 0.98524761, 0.97805594, 0.98280815],

+           [ 0.98323731, 0.97473672, 0.97825805],

+           [ 0.98120432, 0.97143086, 0.97375773],

+           [ 0.97914854, 0.96813832, 0.96930811],

+           [ 0.97707118, 0.9648587 , 0.96490859],

+           [ 0.97496489, 0.96159373, 0.96057068],

+           [ 0.97282361, 0.95834495, 0.95630336],

+           [ 0.9706266 , 0.95511878, 0.95212518],

+           [ 0.96829617, 0.95194965, 0.94800577],

+           [ 0.96603689, 0.94881665, 0.94335712],

+           [ 0.9642474 , 0.94554816, 0.93828198],

+           [ 0.96264371, 0.94222348, 0.93310824],

+           [ 0.96112591, 0.93887825, 0.92787613],

+           [ 0.9596737 , 0.93552033, 0.92258582],

+           [ 0.95826857, 0.93215528, 0.91725354],

+           [ 0.95691078, 0.92878338, 0.91187323],

+           [ 0.95558706, 0.92540812, 0.90646129],

+           [ 0.95429677, 0.92202971, 0.90101648],

+           [ 0.95303987, 0.9186481 , 0.89553731],

+           [ 0.95180895, 0.91526515, 0.89003356],

+           [ 0.95060268, 0.91188113, 0.88450626],

+           [ 0.94942096, 0.90849598, 0.87895478],

+           [ 0.94826527, 0.9051092 , 0.87337621],

+           [ 0.94713012, 0.90172214, 0.86777815],

+           [ 0.9460154 , 0.89833471, 0.86216021],

+           [ 0.94491934, 0.89494729, 0.85652453],

+           [ 0.94384211, 0.89155971, 0.85087045],

+           [ 0.94278282, 0.88817211, 0.84519885],

+           [ 0.9417403 , 0.88478469, 0.83951108],

+           [ 0.94071449, 0.88139736, 0.83380686],

+           [ 0.93970482, 0.87801015, 0.8280867 ],

+           [ 0.93869093, 0.87463154, 0.82235188],

+           [ 0.93767439, 0.87126198, 0.81659016],

+           [ 0.93665457, 0.86790165, 0.81080186],

+           [ 0.93563073, 0.86455072, 0.80498748],

+           [ 0.9346022 , 0.86120939, 0.79914753],

+           [ 0.93356835, 0.85787779, 0.7932825 ],

+           [ 0.93252858, 0.85455608, 0.78739289],

+           [ 0.93148232, 0.8512444 , 0.78147919],

+           [ 0.93042904, 0.84794286, 0.77554188],

+           [ 0.92936825, 0.84465159, 0.76958144],

+           [ 0.92829947, 0.84137068, 0.76359835],

+           [ 0.92722226, 0.83810024, 0.75759308],

+           [ 0.92613621, 0.83484035, 0.75156611],

+           [ 0.92504093, 0.83159108, 0.74551788],

+           [ 0.92393607, 0.8283525 , 0.73944885],

+           [ 0.92282127, 0.82512468, 0.73335947],

+           [ 0.92169623, 0.82190767, 0.72725019],

+           [ 0.92056065, 0.81870151, 0.72112144],

+           [ 0.91941428, 0.81550625, 0.71497363],

+           [ 0.91825684, 0.8123219 , 0.70880719],

+           [ 0.91708812, 0.80914851, 0.70262254],

+           [ 0.91590788, 0.80598608, 0.69642011],

+           [ 0.91471593, 0.80283464, 0.69020029],

+           [ 0.91351208, 0.79969419, 0.6839635 ],

+           [ 0.91229619, 0.79656474, 0.67771009],

+           [ 0.9110681 , 0.79344626, 0.67144042],

+           [ 0.9098277 , 0.79033876, 0.66515486],

+           [ 0.90857484, 0.78724222, 0.65885379],

+           [ 0.90730944, 0.78415662, 0.65253753],

+           [ 0.90603137, 0.78108195, 0.64620649],

+           [ 0.9047406 , 0.77801816, 0.63986091],

+           [ 0.90343706, 0.77496521, 0.63350109],

+           [ 0.90212043, 0.77192324, 0.62712707],

+           [ 0.9007902 , 0.76889244, 0.62073872],

+           [ 0.89944703, 0.76587239, 0.61433695],

+           [ 0.8980909 , 0.76286305, 0.60792204],

+           [ 0.8967218 , 0.75986435, 0.60149417],

+           [ 0.89533975, 0.75687624, 0.59505355],

+           [ 0.89394477, 0.75389866, 0.58860036],

+           [ 0.89253687, 0.75093154, 0.58213482],

+           [ 0.89111598, 0.74797484, 0.57565735],

+           [ 0.88968224, 0.74502847, 0.56916787],

+           [ 0.88823569, 0.74209235, 0.56266646],

+           [ 0.88677639, 0.7391664 , 0.55615322],

+           [ 0.88530439, 0.73625056, 0.5496283 ],

+           [ 0.88381954, 0.73334485, 0.54309156],

+           [ 0.88232071, 0.7304499 , 0.53654178],

+           [ 0.88080928, 0.72756485, 0.52998058],

+           [ 0.87928542, 0.72468959, 0.52340769],

+           [ 0.87774918, 0.72182404, 0.51682313],

+           [ 0.87620068, 0.71896812, 0.51022682],

+           [ 0.87464002, 0.71612171, 0.50361862],

+           [ 0.8730672 , 0.71328477, 0.49699875],

+           [ 0.87148244, 0.71045715, 0.49036671],

+           [ 0.86988582, 0.70763879, 0.48372234],

+           [ 0.86827546, 0.70483075, 0.47706323],

+           [ 0.8666533 , 0.70203183, 0.47039155],

+           [ 0.86501961, 0.69924189, 0.46370656],

+           [ 0.86337451, 0.69646083, 0.45700802],

+           [ 0.86171805, 0.69368856, 0.45029569],

+           [ 0.86005037, 0.69092498, 0.4435692 ],

+           [ 0.8583717 , 0.68816997, 0.43682759],

+           [ 0.85668069, 0.68542427, 0.43006904],

+           [ 0.85497774, 0.68268759, 0.42329375],

+           [ 0.85326416, 0.6799592 , 0.41650135],

+           [ 0.85154001, 0.67723901, 0.40969136],

+           [ 0.84980535, 0.67452693, 0.40286337],

+           [ 0.84806043, 0.67182284, 0.39601592],

+           [ 0.84630751, 0.66912636, 0.38913447],

+           [ 0.84459103, 0.66641219, 0.38224671],

+           [ 0.84291325, 0.66368084, 0.37532092],

+           [ 0.84127642, 0.66093085, 0.36835537],

+           [ 0.8396827 , 0.65816075, 0.36134959],

+           [ 0.83813404, 0.65536906, 0.3543045 ],

+           [ 0.83663304, 0.65255404, 0.34721892],

+           [ 0.83518227, 0.64971389, 0.34009261],

+           [ 0.83378463, 0.64684664, 0.33292467],

+           [ 0.83244282, 0.64395026, 0.32571585],

+           [ 0.83116067, 0.64102228, 0.31846359],

+           [ 0.82994084, 0.63806051, 0.31117051],

+           [ 0.82878668, 0.63506246, 0.30383759],

+           [ 0.82770268, 0.63202504, 0.29646294],

+           [ 0.82669177, 0.62894569, 0.28905098],

+           [ 0.8257594 , 0.62581985, 0.28161649],

+           [ 0.82491057, 0.62264328, 0.27417247],

+           [ 0.82414639, 0.61941489, 0.26670947],

+           [ 0.82346948, 0.6161318 , 0.25923681],

+           [ 0.82288698, 0.61278636, 0.25181095],

+           [ 0.8223949 , 0.6093803 , 0.24442135],

+           [ 0.82199347, 0.60591095, 0.23710189],

+           [ 0.8216812 , 0.6023757 , 0.22990946],

+           [ 0.82144868, 0.59877883, 0.22284964],

+           [ 0.82128972, 0.59511966, 0.21600609],

+           [ 0.82118758, 0.591407  , 0.20938918],

+           [ 0.82112964, 0.58764517, 0.20306727],

+           [ 0.8210973 , 0.58384411, 0.19706615],

+           [ 0.82107295, 0.58001337, 0.19141507],

+           [ 0.82104144, 0.57616127, 0.18614637],

+           [ 0.820987  , 0.57229831, 0.18125098],

+           [ 0.82089884, 0.56843126, 0.17674506],

+           [ 0.82076836, 0.56456672, 0.17261538],

+           [ 0.82058591, 0.56071287, 0.16882561],

+           [ 0.82035412, 0.55686763, 0.1654141 ],

+           [ 0.82006554, 0.55303913, 0.162307  ],

+           [ 0.81971999, 0.54922882, 0.15948908],

+           [ 0.8193204 , 0.54543542, 0.15696633],

+           [ 0.81886751, 0.54166009, 0.15471108],

+           [ 0.81836021, 0.53790534, 0.15268999],

+           [ 0.81780262, 0.53416961, 0.15088739],

+           [ 0.81719522, 0.53045365, 0.14928762],

+           [ 0.81654166, 0.52675603, 0.1478763 ],

+           [ 0.8158438 , 0.52307639, 0.14663934],

+           [ 0.81510297, 0.51941472, 0.14556293],

+           [ 0.8143219 , 0.51576993, 0.14463469],

+           [ 0.81350256, 0.5121414 , 0.14384267],

+           [ 0.81264676, 0.50852856, 0.14317557],

+           [ 0.81175539, 0.50493143, 0.14262199],

+           [ 0.81083033, 0.50134931, 0.14217199],

+           [ 0.80987366, 0.49778124, 0.14181673],

+           [ 0.80888697, 0.49422657, 0.14154764],

+           [ 0.80787088, 0.49068512, 0.14136263],

+           [ 0.80682868, 0.48715438, 0.1412702 ],

+           [ 0.80576011, 0.48363561, 0.14124257],

+           [ 0.80466553, 0.48012899, 0.14127189],

+           [ 0.80355703, 0.47662617, 0.14134325],

+           [ 0.80244307, 0.47312268, 0.14140455],

+           [ 0.80132288, 0.46961902, 0.14145468],

+           [ 0.80019799, 0.46611382, 0.14149602],

+           [ 0.7990684 , 0.46260695, 0.14152863],

+           [ 0.79793402, 0.45909833, 0.14155245],

+           [ 0.79679443, 0.45558816, 0.14156689],

+           [ 0.79564905, 0.45207679, 0.14157115],

+           [ 0.79449895, 0.44856316, 0.1415669 ],

+           [ 0.79334425, 0.44504703, 0.1415543 ],

+           [ 0.7921849 , 0.44152826, 0.14153333],

+           [ 0.791021  , 0.43800659, 0.14150417],

+           [ 0.7898508 , 0.43448337, 0.14146434],

+           [ 0.78867629, 0.4309567 , 0.14141669],

+           [ 0.78749716, 0.42742668, 0.14136083],

+           [ 0.7863135 , 0.42389299, 0.14129692],

+           [ 0.78512536, 0.42035543, 0.14122502],

+           [ 0.7839324 , 0.41681408, 0.14114469],

+           [ 0.78273406, 0.41326924, 0.14105512],

+           [ 0.78153119, 0.40971994, 0.14095756],

+           [ 0.78032382, 0.4061659 , 0.14085208],

+           [ 0.77911231, 0.40260657, 0.14073912],

+           [ 0.77789631, 0.399042  , 0.14061828],

+           [ 0.77667586, 0.39547192, 0.1404896 ],

+           [ 0.77545064, 0.39189639, 0.14035261],

+           [ 0.77422054, 0.38831521, 0.14020724],

+           [ 0.77298612, 0.38472759, 0.14005425],

+           [ 0.77174749, 0.38113312, 0.13989379],

+           [ 0.77050465, 0.37753149, 0.13972588],

+           [ 0.76925755, 0.37392242, 0.13955047],

+           [ 0.76800622, 0.37030556, 0.13936762],

+           [ 0.76675073, 0.36668048, 0.13917742],

+           [ 0.76549072, 0.36304722, 0.13897942],

+           [ 0.76422604, 0.35940553, 0.13877348],

+           [ 0.76295723, 0.35575449, 0.13856031],

+           [ 0.7616843 , 0.35209368, 0.13833995],

+           [ 0.76040723, 0.34842269, 0.13811243],

+           [ 0.75912615, 0.34474093, 0.1378779 ],

+           [ 0.75784098, 0.34104803, 0.13763632],

+           [ 0.75655175, 0.33734346, 0.13738776],

+           [ 0.75525845, 0.33362676, 0.13713223],

+           [ 0.75396105, 0.32989738, 0.13686976],

+           [ 0.75265958, 0.32615477, 0.13660042],

+           [ 0.75135404, 0.32239833, 0.13632426],

+           [ 0.75004443, 0.31862745, 0.13604133],

+           [ 0.74873043, 0.31484189, 0.13575126],

+           [ 0.7474124 , 0.31104052, 0.1354546 ],

+           [ 0.74609036, 0.30722265, 0.1351514 ],

+           [ 0.74476429, 0.30338754, 0.13484171],

+           [ 0.7434342 , 0.2995344 , 0.13452559],

+           [ 0.74210012, 0.29566234, 0.13420317],

+           [ 0.740762  , 0.29177058, 0.13387443],

+           [ 0.73941986, 0.28785816, 0.13353947],

+           [ 0.73807366, 0.28392417, 0.13319832],

+           [ 0.73672338, 0.27996758, 0.13285105],

+           [ 0.73536902, 0.27598733, 0.13249773],

+           [ 0.73401057, 0.27198225, 0.13213842],

+           [ 0.73264812, 0.267951  , 0.13177331],

+           [ 0.7312816 , 0.26389237, 0.13140239],

+           [ 0.72991091, 0.25980509, 0.13102567],

+           [ 0.72853607, 0.25568766, 0.13064324],

+           [ 0.72715709, 0.25153846, 0.13025522],

+           [ 0.72577404, 0.24735572, 0.12986175],

+           [ 0.72438675, 0.24313784, 0.12946276],

+           [ 0.72299506, 0.23888311, 0.12905819],

+           [ 0.72159891, 0.23458955, 0.12864803],

+           [ 0.72019841, 0.23025473, 0.12823254],

+           [ 0.71879351, 0.22587627, 0.12781177],

+           [ 0.71738445, 0.22145117, 0.12738605],

+           [ 0.71597106, 0.2169768 , 0.12695531],

+           [ 0.71455619, 0.21244509, 0.12652294],

+           [ 0.71313673, 0.20785768, 0.12608696],

+           [ 0.71171317, 0.20320993, 0.12564796],

+           [ 0.71028524, 0.19849817, 0.12520579],

+           [ 0.70885023, 0.19372288, 0.12475808],

+           [ 0.70741049, 0.18887484, 0.12430712],

+           [ 0.70596683, 0.18394695, 0.12385374],

+           [ 0.70451965, 0.17893205, 0.12339841],

+           [ 0.70306509, 0.17383098, 0.12293779],

+           [ 0.70160541, 0.16863153, 0.12247404],

+           [ 0.70014347, 0.16331865, 0.12200979],

+           [ 0.69867317, 0.15789552, 0.12153977],

+           [ 0.69719976, 0.15233907, 0.12106873],

+           [ 0.69571992, 0.14664349, 0.12059389],

+           [ 0.69423655, 0.14078634, 0.12011789],

+           [ 0.69274566, 0.13475932, 0.11963739],

+           [ 0.69125189, 0.12852876, 0.11915653],

+           [ 0.68975221, 0.12207639, 0.1186728 ],

+           [ 0.68824567, 0.11537334, 0.11818552],

+           [ 0.6867347 , 0.10837306, 0.11769691],

+           [ 0.68521904, 0.10102592, 0.11720685],

+           [ 0.68369886, 0.09326599, 0.11671559],

+           [ 0.68217313, 0.08501006, 0.11622239],

+           [ 0.68064138, 0.07614058, 0.11572698],

+           [ 0.67910607, 0.06647368, 0.11523153],

+           [ 0.67756743, 0.05573816, 0.11473634],

+           [ 0.67602649, 0.04347155, 0.11424238],

+           [ 0.67448806, 0.02955075, 0.11375372],

+           [ 0.67299504, 0.0153711 , 0.1133058 ]]

+

+# Rescale colormap to start from white.

+kr_data = np.array(kr_data)

+kr_data = np.clip(kr_data / kr_data[0], 0, 1)

+

+kwant_red = ListedColormap(kr_data, name="kwant red")

@@ -20,6 +20,7 @@ import functools

 import warnings

 import cmath

 import numpy as np

+import scipy

 import tinyarray as ta

 from scipy import spatial, interpolate

 from math import cos, sin, pi, sqrt

@@ -32,7 +33,9 @@ try:

     import matplotlib

     from matplotlib.figure import Figure

     from matplotlib import collections

+    from matplotlib import colors

     from matplotlib.backends.backend_agg import FigureCanvasAgg

+    from . import _colormaps

     mpl_enabled = True

     try:

         from mpl_toolkits import mplot3d

@@ -1581,6 +1584,9 @@ def map(sys, value, colorbar=True, cmap=None, vmin=None, vmax=None, a=None,

     else:

         fig = None

+    if cmap is None:

+        cmap = _colormaps.kwant_red

+

     # Note that we tell imshow to show the array created by mask_interpolate

     # faithfully and not to interpolate by itself another time.

     image = ax.imshow(img.T, extent=(min[0], max[0], min[1], max[1]),

@@ -1794,6 +1800,233 @@ def spectrum(syst, x, y=None, params=None, mask=None, file=None,

         return output_fig(fig, file=file, show=show)

+def interpolate_current(syst, current, sigma=1, gauss_range=2, n=3, a=None):

+    """Interpolate currents in a system onto a regular grid.

+

+    The system graph together with current intensities defines a "discrete"

+    current density field where the current density is non-zero only on the

+    straight lines that connect sites that are coupled by a hopping term.

+

+    To make this vector field easier to visualize and interpret at different

+    length scales, it is smoothed by convoluting it with a kernel function

+    that has a (Gaussian) pulse-like shape.

+

+    This function samples the smoothed field on a regular (square or cubic)

+    grid.

+

+    Parameters

+    ----------

+    syst : A finalized system

+        The system on which we are going to calculate the field.

+    current : '1D array of float'

+        Must contain the intensity on each hoppings in the same order that they

+        appear in syst.graph.

+    sigma : 'float'

+        Parameter of the Gaussian used to generate the field :

+        exp(-x**2/(2*sigma**2)), the unit of sigma is a.

+    gauss_range : int

+        Number of sigma after which we consider the Gaussian equal to 0.

+    n : int

+        Number of point the grid must have per sigma.

+    a : float

+        By default, the length of the shortest hoppings.

+

+    Returns

+    -------

+    region : list of the coordonates of the grid for each dimension

+    field : value of the generated field on the grid points

+    """

+    if not isinstance(syst, builder.FiniteSystem):

+        raise TypeError("The system needs to be finalized.")

+

+    if len(current) != syst.graph.num_edges:

+        raise ValueError("Current and hoppings arrays do not have the same"

+                         " length.")

+

+    # hops: hoppings (pairs of points)

+    dim = len(syst.sites[0].pos)

+    hops = np.empty((syst.graph.num_edges, 2, dim))

+    for k, (i, j) in enumerate(syst.graph):

+        # +ve direction defined as *from* j *to* i

+        hops[k][0] = syst.sites[j].pos

+        hops[k][1] = syst.sites[i].pos

+

+    # lens: scaled lengths of hoppings

+    # dirs: normalized directions of hoppings

+    dirs = hops[:, 1] - hops[:, 0]

+    lens = np.sqrt(np.sum(dirs * dirs, -1))

+    dirs /= lens[:, None]

+    if a == None:

+        a = min(lens)

+    sigma *= a

+    r = sigma * gauss_range

+    factor = 0.5 * pow(sigma * sqrt(2*pi), 1 - dim)

+    scale = 1 / (sigma * sqrt(2))

+    lens *= scale

+

+    # Create field array.

+    field_shape = np.zeros(dim + 1, int)

+    field_shape[dim] = dim

+    min_hops = np.min(hops, 1)

+    max_hops = np.max(hops, 1)

+    bbox_min = np.min(min_hops, 0)

+    bbox_max = np.max(max_hops, 0)

+    for d in range(dim):

+        field_shape[d] = int((bbox_max[d] - bbox_min[d])*n / sigma + 2*gauss_range*n)

+        if field_shape[d] % 2:

+            field_shape[d] += 1

+    field = np.zeros(field_shape)

+

+    region = [np.linspace(bbox_min[d] - gauss_range*sigma, bbox_max[d] + gauss_range*sigma,

+                          field_shape[d])

+              for d in range(dim)]

+

+    grid_density = (field_shape[:dim] - 1) / (bbox_max + 2*r - bbox_min)

+    slices = np.empty((len(hops), dim, 2), int)

+    slices[:, :, 0] = np.floor((min_hops - bbox_min) * grid_density)

+    slices[:, :, 1] = np.ceil((max_hops + 2*r - bbox_min) * grid_density)

+

+    # Interpolate the field for each hopping.

+    erf = scipy.special.erf

+    for i in range(len(hops)):

+        if not np.diff(slices[i]).all():

+            # Zero volume: nothing to do.

+            continue

+

+        field_slice = [slice(*slices[i, d]) for d in range(dim)]

+

+        # Coordinates of the grid points that are within range of the current

+        # hopping.

+        coords = np.meshgrid(*[region[d][field_slice[d]] for d in range(dim)],

+                             sparse=True, indexing='ij')

+        coords -= hops[i, 0]

+

+        # Convert "coords" into scaled cylindrical coordinates with regard to

+        # the hopping.

+        z = np.dot(coords, dirs[i])

+        rho = np.sqrt(np.abs(np.sum(coords * coords) - z*z))

+        z *= scale

+        rho *= scale

+

+        magns = current[i] * factor * np.exp(-rho * rho)

+        magns *= erf(z) - erf(z - lens[i])

+

+        field[field_slice] += dirs[i] * magns[..., None]

+

+    # 'field' contains contributions from both hoppings (i, j) and (j, i)

+    return region, field / 2

+

+

+def current(syst, current, sigma=1, gauss_range=6, n=3, a=None,

+            colorbar=True, cmap=None, file=None, show=True, dpi=None,

+            fig_size=None, ax=None):

+    """Show interpolated map of a function defined for the hoppings of a system.

+

+    The system graph together with current intensities defines a "discrete"

+    current density field where the current density is non-zero only on the

+    straight lines that connect sites that are coupled by a hopping term.

+

+    To make this vector field easier to visualize and interpret at different

+    length scales, it is smoothed by convoluting it with a kernel function

+    that has a (Gaussian) pulse-like shape.

+

+    This function samples the smoothed field on a regular (square or cubic)

+    grid.

+

+    Parameters

+    ----------

+    syst : `kwant.system.FiniteSystem`

+        The system for which to plot the ``current``.

+    current : sequence of float

+        Sequence of values defining currents on each hopping of the system.

+        Ordered in the same way as ``syst.graph``. This typically will be

+        the result of evaluating a `~kwant.operator.Current` operator.

+    sigma : 'float'

+        Parameter of the Gaussian used to generate the field:

+        exp(-x**2/(2*sigma**2)), the unit of sigma is ``a``.

+    gauss_range : int

+        Number of sigma after which we consider the Gaussian equal to 0.

+    n : int

+        Number of points the grid must have per sigma.

+    a : float

+        By default, the length of the shortest hoppings.

+    colorbar : bool, optional

+        Whether to show a color bar if numerical data has to be plotted.

+        Defaults to `True`. If `ax` is provided, the colorbar is never plotted.

+    cmap : ``matplotlib`` color map or `None`

+        The color map used for sites and optionally hoppings, if `None`,

+        ``matplotlib`` default is used.

+    file : string or file object or `None`

+        The output file.  If `None`, output will be shown instead.

+    show : bool

+        Whether ``matplotlib.pyplot.show()`` is to be called, and the output is

+        to be shown immediately.  Defaults to `True`.

+    dpi : float or `None`

+        Number of pixels per inch.  If not set the ``matplotlib`` default is

+        used.

+    fig_size : tuple or `None`

+        Figure size `(width, height)` in inches.  If not set, the default

+        ``matplotlib`` value is used.

+    ax : ``matplotlib.axes.Axes`` instance or `None`

+        If `ax` is not `None`, no new figure is created, but the plot is done

+        within the existing Axes `ax`. in this case, `file`, `show`, `dpi`

+        and `fig_size` are ignored.

+

+    Returns

+    -------

+    fig : matplotlib figure

+        A figure with the output if `ax` is not set, else None.

+    """

+

+    if not mpl_enabled:

+        raise RuntimeError("matplotlib was not found, but is required "

+                           "for current()")

+

+    region, field = interpolate_current(syst, current, sigma=sigma,

+                                        gauss_range=gauss_range, n=n, a=a)

+

+    if len(region) != 2:

+        raise ValueError("Only 2D field can be plotted.")

+

+    # The default colormap is extremely ugly with streamplot.

+    if cmap is None:

+        cmap = _colormaps.kwant_red

+

+    if ax is None:

+        fig = Figure()

+        if dpi is not None:

+            fig.set_dpi(dpi)

+        if fig_size is not None:

+            fig.set_figwidth(fig_size[0])

+            fig.set_figheight(fig_size[1])

+        ax = fig.add_subplot(1, 1, 1, aspect='equal')

+    else:

+        fig = None

+

+    Y, X = np.ogrid[region[0][0] : region[0][-1] : 1j * len(region[0]),

+                    region[1][0] : region[1][-1] : 1j * len(region[1])]

+

+    field = field.transpose(1, 0, 2)  # reverse X and Y axes

+    speed = np.linalg.norm(field, axis=-1)

+    image = ax.imshow(speed[::-1], cmap=cmap,

+                      interpolation='bicubic',

+                      extent=(region[0][0], region[0][-1],

+                              region[1][0], region[1][-1]))

+    linewidth = 2 * (speed / (np.max(speed) or 1))

+    ax.streamplot(Y.T, X.T, field[:,:,0], field[:,:,1],

+                  linewidth=linewidth, cmap=cmap, color=speed,

+                  norm=colors.Normalize(vmin=0, vmax=np.max(speed) * 1.5))

+

+    ax.set_xlim(region[0][0], region[0][-1])

+    ax.set_ylim(region[1][0], region[1][-1])

+

+    if colorbar and fig is not None:

+        fig.colorbar(image)

+

+    if fig is not None:

+        return output_fig(fig, file=file, show=show)

+

+

 # TODO (Anton): Fix plotting of parts of the system using color = np.nan.

 # Not plotting sites currently works, not plotting hoppings does not.

 # TODO (Anton): Allow a more flexible treatment of position than pos_transform

@@ -8,10 +8,16 @@

 import tempfile

 import warnings

+import itertools

 import numpy as np

+import tinyarray as ta

+from math import cos, sin, pi

+import scipy.integrate

+import pytest

+

 import kwant

 from kwant import plotter

-import pytest

+from kwant._common import ensure_rng

 if plotter.mpl_enabled:

     from mpl_toolkits import mplot3d

@@ -40,7 +46,7 @@ def test_importable_without_matplotlib():

 def syst_2d(W=3, r1=3, r2=8):

     a = 1

     t = 1.0

-    lat = kwant.lattice.square(a)

+    lat = kwant.lattice.square(a, norbs=1)

     syst = kwant.Builder()

     def ring(pos):

@@ -237,3 +243,176 @@ def test_spectrum():

     with tempfile.TemporaryFile('w+b') as out:

         plotter.spectrum(ham, ('a', vals), ('b', 2 * vals), params=dict(c=1),

                          mask=mask, file=out)

+

+

+def syst_rect(lat, salt, W=3, L=50):

+    syst = kwant.Builder()

+

+    ll = L//2

+    ww = W//2

+

+    def onsite(site):

+        return 4 + 0.1 * kwant.digest.gauss(site.tag, salt=salt)

+

+    syst[(lat(i, j) for i in range(-ll, ll+1)

+         for j in range(-ww, ww+1))] = onsite

+    syst[lat.neighbors()] = -1

+

+    sym = kwant.TranslationalSymmetry(lat.vec((-1, 0)))

+    lead = kwant.Builder(sym)

+    lead[(lat(0, j) for j in range(-ww, ww + 1))] = 4

+    lead[lat.neighbors()] = -1

+

+    syst.attach_lead(lead)

+    syst.attach_lead(lead.reversed())

+

+    return syst

+

+

+def div(F):

+    """Calculate the divergence of a vector field F."""

+    assert len(F.shape[:-1]) == F.shape[-1]

+    return sum(np.gradient(F[..., i])[i] for i in range(F.shape[-1]))

+

+

+def rotational_currents(g):

+    """Return a basis of divergence-free currents for a closed graph.

+

+    Given the graph 'g' of a Kwant system, returns a sequence of arrays

+    which are linearly independent, divergence-free currents on the graph.

+    """

+    #'A' represents the set of expressions that give the net current flow

+    # into the system sites. 'perm' is a map from the edges of a graph

+    # with only 1 edge per hopping to the proper Kwant graph (2 edges

+    # per hopping).

+    A = np.zeros((g.num_nodes, g.num_edges // 2))

+    hoppings = dict()

+    perm_data = np.zeros(g.num_edges)

+    perm_ij = np.zeros((2, g.num_edges))

+    i = 0

+    for k, (a, b) in enumerate(g):

+        hop = frozenset((a, b))

+        if hop not in hoppings:

+            A[a, i] = 1

+            A[b, i] = -1

+            hoppings[hop] = i

+            perm_data[k] = 1

+            perm_ij[:, k] = (k, i)

+            i += 1

+        else:

+            perm_data[k] = -1

+            perm_ij[:, k] = (k, hoppings[hop])

+

+    perm = scipy.sparse.coo_matrix((perm_data, perm_ij))

+

+    # Get the row vectors of V with singular value 0. These form

+    # a basis for the right null space of 'A'.

+    U, S, V = np.linalg.svd(A)

+    tol = S.max() * max(A.shape) * np.finfo(S.dtype).eps

+    rank = sum(S > tol)

+    # Transform null space basis into vectors defined over the full

+    # hopping space (both hopping directions).

+    null_space_basis = V[-(len(V) - rank):].transpose()

+    null_space_basis = perm.dot(null_space_basis).transpose()

+    return null_space_basis

+

+

+def test_current_interpolation():

+

+    ## Passing a Builder will raise an error

+    pytest.raises(TypeError, plotter.interpolate_current, syst_2d(), None)

+

+    def R(theta):

+        return ta.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

+

+    def make_lattice(a, theta):

+        x = ta.dot(R(theta), (a, 0))

+        y = ta.dot(R(theta), (0, a))

+        return kwant.lattice.general([x, y], norbs=1)

+

+    angles = (0, pi/6, pi/4)

+    lat_constants = (1, 2)

+

+    ## Check current through cross section is same for different lattice

+    ## parameters and orientations of the system wrt. the discretization grid

+    for a, theta in itertools.product(lat_constants, angles):

+        lat = make_lattice(a, theta)

+        syst = syst_rect(lat, salt='0').finalized()

+        psi = kwant.wave_function(syst, energy=3)(0)

+

+        def cut(a, b):

+            return b.tag[0] < 0 and a.tag[0] >= 0

+

+        J = kwant.operator.Current(syst).bind()

+        J_cut = kwant.operator.Current(syst, where=cut, sum=True).bind()

+        (x, y), j0 = plotter.interpolate_current(syst, J(psi[0]), gauss_range=5)

+

+        # slice field perpendicular to a cut along the y axis

+        y_axis = (np.argmin(np.abs(x)), slice(None), 0)

+        ## Integrate and compare with summed current.

+        assert np.isclose(J_cut(psi[0]),

+                          scipy.integrate.simps(j0[y_axis], y))

+

+    ## Check that taking a finer grid or changing the broadening does not

+    ## affect the total integrated current.

+    n_s = (3, 5)

+    sigma_s = (1, 0.5)

+

+    for n, sigma in zip(n_s, sigma_s):

+        (x, y), j0 = plotter.interpolate_current(syst, J(psi[0]), n=n,

+                                                 sigma=sigma, gauss_range=5)

+        # slice field perpendicular to a cut along the y axis

+        y_axis = (np.argmin(np.abs(x)), slice(None), 0)

+        assert np.isclose(J_cut(psi[0]),

+                          scipy.integrate.simps(j0[y_axis], y))

+

+    ### Tests on a divergence-free current (closed system)

+

+    lat = kwant.lattice.general([(1, 0), (0.5, np.sqrt(3) / 2)])

+    syst = kwant.Builder()

+    sites = [lat(0, 0), lat(1, 0), lat(0, 1), lat(2, 2)]

+    syst[sites] = None

+    syst[((s, t) for s, t in itertools.product(sites, sites) if s != t)] = None

+    del syst[lat(0, 0), lat(2, 2)]

+    syst = syst.finalized()

+

+    # generate random divergence-free currents

+    Js = rotational_currents(syst.graph)

+    rng = ensure_rng(3)

+    J0 = sum(rng.rand(len(Js))[:, None] * Js)

+    J1 = sum(rng.rand(len(Js))[:, None] * Js)

+

+    # Sanity check that diverence on the graph is 0

+    divergence = np.zeros(len(syst.sites))

+    for (a, _), current in zip(syst.graph, J0):

+        divergence[a] += current

+    assert np.allclose(divergence, 0)

+

+    _, j0 = plotter.interpolate_current(syst, J0)

+    _, j1 = plotter.interpolate_current(syst, J1)

+

+    ## Test linearity of interpolation.

+    _, j_tot = plotter.interpolate_current(syst, J0 + 2 * J1)

+    assert np.allclose(j_tot, j0 + 2 * j1)

+

+    ## Test that divergence of interpolated current is approximately zero.

+    # For currents not aligned with the interpolation grid this is only

+    # 1/a**2 accurate.

+    _, j = plotter.interpolate_current(syst, J0, n=20, gauss_range=4)

+    div_j = np.max(np.abs(div(j)))

+    assert np.isclose(div_j, 0, atol=1E-4)

+

+

+@pytest.mark.skipif(not plotter.mpl_enabled, reason="No matplotlib available.")

+def test_current():

+    syst = syst_2d().finalized()

+    J = kwant.operator.Current(syst)

+    current = J(kwant.wave_function(syst, energy=1)(1)[0])

+

+    # Test good codepath

+    with tempfile.TemporaryFile('w+b') as out:

+        plotter.current(syst, current, file=out)

+

+        fig = pyplot.Figure()

+        ax = fig.add_subplot(1, 1, 1)

+        plotter.current(syst, current, ax=ax, file=out)

new file mode 100644

@@ -0,0 +1,38 @@

+{

+    "license": "http://creativecommons.org/publicdomain/zero/1.0/",

+    "usage-hints": [

+        "red-green-colorblind-safe",

+        "greyscale-safe",

+        "sequential"

+    ],

+    "colors": "fcfbfdfcfafcfbf9fbfbf9f9faf8f8faf7f7f9f6f6f9f5f5f8f4f4f8f4f3f7f3f2f6f2f1f6f1eff5f0eef5efedf5efebf4eeeaf4ede9f4ece7f3ebe6f3eae4f3e9e3f2e9e2f2e8e0f2e7dff2e6ddf1e5dcf1e4daf1e3d9f0e2d8f0e2d6f0e1d5f0e0d3efdfd2efded0efddcfefdccdeedccceedbcaeedac9eed9c7edd8c6edd7c4edd7c3ecd6c1ecd5c0ecd4beecd3bdebd2bbebd2b9ebd1b8ead0b6eacfb5eaceb3eaceb2e9cdb0e9ccaee9cbade8caabe8caaae8c9a8e7c8a6e7c7a5e7c6a3e6c6a2e6c5a0e6c49ee5c39de5c39be5c299e4c198e4c096e4bf94e3bf93e3be91e3bd8fe2bc8ee2bc8ce1bb8ae1ba89e1ba87e0b985e0b884dfb782dfb780dfb67fdeb57ddeb47bddb47addb378ddb276dcb275dcb173dbb071dbaf6fdaaf6edaae6cdaad6ad9ad68d9ac67d8ab65d8ab63d7aa61d7a960d7a95ed6a85cd6a75ad5a659d5a657d5a555d4a453d4a351d4a34fd3a24dd3a14cd3a04ad3a048d29f46d29e44d29d42d29c40d29b3ed29b3cd29a3bd19939d19837d19735d19634d19532d19431d1932fd1922ed1912dd1902cd18f2bd18e2ad18d29d18c29d18b28d18a27d18927d18826d08726d08626d08525d08425d08425cf8325cf8225cf8124cf8024cf7f24ce7e24ce7d24ce7c24cd7b24cd7a24cd7a24cd7924cc7824cc7724cc7624cb7524cb7424cb7324cb7224ca7124ca7124ca7024c96f24c96e24c96d24c96c24c86b24c86a24c86924c76824c76824c76724c66624c66524c66424c56324c56224c56124c46024c45f24c45e24c45e23c35d23c35c23c35b23c25a23c25923c25823c15723c15623c15523c05423c05323c05223bf5123bf5023bf4f23be4e22be4d22be4c22bd4b22bd4a22bd4922bc4822bc4722bc4622bb4522bb4422ba4322ba4221ba4121b94021b93f21b93e21b83d21b83c21b83b21b73a21b73820b73720b63620b63520b53420b53320b53120b43020b42f20b42e1fb32c1fb32b1fb32a1fb2281fb2271fb1251fb1241fb1221fb0211eb01f1eb01d1eaf1c1eaf1a1eae181eae161eae131ead111dad0e1dac0b1dac081dac041d",

+    "colorspace": "sRGB",

+    "name": "kwant red",

+    "extensions": {

+        "https://matplotlib.org/viscm": {

+            "yp": [

+                -2.0153712378334347,

+                10.887492636952686,

+                28.27830916383833,

+                20.84513758379849,

+                15.655942329808425

+            ],

+            "max_Jp": 41,

+            "cmtype": "linear",

+            "uniform_colorspace": "CAM02-UCS",

+            "spline_method": "CatmulClark",

+            "fixed": -1,

+            "filter_k": 100,

+            "min_Jp": 99,

+            "xp": [

+                -1.2674960029171558,

+                5.464432975232114,

+                4.482693332585356,

+                21.45276429833666,

+                33.65438557123224

+            ]

+        }

+    },

+    "domain": "continuous",

+    "content-type": "application/vnd.matplotlib.colormap-v1+json"

+}