diff --git a/kwant/_common.py b/kwant/_common.py index 48c45eaac9cb1e4ed4a2971c3bd58e0cebcb1c43..eda672a14e515a74e8bf5937a1a2ffb8d2dbc004 100644 --- a/kwant/_common.py +++ b/kwant/_common.py @@ -137,7 +137,6 @@ class lazy_import: def __getattr__(self, name): if self.__deprecation_warning: - absolute_module = '.'.join((self.__package, self.__module)) msg = ("Accessing {0} without an explicit import is deprecated. " "Instead, explicitly 'import {0}'." ).format('.'.join((self.__package, self.__module))) diff --git a/kwant/_plotter.py b/kwant/_plotter.py new file mode 100644 index 0000000000000000000000000000000000000000..4db45843b6b8c3a185c3807a5abd43ecea3b6f0e --- /dev/null +++ b/kwant/_plotter.py @@ -0,0 +1,302 @@ +# -*- coding: utf-8 -*- +# Copyright 2011-2018 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. + +# This module is imported by plotter.py. It contains all the expensive imports +# that we want to remove from plotter.py + +# All matplotlib imports must be isolated in a try, because even without +# matplotlib iterators remain useful. Further, mpl_toolkits used for 3D +# plotting are also imported separately, to ensure that 2D plotting works even +# if 3D does not. +try: + import matplotlib + import matplotlib.colors + import matplotlib.cm + from matplotlib.figure import Figure + from matplotlib import collections + from . import _colormaps + mpl_available = True + try: + from mpl_toolkits import mplot3d + has3d = True + except ImportError: + warnings.warn("3D plotting not available.", RuntimeWarning) + has3d = False +except ImportError: + warnings.warn("matplotlib is not available, only iterator-providing " + "functions will work.", RuntimeWarning) + mpl_available = False + + +if mpl_available: + class LineCollection(collections.LineCollection): + def __init__(self, segments, reflen=None, **kwargs): + super().__init__(segments, **kwargs) + self.reflen = reflen + + def set_linewidths(self, linewidths): + self.linewidths_orig = nparray_if_array(linewidths) + + def draw(self, renderer): + 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) + else: + factor = 1 + + super().set_linewidths(self.linewidths_orig * + factor) + return super().draw(renderer) + + + class PathCollection(collections.PathCollection): + def __init__(self, paths, sizes=None, reflen=None, **kwargs): + super().__init__(paths, sizes=sizes, **kwargs) + + self.reflen = reflen + self.linewidths_orig = nparray_if_array(self.get_linewidths()) + + self.transforms = np.array( + [matplotlib.transforms.Affine2D().scale(x).get_matrix() + 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.0 + + + class Line3DCollection(mplot3d.art3d.Line3DCollection): + def __init__(self, segments, reflen=None, zorder=0, **kwargs): + super().__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().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): + 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 + else: + factor = 1 + + super().set_linewidths( + self.linewidths_orig * factor) + super().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().__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).get_matrix() for x in sizes]) + self.transforms = self.orig_transforms + + def set_array(self, array): + self.array_orig = nparray_if_array(array) + super().set_array(array) + + def set_color(self, colors): + self.facecolors_orig = nparray_if_array(colors) + self.edgecolors_orig = self.facecolors_orig + super().set_color(colors) + + def set_edgecolors(self, colors): + colors = matplotlib.colors.colorConverter.to_rgba_array(colors) + self.edgecolors_orig = nparray_if_array(colors) + super().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 = 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[1] + super().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[1] + super().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. + + bbox = np.asarray(self.axes.get_w_lims()) + + proj = mplot3d.proj3d.proj_transform_clip + cz = proj(*(list(np.dot(corners, bbox)) + [renderer.M]))[2] + + 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 + + self.set_linewidths(self.linewidths_orig2 * factor) + + super().draw(renderer) diff --git a/kwant/plotter.py b/kwant/plotter.py index 2112e5b98334cd15e22ef1dcebc09aaa59ed61c9..7adb88d79274722a5840e4dfc83f1f9f99a79bf3 100644 --- a/kwant/plotter.py +++ b/kwant/plotter.py @@ -26,29 +26,6 @@ import tinyarray as ta from scipy import spatial, interpolate from math import cos, sin, pi, sqrt -# All matplotlib imports must be isolated in a try, because even without -# matplotlib iterators remain useful. Further, mpl_toolkits used for 3D -# plotting are also imported separately, to ensure that 2D plotting works even -# if 3D does not. -try: - import matplotlib - import matplotlib.colors - import matplotlib.cm - from matplotlib.figure import Figure - from matplotlib import collections - from . import _colormaps - mpl_available = True - try: - from mpl_toolkits import mplot3d - has3d = True - except ImportError: - warnings.warn("3D plotting not available.", RuntimeWarning) - has3d = False -except ImportError: - warnings.warn("matplotlib is not available, only iterator-providing " - "functions will work.", RuntimeWarning) - mpl_available = False - from . import system, builder, _common @@ -57,6 +34,10 @@ __all__ = ['plot', 'map', 'bands', 'spectrum', 'current', 'sys_leads_sites', 'sys_leads_hoppings', 'sys_leads_pos', 'sys_leads_hopping_pos', 'mask_interpolate'] +# All the expensive imports are done in _plotter.py. We lazy load the module +# to avoid slowing down the initial import of Kwant. +_p = _common.lazy_import('_plotter') + # Collections that allow for symbols and linewiths to be given in data space # (not for general use, only implement what's needed for plotter) @@ -77,274 +58,6 @@ def _sample_array(array, n_samples, rng=None): return array[rng.choice(range(la), min(n_samples, la))] -if mpl_available: - class LineCollection(collections.LineCollection): - def __init__(self, segments, reflen=None, **kwargs): - super().__init__(segments, **kwargs) - self.reflen = reflen - - def set_linewidths(self, linewidths): - self.linewidths_orig = nparray_if_array(linewidths) - - def draw(self, renderer): - 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) - else: - factor = 1 - - super().set_linewidths(self.linewidths_orig * - factor) - return super().draw(renderer) - - - class PathCollection(collections.PathCollection): - def __init__(self, paths, sizes=None, reflen=None, **kwargs): - super().__init__(paths, sizes=sizes, **kwargs) - - self.reflen = reflen - self.linewidths_orig = nparray_if_array(self.get_linewidths()) - - self.transforms = np.array( - [matplotlib.transforms.Affine2D().scale(x).get_matrix() - 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.0 - - - class Line3DCollection(mplot3d.art3d.Line3DCollection): - def __init__(self, segments, reflen=None, zorder=0, **kwargs): - super().__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().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): - 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 - else: - factor = 1 - - super().set_linewidths( - self.linewidths_orig * factor) - super().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().__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).get_matrix() for x in sizes]) - self.transforms = self.orig_transforms - - def set_array(self, array): - self.array_orig = nparray_if_array(array) - super().set_array(array) - - def set_color(self, colors): - self.facecolors_orig = nparray_if_array(colors) - self.edgecolors_orig = self.facecolors_orig - super().set_color(colors) - - def set_edgecolors(self, colors): - colors = matplotlib.colors.colorConverter.to_rgba_array(colors) - self.edgecolors_orig = nparray_if_array(colors) - super().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 = 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[1] - super().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[1] - super().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. - - bbox = np.asarray(self.axes.get_w_lims()) - - proj = mplot3d.proj3d.proj_transform_clip - cz = proj(*(list(np.dot(corners, bbox)) + [renderer.M]))[2] - - 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 - - self.set_linewidths(self.linewidths_orig2 * factor) - - super().draw(renderer) - - # matplotlib helper functions. def _make_figure(dpi, fig_size): diff --git a/pytest.ini b/pytest.ini index 6c8258daf2ed83d5fb6da872d379c2ad36d42c25..c7f56a4768a2f6f410c50545cc5740b85baef3d7 100644 --- a/pytest.ini +++ b/pytest.ini @@ -2,5 +2,6 @@ testpaths = kwant flakes-ignore = __init__.py UnusedImport + kwant/_plotter.py UnusedImport graph/tests/test_scotch.py UndefinedName graph/tests/test_dissection.py UndefinedName