From ec9d0fca70166e54f0c95d89408edb5277dcdb20 Mon Sep 17 00:00:00 2001
From: Christoph Groth <christoph.groth@cea.fr>
Date: Tue, 6 Feb 2018 18:03:10 +0100
Subject: [PATCH] add vmax option to streamplot
Thanks to Tibor Sekera for suggesting this feature.
---
doc/source/pre/whatsnew/1.4.rst | 7 +++++++
kwant/plotter.py | 18 +++++++++++++-----
2 files changed, 20 insertions(+), 5 deletions(-)
diff --git a/doc/source/pre/whatsnew/1.4.rst b/doc/source/pre/whatsnew/1.4.rst
index 92e4c1c0..ab00cd14 100644
--- a/doc/source/pre/whatsnew/1.4.rst
+++ b/doc/source/pre/whatsnew/1.4.rst
@@ -5,3 +5,10 @@ This article explains the user-visible changes in Kwant 1.4.0.
See also the `full list of changes up to the most recent bugfix
release of the 1.4 series
<https://gitlab.kwant-project.org/kwant/kwant/compare/v1.4.0...latest-1.4>`_.
+
+Configurable maximum velocity in stream plots
+---------------------------------------------
+The function `~kwant.plotter.streamplot` has got a new option ``vmax``. Note
+that this option is not available in `~kwant.plotter.current`. In order to use
+it, one has to call ``streamplot`` directly as shown in the docstring of
+``current``.
diff --git a/kwant/plotter.py b/kwant/plotter.py
index 0752ed13..2112e5b9 100644
--- a/kwant/plotter.py
+++ b/kwant/plotter.py
@@ -2011,7 +2011,8 @@ def _linear_cmap(a, b):
def streamplot(field, box, cmap=None, bgcolor=None, linecolor='k',
max_linewidth=3, min_linewidth=1, density=2/9,
colorbar=True, file=None,
- show=True, dpi=None, fig_size=None, ax=None):
+ show=True, dpi=None, fig_size=None, ax=None,
+ vmax=None):
"""Draw streamlines of a flow field in Kwant style
Solid colored streamlines are drawn, superimposed on a color plot of
@@ -2063,6 +2064,10 @@ def streamplot(field, box, cmap=None, bgcolor=None, linecolor='k',
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.
+ vmax : float or `None`
+ The upper saturation limit for the colormap; flows higher than
+ this will saturate. Note that there is no corresponding vmin
+ option, vmin being fixed at zero.
Returns
-------
@@ -2102,6 +2107,8 @@ def streamplot(field, box, cmap=None, bgcolor=None, linecolor='k',
Y = np.linspace(*box[1], num=field.shape[0])
speed = np.linalg.norm(field, axis=-1)
+ if vmax is None:
+ vmax = np.max(speed) or 1
if cmap is None:
ax.set_axis_bgcolor(bgcolor)
@@ -2109,12 +2116,13 @@ def streamplot(field, box, cmap=None, bgcolor=None, linecolor='k',
image = ax.imshow(speed, cmap=cmap,
interpolation='bicubic',
extent=[e for c in box for e in c],
- origin='lower')
+ origin='lower', vmin=0, vmax=vmax)
- linewidth = max_linewidth / (np.max(speed) or 1) * speed
+ linewidth = max_linewidth / vmax * speed
color = linewidth / min_linewidth
- linewidth[linewidth < min_linewidth] = min_linewidth
- color[color > 1] = 1
+ thin = linewidth < min_linewidth
+ linewidth[thin] = min_linewidth
+ color[~ thin] = 1
line_cmap = _linear_cmap(linecolor, bgcolor)
--
GitLab