Source code for mplcursors._pick_info

# Unsupported Artist classes: subclasses of AxesImage, QuadMesh (upstream could
# have a `format_coord`-like method); PolyCollection (picking is not well
# defined).

from collections import namedtuple
from contextlib import suppress
import copy
import functools
import inspect
from inspect import Signature
import itertools
from numbers import Integral
import re
import warnings
from weakref import WeakSet

from matplotlib import cbook
from matplotlib.axes import Axes
from matplotlib.backend_bases import RendererBase
from matplotlib.collections import (
    LineCollection, PatchCollection, PathCollection)
from matplotlib.container import BarContainer, ErrorbarContainer, StemContainer
from matplotlib.figure import Figure
from matplotlib.image import AxesImage
from matplotlib.lines import Line2D
from matplotlib.patches import Patch, PathPatch, Polygon, Rectangle
from matplotlib.quiver import Barbs, Quiver
from matplotlib.text import Text
from matplotlib.transforms import Affine2D
import numpy as np


PATCH_PICKRADIUS = 5  # FIXME Patches do not provide `pickradius`.


def _register_scatter():
    """
    Patch `PathCollection` and `scatter` to register their return values.

    This registration allows us to distinguish `PathCollection`s created by
    `Axes.scatter`, which should use point-like picking, from others, which
    should use path-like picking.  The former is more common, so we store the
    latter instead; this also lets us guess the type better if this module is
    imported late.
    """

    @functools.wraps(PathCollection.__init__)
    def __init__(self, *args, **kwargs):
        _nonscatter_pathcollections.add(self)
        return __init__.__wrapped__(self, *args, **kwargs)
    PathCollection.__init__ = __init__

    @functools.wraps(Axes.scatter)
    def scatter(*args, **kwargs):
        paths = scatter.__wrapped__(*args, **kwargs)
        with suppress(KeyError):
            _nonscatter_pathcollections.remove(paths)
        return paths
    Axes.scatter = scatter


_nonscatter_pathcollections = WeakSet()
_register_scatter()


def _is_scatter(artist):
    return (isinstance(artist, PathCollection)
            and artist not in _nonscatter_pathcollections)


def _artist_in_container(container):
    return next(filter(None, container.get_children()))


class ContainerArtist:
    """Workaround to make containers behave more like artists."""

    def __init__(self, container):
        self.container = container  # Guaranteed to be nonempty.
        # We can't weakref the Container (which subclasses tuple), so
        # we instead create a reference cycle between the Container and
        # the ContainerArtist; as no one else strongly references the
        # ContainerArtist, it will get GC'd whenever the Container is.
        vars(container).setdefault(
            f"_{__class__.__name__}__keep_alive", []).append(self)

    def __str__(self):
        return f"<{type(self).__name__}({self.container})>"

    def __repr__(self):
        return f"<{type(self).__name__}({self.container!r})>"

    figure = property(lambda self: _artist_in_container(self.container).figure)
    axes = property(lambda self: _artist_in_container(self.container).axes)


class AttrArray(np.ndarray):
    """An array subclass that can store additional attributes."""

    def __new__(cls, array):
        return np.asarray(array).view(cls)


def _with_attrs(array, **kwargs):
    array = AttrArray(array)
    for k, v in kwargs.items():
        setattr(array, k, v)
    return array


Selection = namedtuple("Selection", "artist target dist annotation extras")
# Override equality to identity: Selections should be considered immutable
# (with mutable fields though) and we don't want to trigger casts of array
# equality checks to booleans.  We don't need to override comparisons because
# artists are already non-comparable.
Selection.__eq__ = lambda self, other: self is other
Selection.__ne__ = lambda self, other: self is not other
Selection.artist.__doc__ = (
    "The selected artist.")
Selection.target.__doc__ = (
    "The point picked within the artist, in data coordinates.")
Selection.dist.__doc__ = (
    "The distance from the click to the target, in pixels.")
Selection.annotation.__doc__ = (
    "The instantiated `matplotlib.text.Annotation`.")
Selection.extras.__doc__ = (
    "An additional list of artists (e.g., highlighters) that will be cleared "
    "at the same time as the annotation.")


@functools.singledispatch
[docs]def compute_pick(artist, event): """ Find whether *artist* has been picked by *event*. If it has, return the appropriate `Selection`; otherwise return ``None``. This is a single-dispatch function; implementations for various artist classes follow. """ warnings.warn(f"Pick support for {type(artist).__name__} is missing.")
class Index: def __init__(self, i, x, y): self.int = i self.x = x self.y = y def __floor__(self): return self.int def __ceil__(self): return self.int if max(self.x, self.y) == 0 else self.int + 1 # numpy<1.17 backcompat. floor = __floor__ ceil = __ceil__ def __format__(self, fmt): return f"{self.int}.(x={self.x:{fmt}}, y={self.y:{fmt}})" def __str__(self): return format(self, "") @classmethod def pre_index(cls, n_pts, index): i, frac = divmod(index, 1) i, odd = divmod(i, 2) x, y = (0, frac) if not odd else (frac, 1) return cls(i, x, y) @classmethod def post_index(cls, n_pts, index): i, frac = divmod(index, 1) i, odd = divmod(i, 2) x, y = (frac, 0) if not odd else (1, frac) return cls(i, x, y) @classmethod def mid_index(cls, n_pts, index): i, frac = divmod(index, 1) if i == 0: frac = .5 + frac / 2 elif i == 2 * n_pts - 2: # One less line than points. frac = frac / 2 quot, odd = divmod(i, 2) if not odd: if frac < .5: i = quot - 1 x, y = frac + .5, 1 else: i = quot x, y = frac - .5, 0 else: i = quot x, y = .5, frac return cls(i, x, y) def _compute_projection_pick(artist, path, xy): """ Project *xy* on *path* to obtain a `Selection` for *artist*. *path* is first transformed to screen coordinates using the artist transform, and the target of the returned `Selection` is transformed back to data coordinates using the artist *axes* inverse transform. The `Selection` `index` is returned as a float. This function returns ``None`` for degenerate inputs. The caller is responsible for converting the index to the proper class if needed. """ transform = artist.get_transform().frozen() tpath = (path.cleaned(transform) if transform.is_affine # `cleaned` only handles affine transforms. else transform.transform_path(path).cleaned()) # `cleaned` should return a path where the first element is `MOVETO`, the # following are `LINETO` or `CLOSEPOLY`, and the last one is `STOP`, i.e. # codes = path.codes # assert (codes[0], codes[-1]) == (path.MOVETO, path.STOP) # assert np.in1d(codes[1:-1], [path.LINETO, path.CLOSEPOLY]).all() vertices = tpath.vertices[:-1] codes = tpath.codes[:-1] vertices[codes == tpath.CLOSEPOLY] = vertices[0] # Unit vectors for each segment. us = vertices[1:] - vertices[:-1] ls = np.hypot(*us.T) with np.errstate(invalid="ignore"): # Results in 0/0 for repeated consecutive points. us /= ls[:, None] # Vectors from each vertex to the event (overwritten below). vs = xy - vertices[:-1] # Clipped dot products -- `einsum` cannot be done in place, `clip` can. # `clip` can trigger invalid comparisons if there are nan points. with np.errstate(invalid="ignore"): dot = np.clip(np.einsum("ij,ij->i", vs, us), 0, ls, out=vs[:, 0]) # Projections. projs = vertices[:-1] + dot[:, None] * us ds = np.hypot(*(xy - projs).T, out=vs[:, 1]) try: argmin = np.nanargmin(ds) dmin = ds[argmin] except (ValueError, IndexError): # See above re: exceptions caught. return else: target = AttrArray( artist.axes.transData.inverted().transform(projs[argmin])) target.index = ( (argmin + dot[argmin] / ls[argmin]) / (path._interpolation_steps / tpath._interpolation_steps)) return Selection(artist, target, dmin, None, None) def _untransform(orig_xy, screen_xy, ax): """ Return data coordinates to place an annotation at screen coordinates *screen_xy* in axes *ax*. *orig_xy* are the "original" coordinates as stored by the artist; they are transformed to *screen_xy* by whatever transform the artist uses. If the artist uses ``ax.transData``, just return *orig_xy*; else, apply ``ax.transData.inverse()`` to *screen_xy*. (The first case is more accurate than always applying ``ax.transData.inverse()``.) """ tr_xy = ax.transData.transform(orig_xy) return ( orig_xy if ((tr_xy == screen_xy) | np.isnan(tr_xy) & np.isnan(screen_xy)).all() else ax.transData.inverted().transform(screen_xy)) @compute_pick.register(Line2D) def _(artist, event): # No need to call `line.contains` as we're going to redo the work anyways # (also see matplotlib/matplotlib#6645, though that's fixed in mpl2.1). # Always work in screen coordinates, as this is how we need to compute # distances. Note that the artist transform may be different from the axes # transform (e.g., for axvline). xy = event.x, event.y data_xy = artist.get_xydata() data_screen_xy = artist.get_transform().transform(data_xy) sels = [] # If markers are visible, find the closest vertex. if artist.get_marker() not in ["None", "none", " ", "", None]: ds = np.hypot(*(xy - data_screen_xy).T) try: argmin = np.nanargmin(ds) dmin = ds[argmin] except (ValueError, IndexError): # numpy 1.7.0's `nanargmin([nan])` returns nan, so # `ds[argmin]` raises IndexError. In later versions of numpy, # `nanargmin([nan])` raises ValueError (the release notes for 1.8.0 # are incorrect on this topic). pass else: # More precise than transforming back. target = _with_attrs( _untransform( data_xy[argmin], data_screen_xy[argmin], artist.axes), index=argmin) sels.append(Selection(artist, target, dmin, None, None)) # If lines are visible, find the closest projection. if (artist.get_linestyle() not in ["None", "none", " ", "", None] and len(artist.get_xydata()) > 1): sel = _compute_projection_pick(artist, artist.get_path(), xy) if sel is not None: sel.target.index = { "_draw_lines": lambda _, index: index, "_draw_steps_pre": Index.pre_index, "_draw_steps_mid": Index.mid_index, "_draw_steps_post": Index.post_index}[ Line2D.drawStyles[artist.get_drawstyle()]]( len(data_xy), sel.target.index) sels.append(sel) sel = min(sels, key=lambda sel: sel.dist, default=None) return sel if sel and sel.dist < artist.get_pickradius() else None @compute_pick.register(PathPatch) @compute_pick.register(Polygon) @compute_pick.register(Rectangle) def _(artist, event): sel = _compute_projection_pick( artist, artist.get_path(), (event.x, event.y)) if sel and sel.dist < PATCH_PICKRADIUS: return sel @compute_pick.register(LineCollection) @compute_pick.register(PatchCollection) @compute_pick.register(PathCollection) def _(artist, event): # Use the C implementation to prune the list of segments. contains, info = artist.contains(event) if not contains: return offsets = artist.get_offsets() paths = artist.get_paths() if _is_scatter(artist): inds = info["ind"] offsets = artist.get_offsets()[inds] offsets_screen = artist.get_offset_transform().transform(offsets) ds = np.hypot(*(offsets_screen - [event.x, event.y]).T) argmin = ds.argmin() target = _with_attrs( _untransform(offsets[argmin], offsets_screen[argmin], artist.axes), index=inds[argmin]) return Selection(artist, target, ds[argmin], None, None) else: # Note that this won't select implicitly closed paths. sels = [ _compute_projection_pick( artist, Affine2D().translate(*offsets[ind % len(offsets)]) .transform_path(paths[ind % len(paths)]), (event.x, event.y)) for ind in info["ind"]] sel, index = min( ((sel, info["ind"][idx]) for idx, sel in enumerate(sels) if sel), key=lambda sel_idx: sel_idx[0].dist, default=(None, None)) if sel: sel = sel._replace(artist=artist) sel.target.index = (index, sel.target.index) if (isinstance(artist, PatchCollection) and sel.dist >= PATCH_PICKRADIUS): sel = None return sel @compute_pick.register(AxesImage) def _(artist, event): if type(artist) != AxesImage: # Skip and warn on subclasses (`NonUniformImage`, `PcolorImage`) as # they do not implement `contains` correctly. Even if they did, they # would not support moving as we do not know where a given index maps # back physically. return compute_pick.dispatch(object)(artist, event) contains, _ = artist.contains(event) if not contains: return ns = np.asarray(artist.get_array().shape[:2])[::-1] # (y, x) -> (x, y) xy = np.array([event.xdata, event.ydata]) xmin, xmax, ymin, ymax = artist.get_extent() # Handling of "upper" origin copied from AxesImage.get_cursor_data. if artist.origin == "upper": ymin, ymax = ymax, ymin low, high = np.array([[xmin, ymin], [xmax, ymax]]) idxs = ((xy - low) / (high - low) * ns).astype(int)[::-1] target = _with_attrs(xy, index=tuple(idxs)) return Selection(artist, target, 0, None, None) @compute_pick.register(Barbs) @compute_pick.register(Quiver) def _(artist, event): offsets = artist.get_offsets() offsets_screen = artist.get_offset_transform().transform(offsets) ds = np.hypot(*(offsets_screen - [event.x, event.y]).T) argmin = np.nanargmin(ds) if ds[argmin] < artist.get_pickradius(): target = _with_attrs( _untransform(offsets[argmin], offsets_screen[argmin], artist.axes), index=argmin) return Selection(artist, target, ds[argmin], None, None) else: return None @compute_pick.register(Text) def _(artist, event): return @compute_pick.register(ContainerArtist) def _(artist, event): return compute_pick(artist.container, event) @compute_pick.register(BarContainer) def _(container, event): try: (idx, patch), = { (idx, patch) for idx, patch in enumerate(container.patches) if patch.contains(event)[0]} except ValueError: return target = _with_attrs([event.xdata, event.ydata], index=idx) if patch.sticky_edges.x: target[0], = ( x for x in [patch.get_x(), patch.get_x() + patch.get_width()] if x not in patch.sticky_edges.x) if patch.sticky_edges.y: target[1], = ( y for y in [patch.get_y(), patch.get_y() + patch.get_height()] if y not in patch.sticky_edges.y) return Selection(container, target, 0, None, None) @compute_pick.register(ErrorbarContainer) def _(container, event): data_line, cap_lines, err_lcs = container sel_data = compute_pick(data_line, event) if data_line else None sel_err = min( filter(None, (compute_pick(err_lc, event) for err_lc in err_lcs)), key=lambda sel: sel.dist, default=None) if (sel_data and sel_data.dist < getattr(sel_err, "dist", np.inf)): return sel_data elif sel_err: idx, _ = sel_err.target.index if data_line: target = _with_attrs(data_line.get_xydata()[idx], index=idx) else: # We can't guess the original data in that case! return return Selection(container, target, 0, None, None) else: return @compute_pick.register(StemContainer) def _(container, event): sel = compute_pick(container.markerline, event) if sel: return sel idx_sel = min(filter(lambda idx_sel: idx_sel[1] is not None, ((idx, compute_pick(line, event)) for idx, line in enumerate(container.stemlines))), key=lambda idx_sel: idx_sel[1].dist, default=None) if idx_sel: idx, _ = idx_sel target = _with_attrs( container.stemlines[idx].get_xydata()[-1], index=idx) return Selection(container, target, 0, None, None) def _call_with_selection(func): """Decorator that passes a `Selection` built from the non-kwonly args.""" wrapped_kwonly_params = [ param for param in inspect.signature(func).parameters.values() if param.kind == param.KEYWORD_ONLY] sel_sig = inspect.signature(Selection) default_sel_sig = sel_sig.replace( parameters=[param.replace(default=None) if param.default is param.empty else param for param in sel_sig.parameters.values()]) @functools.wraps(func) def wrapper(*args, **kwargs): extra_kw = {param.name: kwargs.pop(param.name) for param in wrapped_kwonly_params if param.name in kwargs} ba = default_sel_sig.bind(*args, **kwargs) ba.apply_defaults() sel = Selection(*ba.args, **ba.kwargs) return func(sel, **extra_kw) wrapper.__signature__ = Signature( [*sel_sig.parameters.values(), *wrapped_kwonly_params]) return wrapper def _format_coord_unspaced(ax, xy): # Un-space-pad, remove empty coordinates from the output of # `format_{x,y}data`, and rejoin with newlines. return "\n".join( line for line, empty in zip( re.split(",? +", ax.format_coord(*xy)), itertools.chain(["x=", "y=", "z="], itertools.repeat(None))) if line != empty).rstrip() @functools.singledispatch @_call_with_selection
[docs]def get_ann_text(sel): """ Compute an annotating text for a `Selection` (passed **unpacked**). This is a single-dispatch function; implementations for various artist classes follow. """ warnings.warn( f"Annotation support for {type(sel.artist).__name__} is missing") return ""
def _strip_math(s): return cbook.strip_math(s) if len(s) >= 2 and s[0] == s[-1] == "$" else s def _format_scalarmappable_value(artist, idx): # matplotlib/matplotlib#12473. data = artist.get_array()[idx] if np.ndim(data) == 0: if not artist.colorbar: fig = Figure() ax = fig.subplots() artist.colorbar = fig.colorbar(artist, cax=ax) # This hack updates the ticks without actually paying the cost of # drawing (RendererBase.draw_path raises NotImplementedError). try: ax.yaxis.draw(RendererBase()) except NotImplementedError: pass fmt = artist.colorbar.formatter.format_data_short return "[" + _strip_math(fmt(data).strip()) + "]" else: return artist.format_cursor_data(data) # Includes brackets. @get_ann_text.register(Line2D) @get_ann_text.register(LineCollection) @get_ann_text.register(PatchCollection) @get_ann_text.register(PathCollection) @get_ann_text.register(Patch) @_call_with_selection def _(sel): artist = sel.artist label = artist.get_label() or "" text = _format_coord_unspaced(artist.axes, sel.target) if (_is_scatter(artist) # Heuristic: is the artist colormapped? # Note that this doesn't handle size-mapping (which is more likely # to involve an arbitrary scaling). and artist.get_array() is not None and len(artist.get_array()) == len(artist.get_offsets())): value = _format_scalarmappable_value(artist, sel.target.index) text = f"{text}\n{value}" if re.match("[^_]", label): text = f"{label}\n{text}" return text _Event = namedtuple("_Event", "xdata ydata") @get_ann_text.register(AxesImage) @_call_with_selection def _(sel): artist = sel.artist text = _format_coord_unspaced(artist.axes, sel.target) cursor_text = _format_scalarmappable_value(artist, sel.target.index) return f"{text}\n{cursor_text}" @get_ann_text.register(Barbs) @_call_with_selection def _(sel): artist = sel.artist text = "{}\n{}".format( _format_coord_unspaced(artist.axes, sel.target), (artist.u[sel.target.index], artist.v[sel.target.index])) return text @get_ann_text.register(Quiver) @_call_with_selection def _(sel): artist = sel.artist text = "{}\n{}".format( _format_coord_unspaced(artist.axes, sel.target), (artist.U[sel.target.index], artist.V[sel.target.index])) return text @get_ann_text.register(ContainerArtist) @_call_with_selection def _(sel): return get_ann_text(*sel._replace(artist=sel.artist.container)) @get_ann_text.register(BarContainer) @_call_with_selection def _(sel): return _format_coord_unspaced( _artist_in_container(sel.artist).axes, sel.target) @get_ann_text.register(ErrorbarContainer) @_call_with_selection def _(sel): data_line, cap_lines, err_lcs = sel.artist ann_text = get_ann_text(*sel._replace(artist=data_line)) if isinstance(sel.target.index, Integral): err_lcs = iter(err_lcs) for idx, (dir, has) in enumerate( zip("xy", [sel.artist.has_xerr, sel.artist.has_yerr])): if has: err = (next(err_lcs).get_paths()[sel.target.index].vertices - data_line.get_xydata()[sel.target.index])[:, idx] err_s = [getattr(_artist_in_container(sel.artist).axes, f"format_{dir}data")(e).rstrip() for e in err] # We'd normally want to check err.sum() == 0, but that can run # into fp inaccuracies. if len({s.lstrip("+-") for s in err_s}) == 1: repl = rf"\1=$\2\\pm{err_s[1]}$\3" else: err_s = [("+" if not s.startswith(("+", "-")) else "") + s for s in err_s] repl = r"\1=$\2_{%s}^{%s}$\3" % tuple(err_s) ann_text = re.sub(f"({dir})=(.*)(\n?)", repl, ann_text) return ann_text @get_ann_text.register(StemContainer) @_call_with_selection def _(sel): return get_ann_text(*sel._replace(artist=sel.artist.markerline)) @functools.singledispatch @_call_with_selection def move(sel, *, key): """ Move a `Selection` (passed **unpacked**) following a keypress. This function is used to implement annotation displacement through the keyboard. This is a single-dispatch function; implementations for various artist classes follow. """ return sel def _move_within_points(sel, xys, *, key): # Avoid infinite loop in case everything became nan at some point. for _ in range(len(xys)): if key == "left": new_idx = int(np.ceil(sel.target.index) - 1) % len(xys) elif key == "right": new_idx = int(np.floor(sel.target.index) + 1) % len(xys) else: return sel target = _with_attrs(xys[new_idx], index=new_idx) sel = sel._replace(target=target, dist=0) if np.isfinite(target).all(): return sel @move.register(Line2D) @_call_with_selection def _(sel, *, key): data_xy = sel.artist.get_xydata() return _move_within_points( sel, _untransform(data_xy, sel.artist.get_transform().transform(data_xy), sel.artist.axes), key=key) @move.register(PathCollection) @_call_with_selection def _(sel, *, key): if _is_scatter(sel.artist): offsets = sel.artist.get_offsets() return _move_within_points( sel, _untransform( offsets, sel.artist.get_offset_transform().transform(offsets), sel.artist.axes), key=key) else: return sel @move.register(AxesImage) @_call_with_selection def _(sel, *, key): ns = sel.artist.get_array().shape[:2] idxs = (np.asarray(sel.target.index) + {"left": [0, -1], "right": [0, 1], "up": {"lower": [1, 0], "upper": [-1, 0]}[sel.artist.origin], "down": {"lower": [-1, 0], "upper": [1, 0]}[sel.artist.origin]}[ key]) % ns xmin, xmax, ymin, ymax = sel.artist.get_extent() if sel.artist.origin == "upper": ymin, ymax = ymax, ymin low, high = np.array([[xmin, ymin], [xmax, ymax]]) target = _with_attrs(((idxs + .5) / ns)[::-1] * (high - low) + low, index=tuple(idxs)) return sel._replace(target=target) @move.register(ContainerArtist) @_call_with_selection def _(sel, *, key): return (move(*sel._replace(artist=sel.artist.container), key=key) ._replace(artist=sel.artist)) @move.register(ErrorbarContainer) @_call_with_selection def _(sel, *, key): data_line, cap_lines, err_lcs = sel.artist return _move_within_points(sel, data_line.get_xydata(), key=key) @functools.singledispatch @_call_with_selection
[docs]def make_highlight(sel, *, highlight_kwargs): """ Create a highlight for a `Selection`. This is a single-dispatch function; implementations for various artist classes follow. """ warnings.warn( f"Highlight support for {type(sel.artist).__name__} is missing")
def _set_valid_props(artist, kwargs): """Set valid properties for the artist, dropping the others.""" artist.set(**{k: kwargs[k] for k in kwargs if hasattr(artist, "set_" + k)}) return artist @make_highlight.register(Line2D) @_call_with_selection def _(sel, *, highlight_kwargs): hl = copy.copy(sel.artist) _set_valid_props(hl, highlight_kwargs) return hl @make_highlight.register(PathCollection) @_call_with_selection def _(sel, *, highlight_kwargs): hl = copy.copy(sel.artist) offsets = hl.get_offsets() hl.set_offsets(np.where( np.arange(len(offsets))[:, None] == sel.target.index, offsets, np.nan)) _set_valid_props(hl, highlight_kwargs) return hl