diff --git a/kwant/_colormaps.py b/kwant/_colormaps.py
new file mode 100644
index 0000000000000000000000000000000000000000..86b11cb52ba83023b68b500fce58aaa5aa1adfe1
--- /dev/null
+++ b/kwant/_colormaps.py
@@ -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")
diff --git a/kwant/plotter.py b/kwant/plotter.py
index 962a7568a76b63242a7aa710a93dab2aab61417f..5920819e024e2f762db7a4f891405e7845860d7e 100644
--- a/kwant/plotter.py
+++ b/kwant/plotter.py
@@ -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
diff --git a/kwant/tests/test_plotter.py b/kwant/tests/test_plotter.py
index 69f019993c1cc4cfdc42b957159685c574e307a8..bec75f2d8ac3aa83e0e3a800cf698691df47b08d 100644
--- a/kwant/tests/test_plotter.py
+++ b/kwant/tests/test_plotter.py
@@ -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)
diff --git a/kwant_red.jscm b/kwant_red.jscm
new file mode 100644
index 0000000000000000000000000000000000000000..fd1f6485331d743bf8d1a512ef5db1cd7881f251
--- /dev/null
+++ b/kwant_red.jscm
@@ -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"
+}