diff --git a/kwant/plotter.py b/kwant/plotter.py
index c096ddd0bd058cd3bbc4e422398d9be151ec3b0d..04e3f1c235ce7d2f2532c7b61f42576ffd4639bd 100644
--- a/kwant/plotter.py
+++ b/kwant/plotter.py
@@ -60,270 +60,271 @@ def isarray(var):
def nparray_if_array(var):
return np.asarray(var) if isarray(var) else var
-
-class LineCollection(collections.LineCollection):
- def __init__(self, segments, reflen=None, **kwargs):
- super(LineCollection, self).__init__(segments, **kwargs)
- self.reflen = reflen
-
- def set_linewidths(self, linewidths):
- self._linewidths_orig = nparray_if_array(linewidths)
-
- def draw(self, renderer):
- linewidths = self._linewidths_orig
- if self.reflen is not None:
- # Note: only works for aspect ratio 1!
- # 72.0 - there is 72 points in an inch
- factor = (self.axes.transData.frozen().to_values()[0] * 72.0 *
- self.reflen / self.figure.dpi)
- linewidths *= factor
-
- super(LineCollection, self).set_linewidths(linewidths)
- return super(LineCollection, self).draw(renderer)
-
-
-class PathCollection(collections.PathCollection):
- def __init__(self, paths, sizes=None, reflen=None, **kwargs):
- super(PathCollection, self).__init__(paths, sizes=sizes, **kwargs)
-
- self.reflen = reflen
- self._linewidths_orig = nparray_if_array(self.get_linewidths())
-
- self._transforms = [matplotlib.transforms.Affine2D().scale(x) for x in
- sizes]
-
- def get_transforms(self):
- return self._transforms
-
- def get_transform(self):
- Affine2D = matplotlib.transforms.Affine2D
- if self.reflen is not None:
- # For the paths, use the data transformation but strip the offset
- # (will be added later with offsets)
- args = self.axes.transData.frozen().to_values()[:4] + (0, 0)
- return Affine2D().from_values(*args).scale(self.reflen)
- else:
- return Affine2D().scale(self.figure.dpi / 72.0)
-
- def draw(self, renderer):
- if self.reflen:
- # Note: only works for aspect ratio 1!
- factor = (self.axes.transData.frozen().to_values()[0] /
- self.figure.dpi * 72.0 * self.reflen)
- self.set_linewidths(self._linewidths_orig * factor)
-
- return collections.Collection.draw(self, renderer)
-
-
-if has3d:
- # Sorting is optional.
- sort3d = True
-
- # Compute the projection of a 3D length into 2D data coordinates
- # for this we use 2 3D half-circles that are projected into 2D.
- # (This gives the same length as projecting the full unit sphere.)
-
- phi = np.linspace(0, pi, 21)
- xyz = np.c_[np.cos(phi), np.sin(phi), 0 * phi].T.reshape(-1, 1, 21)
- unit_sphere = np.bmat([[xyz[0], xyz[2]], [xyz[1], xyz[0]],
- [xyz[2], xyz[1]]])
- unit_sphere = np.asarray(unit_sphere)
-
- def projected_length(ax, length):
- rc = np.array([ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()])
- rc = np.apply_along_axis(np.sum, 1, rc) / 2.
-
- rs = unit_sphere * length + rc.reshape(-1, 1)
-
- transform = mplot3d.proj3d.proj_transform
- rp = np.asarray(transform(*(list(rs) + [ax.get_proj()]))[:2])
- rc[:2] = transform(*(list(rc) + [ax.get_proj()]))[:2]
-
- coords = rp - np.repeat(rc[:2].reshape(-1, 1), len(rs[0]), axis=1)
- return sqrt(np.sum(coords**2, axis=0).max())
-
-
- # Auxiliary array for calculating corners of a cube.
- corners = np.zeros((3, 8, 6), np.float_)
- corners[0, [0, 1, 2, 3], 0] = corners[0, [4, 5, 6, 7], 1] = \
- corners[0, [0, 1, 4, 5], 2] = corners[0, [2, 3, 6, 7], 3] = \
- corners[0, [0, 2, 4, 6], 4] = corners[0, [1, 3, 5, 7], 5] = 1.
-
-
- class Line3DCollection(mplot3d.art3d.Line3DCollection):
- def __init__(self, segments, reflen=None, zorder=0, **kwargs):
- super(Line3DCollection, self).__init__(segments, **kwargs)
+if mpl_enabled:
+ class LineCollection(collections.LineCollection):
+ def __init__(self, segments, reflen=None, **kwargs):
+ super(LineCollection, self).__init__(segments, **kwargs)
self.reflen = reflen
- self._zorder3d = zorder
def set_linewidths(self, linewidths):
self._linewidths_orig = nparray_if_array(linewidths)
- def do_3d_projection(self, renderer):
- super(Line3DCollection, self).do_3d_projection(renderer)
- # The whole 3D ordering is flawed in mplot3d when several
- # collections are added. We just use normal zorder. Note the
- # "-" due to the different logic in the 3d plotting, we still want
- # larger zorder values to be plotted on top of smaller ones.
- return -self._zorder3d
-
def draw(self, renderer):
linewidths = self._linewidths_orig
- if self.reflen:
- proj_len = projected_length(self.axes, self.reflen)
- args = self.axes.transData.frozen().to_values()
- # Note: unlike in the 2D case, where we can enforce equal
- # aspect ratio, this (currently) does not work with
- # 3D plots in matplotlib. As an approximation, we
- # thus scale with the average of the x- and y-axis
- # transformation.
- factor = proj_len * (args[0] +
- args[3]) * 0.5 * 72.0 / self.figure.dpi
+ if self.reflen is not None:
+ # Note: only works for aspect ratio 1!
+ # 72.0 - there is 72 points in an inch
+ factor = (self.axes.transData.frozen().to_values()[0] * 72.0 *
+ self.reflen / self.figure.dpi)
linewidths *= factor
- super(Line3DCollection, self).set_linewidths(linewidths)
- super(Line3DCollection, self).draw(renderer)
-
-
- class Path3DCollection(mplot3d.art3d.Patch3DCollection):
- def __init__(self, paths, sizes, reflen=None, zorder=0, offsets=None,
- **kwargs):
- paths = [matplotlib.patches.PathPatch(path) for path in paths]
+ super(LineCollection, self).set_linewidths(linewidths)
+ return super(LineCollection, self).draw(renderer)
- if offsets is not None:
- kwargs['offsets'] = offsets[:, :2]
- super(Path3DCollection, self).__init__(paths, **kwargs)
-
- if offsets is not None:
- self.set_3d_properties(zs=offsets[:, 2], zdir="z")
+ class PathCollection(collections.PathCollection):
+ def __init__(self, paths, sizes=None, reflen=None, **kwargs):
+ super(PathCollection, self).__init__(paths, sizes=sizes, **kwargs)
self.reflen = reflen
- self._zorder3d = zorder
-
- self._paths_orig = np.array(paths, dtype='object')
self._linewidths_orig = nparray_if_array(self.get_linewidths())
- self._linewidths_orig2 = self._linewidths_orig
- self._array_orig = nparray_if_array(self.get_array())
- self._facecolors_orig = nparray_if_array(self.get_facecolors())
- self._edgecolors_orig = nparray_if_array(self.get_edgecolors())
-
- Affine2D = matplotlib.transforms.Affine2D
- self._orig_transforms = np.array([Affine2D().scale(x) for x in
- sizes], dtype='object')
- self._transforms = self._orig_transforms
- def set_array(self, array):
- self._array_orig = nparray_if_array(array)
- super(Path3DCollection, self).set_array(array)
-
- def set_color(self, colors):
- self._facecolors_orig = nparray_if_array(colors)
- self._edgecolors_orig = self._facecolors_orig
- super(Path3DCollection, self).set_color(colors)
-
- def set_edgecolors(self, colors):
- colors = matplotlib.colors.colorConverter.to_rgba_array(colors)
- self._edgecolors_orig = nparray_if_array(colors)
- super(Path3DCollection, self).set_edgecolors(colors)
+ self._transforms = [matplotlib.transforms.Affine2D().scale(x) for x
+ in sizes]
def get_transforms(self):
- # this is exact only for an isometric projection, for the
- # perspective projection used in mplot3d it's an approximation
return self._transforms
def get_transform(self):
Affine2D = matplotlib.transforms.Affine2D
- if self.reflen:
- proj_len = projected_length(self.axes, self.reflen)
-
+ if self.reflen is not None:
# For the paths, use the data transformation but strip the
- # offset (will be added later with the offsets).
+ # offset (will be added later with offsets)
args = self.axes.transData.frozen().to_values()[:4] + (0, 0)
- return Affine2D().from_values(*args).scale(proj_len)
+ return Affine2D().from_values(*args).scale(self.reflen)
else:
return Affine2D().scale(self.figure.dpi / 72.0)
- def do_3d_projection(self, renderer):
- xs, ys, zs = self._offsets3d
+ def draw(self, renderer):
+ if self.reflen:
+ # Note: only works for aspect ratio 1!
+ factor = (self.axes.transData.frozen().to_values()[0] /
+ self.figure.dpi * 72.0 * self.reflen)
+ self.set_linewidths(self._linewidths_orig * factor)
+
+ return collections.Collection.draw(self, renderer)
+
+
+ if has3d:
+ # Sorting is optional.
+ sort3d = True
+
+ # Compute the projection of a 3D length into 2D data coordinates
+ # for this we use 2 3D half-circles that are projected into 2D.
+ # (This gives the same length as projecting the full unit sphere.)
+
+ phi = np.linspace(0, pi, 21)
+ xyz = np.c_[np.cos(phi), np.sin(phi), 0 * phi].T.reshape(-1, 1, 21)
+ unit_sphere = np.bmat([[xyz[0], xyz[2]], [xyz[1], xyz[0]],
+ [xyz[2], xyz[1]]])
+ unit_sphere = np.asarray(unit_sphere)
- # numpy complains about zero-length index arrays
- if len(xs) == 0:
+ def projected_length(ax, length):
+ rc = np.array([ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()])
+ rc = np.apply_along_axis(np.sum, 1, rc) / 2.
+
+ rs = unit_sphere * length + rc.reshape(-1, 1)
+
+ transform = mplot3d.proj3d.proj_transform
+ rp = np.asarray(transform(*(list(rs) + [ax.get_proj()]))[:2])
+ rc[:2] = transform(*(list(rc) + [ax.get_proj()]))[:2]
+
+ coords = rp - np.repeat(rc[:2].reshape(-1, 1), len(rs[0]), axis=1)
+ return sqrt(np.sum(coords**2, axis=0).max())
+
+
+ # Auxiliary array for calculating corners of a cube.
+ corners = np.zeros((3, 8, 6), np.float_)
+ corners[0, [0, 1, 2, 3], 0] = corners[0, [4, 5, 6, 7], 1] = \
+ corners[0, [0, 1, 4, 5], 2] = corners[0, [2, 3, 6, 7], 3] = \
+ corners[0, [0, 2, 4, 6], 4] = corners[0, [1, 3, 5, 7], 5] = 1.
+
+
+ class Line3DCollection(mplot3d.art3d.Line3DCollection):
+ def __init__(self, segments, reflen=None, zorder=0, **kwargs):
+ super(Line3DCollection, self).__init__(segments, **kwargs)
+ self.reflen = reflen
+ self._zorder3d = zorder
+
+ def set_linewidths(self, linewidths):
+ self._linewidths_orig = nparray_if_array(linewidths)
+
+ def do_3d_projection(self, renderer):
+ super(Line3DCollection, self).do_3d_projection(renderer)
+ # The whole 3D ordering is flawed in mplot3d when several
+ # collections are added. We just use normal zorder. Note the
+ # "-" due to the different logic in the 3d plotting, we still
+ # want larger zorder values to be plotted on top of smaller
+ # ones.
return -self._zorder3d
- proj = mplot3d.proj3d.proj_transform_clip
- vs = np.array(proj(xs, ys, zs, renderer.M)[:3])
-
- if sort3d:
- indx = vs[2].argsort()[::-1]
-
- self.set_offsets(vs[:2, indx].T)
-
- if len(self._paths_orig) > 1:
- paths = np.resize(self._paths_orig, (vs.shape[1],))
- self.set_paths(paths[indx])
-
- if len(self._orig_transforms) > 1:
- self._transforms = np.resize(self._orig_transforms,
- (vs.shape[1],))
- self._transforms = self._transforms[indx]
-
- if (isinstance(self._linewidths_orig, np.ndarray) and
- len(self._linewidths_orig) > 1):
- self._linewidths_orig2 = np.resize(self._linewidths_orig,
- (vs.shape[1],))[indx]
-
- # Note: here array, facecolors and edgecolors are guaranteed to
- # be 2d numpy arrays or None. (And array is the same
- # length as the coordinates)
-
- if self._array_orig is not None:
- super(Path3DCollection,
- self).set_array(self._array_orig[indx])
-
- if (self._facecolors_orig is not None and
- self._facecolors_orig.shape[0] > 1):
- shape = list(self._facecolors_orig.shape)
- shape[0] = vs.shape[0]
- super(Path3DCollection, self).set_facecolors(
- np.resize(self._facecolors_orig, shape)[indx])
-
- if (self._edgecolors_orig is not None and
- self._edgecolors_orig.shape[0] > 1):
- shape = list(self._edgecolors_orig.shape)
- shape[0] = vs.shape[0]
- super(Path3DCollection, self).set_edgecolors(
- np.resize(self._edgecolors_orig,
- shape)[indx])
- else:
- self.set_offsets(vs[:2].T)
+ def draw(self, renderer):
+ linewidths = self._linewidths_orig
+ if self.reflen:
+ proj_len = projected_length(self.axes, self.reflen)
+ args = self.axes.transData.frozen().to_values()
+ # Note: unlike in the 2D case, where we can enforce equal
+ # aspect ratio, this (currently) does not work with
+ # 3D plots in matplotlib. As an approximation, we
+ # thus scale with the average of the x- and y-axis
+ # transformation.
+ factor = proj_len * (args[0] +
+ args[3]) * 0.5 * 72.0 / self.figure.dpi
+ linewidths *= factor
+
+ super(Line3DCollection, self).set_linewidths(linewidths)
+ super(Line3DCollection, self).draw(renderer)
+
+
+ class Path3DCollection(mplot3d.art3d.Patch3DCollection):
+ def __init__(self, paths, sizes, reflen=None, zorder=0,
+ offsets=None, **kwargs):
+ paths = [matplotlib.patches.PathPatch(path) for path in paths]
+
+ if offsets is not None:
+ kwargs['offsets'] = offsets[:, :2]
+
+ super(Path3DCollection, self).__init__(paths, **kwargs)
+
+ if offsets is not None:
+ self.set_3d_properties(zs=offsets[:, 2], zdir="z")
+
+ self.reflen = reflen
+ self._zorder3d = zorder
+
+ self._paths_orig = np.array(paths, dtype='object')
+ self._linewidths_orig = nparray_if_array(self.get_linewidths())
+ self._linewidths_orig2 = self._linewidths_orig
+ self._array_orig = nparray_if_array(self.get_array())
+ self._facecolors_orig = nparray_if_array(self.get_facecolors())
+ self._edgecolors_orig = nparray_if_array(self.get_edgecolors())
+
+ Affine2D = matplotlib.transforms.Affine2D
+ self._orig_transforms = np.array([Affine2D().scale(x) for x in
+ sizes], dtype='object')
+ self._transforms = self._orig_transforms
+
+ def set_array(self, array):
+ self._array_orig = nparray_if_array(array)
+ super(Path3DCollection, self).set_array(array)
+
+ def set_color(self, colors):
+ self._facecolors_orig = nparray_if_array(colors)
+ self._edgecolors_orig = self._facecolors_orig
+ super(Path3DCollection, self).set_color(colors)
+
+ def set_edgecolors(self, colors):
+ colors = matplotlib.colors.colorConverter.to_rgba_array(colors)
+ self._edgecolors_orig = nparray_if_array(colors)
+ super(Path3DCollection, self).set_edgecolors(colors)
+
+ def get_transforms(self):
+ # this is exact only for an isometric projection, for the
+ # perspective projection used in mplot3d it's an approximation
+ return self._transforms
+
+ def get_transform(self):
+ Affine2D = matplotlib.transforms.Affine2D
+ if self.reflen:
+ proj_len = projected_length(self.axes, self.reflen)
+
+ # For the paths, use the data transformation but strip the
+ # offset (will be added later with the offsets).
+ args = self.axes.transData.frozen().to_values()[:4] + (0, 0)
+ return Affine2D().from_values(*args).scale(proj_len)
+ else:
+ return Affine2D().scale(self.figure.dpi / 72.0)
+
+ def do_3d_projection(self, renderer):
+ xs, ys, zs = self._offsets3d
+
+ # numpy complains about zero-length index arrays
+ if len(xs) == 0:
+ return -self._zorder3d
+
+ proj = mplot3d.proj3d.proj_transform_clip
+ vs = np.array(proj(xs, ys, zs, renderer.M)[:3])
+
+ if sort3d:
+ indx = vs[2].argsort()[::-1]
+
+ self.set_offsets(vs[:2, indx].T)
+
+ if len(self._paths_orig) > 1:
+ paths = np.resize(self._paths_orig, (vs.shape[1],))
+ self.set_paths(paths[indx])
+
+ if len(self._orig_transforms) > 1:
+ self._transforms = np.resize(self._orig_transforms,
+ (vs.shape[1],))
+ self._transforms = self._transforms[indx]
+
+ lw_orig = self._linewidths_orig
+ if (isinstance(lw_orig, np.ndarray) and len(lw_orig) > 1):
+ self._linewidths_orig2 = np.resize(lw_orig,
+ (vs.shape[1],))[indx]
+
+ # Note: here array, facecolors and edgecolors are
+ # guaranteed to be 2d numpy arrays or None. (And
+ # array is the same length as the coordinates)
+
+ if self._array_orig is not None:
+ super(Path3DCollection,
+ self).set_array(self._array_orig[indx])
+
+ if (self._facecolors_orig is not None and
+ self._facecolors_orig.shape[0] > 1):
+ shape = list(self._facecolors_orig.shape)
+ shape[0] = vs.shape[0]
+ super(Path3DCollection, self).set_facecolors(
+ np.resize(self._facecolors_orig, shape)[indx])
+
+ if (self._edgecolors_orig is not None and
+ self._edgecolors_orig.shape[0] > 1):
+ shape = list(self._edgecolors_orig.shape)
+ shape[0] = vs.shape[0]
+ super(Path3DCollection, self).set_edgecolors(
+ np.resize(self._edgecolors_orig,
+ shape)[indx])
+ else:
+ self.set_offsets(vs[:2].T)
- # the whole 3D ordering is flawed in mplot3d when several
- # collections are added. We just use normal zorder, but correct
- # by the projected z-coord of the "center of gravity", normalized
- # by the projected z-coord of the world coordinates.
- # In doing so, several Path3DCollections are plotted probably
- # in the right order (it's not exact) if they have the same
- # zorder. Still, smaller and larger integer zorders are plotted
- # below or on top.
+ # the whole 3D ordering is flawed in mplot3d when several
+ # collections are added. We just use normal zorder, but correct
+ # by the projected z-coord of the "center of gravity",
+ # normalized by the projected z-coord of the world coordinates.
+ # In doing so, several Path3DCollections are plotted probably
+ # in the right order (it's not exact) if they have the same
+ # zorder. Still, smaller and larger integer zorders are plotted
+ # below or on top.
- bbox = np.asarray(self.axes.get_w_lims())
+ bbox = np.asarray(self.axes.get_w_lims())
- proj = mplot3d.proj3d.proj_transform_clip
- cz = proj(*(list(np.dot(corners, bbox)) + [renderer.M]))[2]
+ proj = mplot3d.proj3d.proj_transform_clip
+ cz = proj(*(list(np.dot(corners, bbox)) + [renderer.M]))[2]
- return -self._zorder3d + vs[2].mean() / cz.ptp()
+ return -self._zorder3d + vs[2].mean() / cz.ptp()
- def draw(self, renderer):
- if self.reflen:
- proj_len = projected_length(self.axes, self.reflen)
- args = self.axes.transData.frozen().to_values()
- factor = proj_len * (args[0] +
- args[3]) * 0.5 * 72.0 / self.figure.dpi
+ def draw(self, renderer):
+ if self.reflen:
+ proj_len = projected_length(self.axes, self.reflen)
+ args = self.axes.transData.frozen().to_values()
+ factor = proj_len * (args[0] +
+ args[3]) * 0.5 * 72.0 / self.figure.dpi
- self.set_linewidths(self._linewidths_orig2 * factor)
+ self.set_linewidths(self._linewidths_orig2 * factor)
- super(Path3DCollection, self).draw(renderer)
+ super(Path3DCollection, self).draw(renderer)
# matplotlib helper functions.
diff --git a/kwant/tests/test_plotter.py b/kwant/tests/test_plotter.py
index bd12ef6f76eb964d742ef45be0c51f46566bdf28..55f87faff44492c7cabbea2dce025a17b7fc506c 100644
--- a/kwant/tests/test_plotter.py
+++ b/kwant/tests/test_plotter.py
@@ -26,7 +26,13 @@ def test_importable_without_matplotlib():
code = f.read()
code = code.replace('from . import', 'from kwant import')
code = code.replace('matplotlib', 'totalblimp')
- exec code
+
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ exec code # Trigger the warning.
+ nose.tools.assert_equal(len(w), 1)
+ assert issubclass(w[0].category, RuntimeWarning)
+ assert "only iterator-providing functions" in str(w[0].message)
def sys_2d(W=3, r1=3, r2=8):