diff --git a/docs/api.rst b/docs/api.rst index a09ffa027..5e2489e2b 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -133,6 +133,16 @@ Demo functions :toctree: api +Text alignment +============== + +.. automodule:: ultraplot.textalign + :no-private-members: + +.. automodsumm:: ultraplot.textalign + :toctree: api + + Miscellaneous functions ======================= diff --git a/ultraplot/__init__.py b/ultraplot/__init__.py index 1aa62d0b0..0632ecfec 100644 --- a/ultraplot/__init__.py +++ b/ultraplot/__init__.py @@ -101,6 +101,7 @@ from .scale import SineLatitudeScale as SineLatitudeScale from .scale import SymmetricalLogScale as SymmetricalLogScale from .text import CurvedText as CurvedText + from .textalign import align_text as align_text from .ultralayout import ColorbarLayoutSolver as ColorbarLayoutSolver from .ultralayout import compute_ultra_positions as compute_ultra_positions from .ultralayout import get_grid_positions_ultra as get_grid_positions_ultra diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index f8373914f..f2e49929b 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -4,6 +4,7 @@ Implements basic shared functionality. """ +import contextlib import copy import inspect import re @@ -970,6 +971,10 @@ def __init__(self, *args, **kwargs): self._altx_parent = None # for cartesian axes only self._alty_parent = None self._colorbar_fill = None + self._align_texts = [] # texts registered for overlap avoidance + self._align_kwargs = {} # solver settings from auto_align_text() + self._align_arrows = [] # connectors drawn back to displaced anchors + self._align_cache = None # last solve, reused while its inputs hold self._inset_parent = None self._inset_bounds = None # for introspection ony self._inset_zoom = False @@ -3381,6 +3386,7 @@ def draw(self, renderer=None, *args, **kwargs): self._colorbar_fill.update_ticks(manual_only=True) # only if needed if self._inset_parent is not None and self._inset_zoom: self.indicate_inset_zoom() + self._apply_align_text(renderer) needs_inset_reflow = bool(getattr(self, "_inset_colorbar_needs_reflow", False)) has_inset_frame = bool( getattr(self, "_inset_colorbar_frame", None) is not None @@ -3418,6 +3424,7 @@ def get_tightbbox(self, renderer, *args, **kwargs): self._colorbar_fill.update_ticks(manual_only=True) # only if needed if self._inset_parent is not None and self._inset_zoom: self.indicate_inset_zoom() + self._apply_align_text(renderer) bbox = super().get_tightbbox(renderer, *args, **kwargs) fig = self.figure if ( @@ -4048,6 +4055,7 @@ def _curve_center(x, y, transform): def text( self, *args, + avoid_overlap=None, border=False, bbox=False, bordercolor="w", @@ -4075,6 +4083,10 @@ def text( Other parameters ---------------- + avoid_overlap : bool, default: :rc:`text.align` + Whether to automatically nudge this text at draw time so it does not + overlap other auto-aligned text or the plotted data. See + `~ultraplot.axes.Axes.auto_align_text` for the solver settings. border : bool, default: False Whether to draw border around text. borderwidth : float, default: 2 @@ -4104,6 +4116,7 @@ def text( See also -------- matplotlib.axes.Axes.text + ultraplot.axes.Axes.auto_align_text """ # Translate positional args # Audo-redirect to text2D for 3D axes if not enough arguments passed @@ -4172,8 +4185,127 @@ def text( "bboxpad": bboxpad, } ) + self._register_align_text(obj, avoid_overlap) + return obj + + def _register_align_text(self, obj, avoid_overlap=None): + """ + Queue a text object for draw-time overlap avoidance, or release it. + """ + # The solver moves a label by inverting its transform, which only means + # anything for text that its transform actually positions. A Text3D is + # placed by the 3D projection instead, and its stored position is the data + # coordinate that feeds it, so a display-space nudge would not move the + # label -- it would rewrite the data point. + if hasattr(obj, "get_position_3d"): + return obj + # An explicit avoid_overlap=False beats the rc setting, so opting out is + # always available. Releasing a label also puts it back where the user + # asked for it, rather than stranding it wherever the solver left it. + if _not_none(avoid_overlap, rc["text.align"]): + if obj not in self._align_texts: + self._align_texts.append(obj) + else: + if obj in self._align_texts: + self._align_texts.remove(obj) + from ..textalign import _reset_label + + if hasattr(obj, "_uplt_align_anchor"): + _reset_label(obj) + self.stale = True return obj + def auto_align_text(self, *objs, **kwargs): + """ + Automatically reposition text so that it does not overlap. + + Labels are relaxed away from each other, from the plotted data and from + the axes edges at draw time, then pulled back towards where you put them. + Because the solver runs on every draw, the layout stays valid when the + figure is resized or the data limits change. + + Parameters + ---------- + *objs : `~matplotlib.text.Text`, optional + The text or annotation objects to align. Default is every text + created with ``avoid_overlap=True`` plus, if none were, all the text + you added to the axes. + pad : float, default: :rc:`text.align.pad` + Padding in points kept around each label. + avoid_points : bool, default: True + Whether labels also repel the data points of lines and scatter plots. + avoid : sequence of `~matplotlib.artist.Artist`, optional + Extra artists whose bounding boxes the labels must stay clear of. + only_move : {'xy', 'x', 'y'}, default: 'xy' + Restrict movement to one axis. Use ``'y'`` when the horizontal + position of a label carries meaning, as on a time series. + max_iter : int, default: :rc:`text.align.maxiter` + Maximum number of relaxation iterations. + spring : float, default: 0.05 + Strength of the pull back towards the original position. Larger + values keep labels closer to their anchors at the cost of overlap. + step : float, default: 0.6 + Damping applied to each iteration's displacement. + clip : bool, default: True + Whether to keep labels inside the axes. + arrows : bool or dict, default: :rc:`text.align.arrows` + Whether to draw a connector from each displaced label back to the + point it labels. A dict is passed to `~matplotlib.patches.FancyArrowPatch`. + min_arrow_dist : float, default: 8.0 + Only draw connectors for labels displaced further than this, in points. + + Examples + -------- + >>> import ultraplot as uplt + >>> fig, ax = uplt.subplots() + >>> ax.scatter(x, y) + >>> for xi, yi, name in zip(x, y, names): + ... ax.text(xi, yi, name) + >>> ax.auto_align_text() + + See also + -------- + ultraplot.axes.Axes.text + ultraplot.axes.Axes.annotate + """ + objs = [obj for arg in objs for obj in (arg if np.iterable(arg) else (arg,))] + if objs: + self._align_texts = list(objs) + elif not self._align_texts: + # Titles and a-b-c labels live in self.texts but are positioned by the + # layout engine, so they are not ours to move. + titles = set(map(id, self._title_dict.values())) + self._align_texts = [ + t for t in self.texts if id(t) not in titles and t.get_text().strip() + ] + self._align_kwargs = kwargs + self.stale = True + return self._align_texts + + def _apply_align_text(self, renderer): + """ + Run the overlap solver for this axes (called on every draw). + """ + if not self._align_texts: + # Every label was released; take their connectors with them + for patch in self._align_arrows: + with contextlib.suppress(Exception): + patch.remove() + self._align_arrows = [] + return + from ..textalign import align_text + + kwargs = { + "pad": rc["text.align.pad"], + "max_iter": rc["text.align.maxiter"], + "arrows": rc["text.align.arrows"], + **self._align_kwargs, + } + try: + align_text(self, self._align_texts, renderer=renderer, **kwargs) + except Exception as err: # never let label polish break a render + warnings._warn_ultraplot(f"Failed to auto-align text: {err}") + @docstring._concatenate_inherited def annotate( self, @@ -4191,16 +4323,24 @@ def annotate( textcoords: Optional[Union[str, mtransforms.Transform]] = None, arrowprops: Optional[dict[str, Any]] = None, annotation_clip: Optional[bool] = None, + avoid_overlap: Optional[bool] = None, **kwargs: Any, ) -> Union[mtext.Annotation, "CurvedText"]: """ Add an annotation. If `xy` is a pair of 1D arrays, draw curved text. For curved input with `arrowprops`, the arrow points to the curve center. + + Parameters + ---------- + avoid_overlap : bool, default: :rc:`text.align` + Whether to automatically nudge this annotation at draw time so it + does not overlap other auto-aligned text or the plotted data. See + `~ultraplot.axes.Axes.auto_align_text`. """ curve_xy = self._coerce_curve_xy_from_xy_arg(xy) if curve_xy is None: - return super().annotate( + obj = super().annotate( text, xy=xy, xytext=xytext, @@ -4210,6 +4350,7 @@ def annotate( annotation_clip=annotation_clip, **kwargs, ) + return self._register_align_text(obj, avoid_overlap) x_curve, y_curve = curve_xy try: diff --git a/ultraplot/internals/rcsetup.py b/ultraplot/internals/rcsetup.py index eedb4ae38..18543a838 100644 --- a/ultraplot/internals/rcsetup.py +++ b/ultraplot/internals/rcsetup.py @@ -1126,6 +1126,28 @@ def _validator_accepts(validator, value): "Join style for text border strokes. Must be one of " "``'miter'``, ``'round'``, or ``'bevel'``.", ), + "text.align": ( + False, + _validate_bool, + "Whether text and annotations avoid overlapping each other and the data " + "by default. Set to ``True`` to opt every label into the solver used by " + "`~ultraplot.axes.Axes.auto_align_text`.", + ), + "text.align.pad": ( + 2.0, + _validate_pt, + "Padding in points kept around auto-aligned text.", + ), + "text.align.maxiter": ( + 60, + _validate_int, + "Maximum number of relaxation iterations used to auto-align text.", + ), + "text.align.arrows": ( + False, + _validate_bool, + "Whether auto-aligned text draws a connector back to the point it labels.", + ), "text.curved.upright": ( True, _validate_bool, diff --git a/ultraplot/tests/test_textalign.py b/ultraplot/tests/test_textalign.py new file mode 100644 index 000000000..132ca9118 --- /dev/null +++ b/ultraplot/tests/test_textalign.py @@ -0,0 +1,620 @@ +#!/usr/bin/env python3 +""" +Test automatic alignment of overlapping text and annotations. +""" + +import numpy as np +import pytest + +import ultraplot as uplt +from ultraplot.textalign import _label_bbox + + +def _count_overlaps(fig, texts) -> int: + # Measure the text box alone: an Annotation's window extent also covers its + # arrow, which legitimately crosses other labels. + renderer = fig._get_renderer() + bboxes = [_label_bbox(t, renderer, 0, 0) for t in texts] + return sum( + bboxes[i].overlaps(bboxes[j]) + for i in range(len(bboxes)) + for j in range(i + 1, len(bboxes)) + ) + + +@pytest.fixture +def crowded(): + """Points clustered tightly enough that their labels collide.""" + rng = np.random.default_rng(0) + x = np.concatenate([rng.normal(c, 0.12, 12) for c in (0.3, 0.7)]) + y = np.concatenate([rng.normal(c, 0.12, 12) for c in (0.4, 0.6)]) + return x, y, [f"sample {i}" for i in range(len(x))] + + +def test_overlapping_text_is_separated(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x, y, names)] + fig.canvas.draw() + assert _count_overlaps(fig, texts) > 0 + + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +def test_align_is_idempotent_across_draws(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + fig.canvas.draw() + first = np.array([t.get_position() for t in texts]) + fig.canvas.draw() + second = np.array([t.get_position() for t in texts]) + assert np.array_equal(first, second) + + +def test_align_survives_resize(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + fig.canvas.draw() + fig.set_size_inches(12, 4) + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +def test_auto_align_text_collects_untagged_text(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x, y, names)] + assert ax.auto_align_text() == texts + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +def test_only_move_preserves_other_axis(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + ax.auto_align_text(only_move="y", avoid_points=False) + fig.canvas.draw() + assert np.allclose([t.get_position()[0] for t in texts], x) + assert not np.allclose([t.get_position()[1] for t in texts], y) + + +def test_rc_opt_in(crowded): + x, y, names = crowded + with uplt.rc.context({"text.align": True}): + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x, y, names)] + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +def test_off_by_default(crowded): + """Nothing moves unless you ask for it.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x, y, names)] + fig.canvas.draw() + assert np.allclose([t.get_position() for t in texts], np.column_stack([x, y])) + assert _count_overlaps(fig, texts) > 0 + + +def test_per_label_opt_out_beats_rc(crowded): + """avoid_overlap=False wins even when the rc setting turns alignment on.""" + x, y, names = crowded + with uplt.rc.context({"text.align": True}): + fig, axs = uplt.subplots() + ax = axs[0] + kept = ax.text(x[0], y[0], names[0], fontsize=7, avoid_overlap=False) + moved = [ + ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x[1:], y[1:], names[1:]) + ] + fig.canvas.draw() + assert kept not in ax._align_texts + assert kept.get_position() == (x[0], y[0]) + assert any(t.get_position() != xy for t, xy in zip(moved, zip(x[1:], y[1:]))) + + +def test_opting_out_restores_position(crowded): + """Releasing a label puts it back where the user asked for it.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + text = ax.text(x[0], y[0], names[0], fontsize=7, avoid_overlap=True) + for xi, yi, nm in zip(x[1:], y[1:], names[1:]): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + fig.canvas.draw() + assert text.get_position() != (x[0], y[0]) # the solver moved it + ax._register_align_text(text, False) + assert text.get_position() == (x[0], y[0]) + + +def test_annotations_are_aligned(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + anns = [ + ax.annotate( + nm, + xy=(xi, yi), + xytext=(4, 4), + textcoords="offset points", + fontsize=7, + arrowprops=dict(arrowstyle="-"), + avoid_overlap=True, + ) + for xi, yi, nm in zip(x, y, names) + ] + fig.canvas.draw() + assert _count_overlaps(fig, anns) == 0 + # The arrow still points at the annotated data point + assert np.allclose([a.xy for a in anns], np.column_stack([x, y])) + + +def test_arrows_connect_displaced_labels(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + for xi, yi, nm in zip(x, y, names): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + ax.auto_align_text(arrows=True, min_arrow_dist=0) + fig.canvas.draw() + assert len(ax._align_arrows) > 0 + # Redrawing replaces the connectors rather than piling up new ones + count = len(ax._align_arrows) + fig.canvas.draw() + assert len(ax._align_arrows) == count + + +def test_arrows_do_not_disturb_data_limits(crowded): + """Connectors are decoration; they must not widen the axes.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + fig.canvas.draw() + before = (ax.get_xlim(), ax.get_ylim()) + for xi, yi, nm in zip(x, y, names): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + ax.auto_align_text(arrows=True, min_arrow_dist=0) + fig.canvas.draw() + assert (ax.get_xlim(), ax.get_ylim()) == before + + +def test_labels_stay_inside_axes(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + for xi, yi, nm in zip(x, y, names): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + fig.canvas.draw() + renderer = fig._get_renderer() + axes_bbox = ax.get_window_extent(renderer=renderer) + for text in ax._align_texts: + bbox = text.get_window_extent(renderer=renderer) + assert axes_bbox.x0 - 1 <= bbox.x0 and bbox.x1 <= axes_bbox.x1 + 1 + assert axes_bbox.y0 - 1 <= bbox.y0 and bbox.y1 <= axes_bbox.y1 + 1 + + +def test_invalid_only_move_raises(): + fig, axs = uplt.subplots() + ax = axs[0] + ax.text(0.5, 0.5, "a", avoid_overlap=True) + ax.text(0.5, 0.5, "b", avoid_overlap=True) + with pytest.raises(ValueError, match="only_move"): + uplt.align_text(ax, only_move="z") + + +def test_log_scale_labels(crowded): + """Movement is computed in pixels, so nonlinear transforms are fine.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.format(yscale="log") + texts = [ + ax.text(xi, 10**yi, nm, fontsize=7, avoid_overlap=True) + for xi, yi, nm in zip(x, y, names) + ] + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +# -------------------------------------------------------------------------- +# Obstacle types: labels repel the data, whichever artist drew it +# -------------------------------------------------------------------------- + + +def _labels_clear_of(fig, texts, points): + """No label box contains any of the given display-space points.""" + renderer = fig._get_renderer() + for text in texts: + bbox = _label_bbox(text, renderer, 0, 0) + for px, py in points: + if bbox.x0 < px < bbox.x1 and bbox.y0 < py < bbox.y1: + return False + return True + + +def test_labels_avoid_line_markers(): + """A Line2D contributes its vertices as obstacles, not just scatter.""" + x = np.linspace(0, 1, 12) + y = np.full_like(x, 0.5) + fig, axs = uplt.subplots() + ax = axs[0] + ax.plot(x, y, marker="o") + texts = [ + ax.text(xi, yi, f"p{i}", fontsize=7, avoid_overlap=True) + for i, (xi, yi) in enumerate(zip(x, y)) + ] + fig.canvas.draw() + pts = ax.transData.transform(np.column_stack([x, y])) + assert _labels_clear_of(fig, texts, pts) + + +def test_labels_avoid_line_collection(): + import matplotlib.collections as mcollections + + segs = [[(0.1, 0.5), (0.9, 0.5)], [(0.1, 0.6), (0.9, 0.6)]] + fig, axs = uplt.subplots() + ax = axs[0] + ax.add_collection(mcollections.LineCollection(segs)) + texts = [ + ax.text(0.5, 0.5, "a", fontsize=7, avoid_overlap=True), + ax.text(0.5, 0.6, "b", fontsize=7, avoid_overlap=True), + ] + fig.canvas.draw() + pts = ax.transData.transform(np.array([p for s in segs for p in s])) + assert _labels_clear_of(fig, texts, pts) + + +def test_hidden_artists_are_not_obstacles(): + """An invisible line should not push labels around.""" + x = np.linspace(0.2, 0.8, 8) + y = np.full_like(x, 0.5) + fig, axs = uplt.subplots() + ax = axs[0] + line = ax.plot(x, y, marker="o")[0] + line.set_visible(False) + text = ax.text(0.5, 0.5, "solo", fontsize=7, avoid_overlap=True) + fig.canvas.draw() + assert text.get_position() == (0.5, 0.5) # nothing to avoid, nothing moved + + +def test_avoid_points_false_ignores_data(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + text = ax.text(x[0], y[0], names[0], fontsize=7, avoid_overlap=True) + ax.auto_align_text(text, avoid_points=False) + fig.canvas.draw() + # The only label has nothing to collide with once the markers are ignored + assert text.get_position() == (x[0], y[0]) + + +# -------------------------------------------------------------------------- +# Solver options +# -------------------------------------------------------------------------- + + +def test_avoid_keeps_labels_off_a_given_artist(): + """The `avoid` argument treats an arbitrary artist as immovable.""" + fig, axs = uplt.subplots() + ax = axs[0] + blocker = ax.text(0.5, 0.5, "BLOCKER", fontsize=20) + text = ax.text(0.5, 0.5, "moved", fontsize=7, avoid_overlap=True) + ax.auto_align_text(text, avoid=[blocker], avoid_points=False) + fig.canvas.draw() + renderer = fig._get_renderer() + assert not _label_bbox(text, renderer, 0, 0).overlaps( + blocker.get_window_extent(renderer=renderer) + ) + + +def test_clip_false_changes_the_layout(crowded): + """Without clipping the solver is free to push labels past the axes edge.""" + x, y, names = crowded + + def solve(clip): + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) + for *xy, nm in zip(x, y, names) + ] + ax.auto_align_text(clip=clip) + fig.canvas.draw() + return np.array([t.get_position() for t in texts]) + + clipped, free = solve(True), solve(False) + assert not np.allclose(clipped, free) + + +def test_only_move_x_preserves_y(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + ax.auto_align_text(only_move="x", avoid_points=False) + fig.canvas.draw() + assert np.allclose([t.get_position()[1] for t in texts], y) + assert not np.allclose([t.get_position()[0] for t in texts], x) + + +def test_many_labels_use_the_reduced_schedule(crowded): + """Past the restart limit the solver still runs and still separates labels.""" + from ultraplot.textalign import _RESTART_LIMIT + + n = _RESTART_LIMIT + 10 + rng = np.random.default_rng(3) + x, y = rng.random(n), rng.random(n) + fig, axs = uplt.subplots(refwidth=8) + ax = axs[0] + texts = [ + ax.text(xi, yi, f"p{i}", fontsize=5, avoid_overlap=True) + for i, (xi, yi) in enumerate(zip(x, y)) + ] + before = _count_overlaps(fig, texts) + fig.canvas.draw() + assert _count_overlaps(fig, texts) < before + + +# -------------------------------------------------------------------------- +# Connectors +# -------------------------------------------------------------------------- + + +def test_arrow_style_dict_is_forwarded(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + for xi, yi, nm in zip(x, y, names): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + ax.auto_align_text(arrows=dict(color="red", linewidth=2.0), min_arrow_dist=0) + fig.canvas.draw() + assert ax._align_arrows + patch = ax._align_arrows[0] + assert patch.get_linewidth() == 2.0 + assert uplt.colors.to_hex(patch.get_edgecolor()) == uplt.colors.to_hex("red") + + +def test_no_connector_for_annotations_with_their_own_arrow(crowded): + """matplotlib already draws that arrow; we must not double it.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + for xi, yi, nm in zip(x, y, names): + ax.annotate( + nm, + xy=(xi, yi), + xytext=(4, 4), + textcoords="offset points", + fontsize=7, + arrowprops=dict(arrowstyle="-"), + avoid_overlap=True, + ) + ax.auto_align_text(arrows=True, min_arrow_dist=0) + fig.canvas.draw() + assert ax._align_arrows == [] + + +def test_connector_points_at_the_annotated_point(crowded): + """An annotation without arrowprops still gets a connector to its xy.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + for xi, yi, nm in zip(x, y, names): + ax.annotate( + nm, + xy=(xi, yi), + xytext=(4, 4), + textcoords="offset points", + fontsize=7, + avoid_overlap=True, + ) + ax.auto_align_text(arrows=True, min_arrow_dist=0) + fig.canvas.draw() + assert ax._align_arrows + targets = {tuple(np.round(p.get_path().vertices[-1], 6)) for p in ax._align_arrows} + assert targets & {tuple(np.round(xy, 6)) for xy in zip(x, y)} + + +def test_min_arrow_dist_suppresses_short_connectors(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + for xi, yi, nm in zip(x, y, names): + ax.text(xi, yi, nm, fontsize=7, avoid_overlap=True) + ax.auto_align_text(arrows=True, min_arrow_dist=10_000) + fig.canvas.draw() + assert ax._align_arrows == [] + + +# -------------------------------------------------------------------------- +# Edge cases +# -------------------------------------------------------------------------- + + +def test_no_labels_is_a_noop(): + fig, axs = uplt.subplots() + ax = axs[0] + assert uplt.align_text(ax) == [] + + +def test_blank_and_hidden_labels_are_skipped(): + fig, axs = uplt.subplots() + ax = axs[0] + blank = ax.text(0.5, 0.5, " ", avoid_overlap=True) + hidden = ax.text(0.5, 0.5, "hidden", avoid_overlap=True) + hidden.set_visible(False) + fig.canvas.draw() + assert blank.get_position() == (0.5, 0.5) + assert hidden.get_position() == (0.5, 0.5) + + +def test_single_label_never_moves(): + fig, axs = uplt.subplots() + ax = axs[0] + text = ax.text(0.5, 0.5, "alone", avoid_overlap=True) + fig.canvas.draw() + assert text.get_position() == (0.5, 0.5) + + +def test_align_text_without_a_renderer(crowded): + """Called directly, the solver fetches its own renderer.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + uplt.align_text(ax) # no renderer= passed + assert _count_overlaps(fig, texts) == 0 + + +def test_solver_failure_warns_but_still_draws(crowded, monkeypatch): + """A broken solve must never take the whole figure down with it.""" + import ultraplot.textalign as textalign + from ultraplot.internals.warnings import UltraPlotWarning + + def boom(*args, **kwargs): + raise RuntimeError("solver exploded") + + monkeypatch.setattr(textalign, "align_text", boom) + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.text(x[0], y[0], names[0], avoid_overlap=True) + with pytest.warns(UltraPlotWarning, match="auto-align"): + fig.canvas.draw() + assert ax.texts # the label is still there, just unaligned + + +def test_auto_align_text_accepts_an_iterable(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + texts = [ax.text(*xy, nm, fontsize=7) for *xy, nm in zip(x, y, names)] + assert ax.auto_align_text(texts) == texts # a list, not *args + fig.canvas.draw() + assert _count_overlaps(fig, texts) == 0 + + +def test_moving_a_label_by_hand_re_anchors_it(crowded): + """set_position() after a draw must not be undone by the next one.""" + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + fig.canvas.draw() + text = texts[0] + solved = np.array(text.get_position()) # where the solver had parked it + text.set_position((0.05, 0.95)) + fig.canvas.draw() + + assert text._uplt_align_anchor == (0.05, 0.95) + # The solver may still nudge it, but it now relaxes away from where the user + # put it rather than from the anchor the user has abandoned + now = np.array(text.get_position()) + assert np.hypot(*(now - [0.05, 0.95])) < np.hypot(*(now - solved)) + assert _count_overlaps(fig, texts) == 0 + + +def test_three_axes_text_is_left_alone(): + """A Text3D is placed by the projection, so a display-space nudge is meaningless.""" + fig = uplt.figure() + ax = fig.subplot(proj="3d") + text = ax.text(0.5, 0.5, 0.5, "label", avoid_overlap=True) + other = ax.text(0.5, 0.5, 0.5, "overlapping", avoid_overlap=True) + fig.canvas.draw() + assert ax._align_texts == [] + assert text.get_position() == (0.5, 0.5) + assert other.get_position() == (0.5, 0.5) + + +def test_unchanged_redraw_reuses_the_solve(crowded, monkeypatch): + """A redraw that changes nothing must not pay for the relaxation again.""" + import ultraplot.textalign as textalign + + calls = [] + seed = textalign._crowd_seed # only reached on a full solve, never on a cache hit + monkeypatch.setattr( + textalign, + "_crowd_seed", + lambda *args, **kwargs: (calls.append(1), seed(*args, **kwargs))[1], + ) + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + fig.canvas.draw() + assert calls # solved at least once + before = np.array([t.get_position() for t in texts]) + + calls.clear() + fig.canvas.draw() + assert not calls # ... and reused it, rather than solving again + assert np.array_equal([t.get_position() for t in texts], before) + + # A resize changes every box in display space, so the cache must not survive it + fig.set_size_inches(12, 4) + fig.canvas.draw() + assert calls + assert _count_overlaps(fig, texts) == 0 + + +def test_releasing_every_label_removes_the_connectors(crowded): + x, y, names = crowded + fig, axs = uplt.subplots() + ax = axs[0] + ax.scatter(x, y) + texts = [ + ax.text(*xy, nm, fontsize=7, avoid_overlap=True) for *xy, nm in zip(x, y, names) + ] + ax.auto_align_text(arrows=True, min_arrow_dist=0) + fig.canvas.draw() + patches = list(ax._align_arrows) + assert patches + for text in texts: + ax._register_align_text(text, False) + fig.canvas.draw() + assert ax._align_arrows == [] + assert not any(p in ax.artists for p in patches) diff --git a/ultraplot/textalign.py b/ultraplot/textalign.py new file mode 100644 index 000000000..a87a47fd6 --- /dev/null +++ b/ultraplot/textalign.py @@ -0,0 +1,640 @@ +#!/usr/bin/env python3 +""" +Automatic repositioning of text and annotation boxes so they do not overlap. + +The solver works entirely in display (pixel) space, so it is agnostic to the +transform attached to each label -- data, axes, log-scaled, polar and +geographic labels are all handled the same way. Labels are reset to their +original anchors before every pass, which keeps the result stable across +repeated draws, resizes and dpi changes. +""" + +from __future__ import annotations + +from typing import Iterable, Optional, Sequence + +import matplotlib.collections as mcollections +import matplotlib.lines as mlines +import matplotlib.patches as mpatches +import matplotlib.text as mtext +import matplotlib.transforms as mtransforms +import numpy as np + +__all__ = ["align_text"] + +#: Above this many labels, fall back to a single relaxation schedule (see below). +_RESTART_LIMIT = 120 + + +def _points_to_pixels(fig, value) -> float: + return float(value) * fig.dpi / 72.0 + + +def _expand_bbox(bbox, padx: float, pady: float): + return mtransforms.Bbox.from_extents( + bbox.x0 - padx, bbox.y0 - pady, bbox.x1 + padx, bbox.y1 + pady + ) + + +def _fits(box, sx, sy, bounds) -> bool: + """ + Whether shifting ``box`` (x0, y0, x1, y1) by half of (sx, sy) keeps it inside + ``bounds`` -- half, because a box only takes half of a pairwise push. + """ + if bounds is None: + return True + dx, dy = 0.5 * sx, 0.5 * sy + return ( + box[0] + dx >= bounds[0] - 1e-6 + and box[2] + dx <= bounds[2] + 1e-6 + and box[1] + dy >= bounds[1] - 1e-6 + and box[3] + dy <= bounds[3] + 1e-6 + ) + + +def _overlap_shift(b1, b2, bounds=None): + """ + Return the translation ``(sx, sy)`` separating box ``b1`` from ``b2``, or zeros. + + Boxes are ``(x0, y0, x1, y1)``. Both axis-aligned escape routes are considered + and the shallower one wins, because every pixel a label moves is a pixel further + from the thing it describes. If that route would push the label out of + ``bounds`` the other one is used instead -- without this, a stack of labels + jammed against the edge of the axes gets shoved straight back into the wall on + every iteration and can never spread out sideways. + + This runs once per colliding pair per iteration, so it deals in plain floats: + building a two-element numpy array here costs more than all the arithmetic. + """ + dx = min(b1[2], b2[2]) - max(b1[0], b2[0]) + if dx <= 0: + return 0.0, 0.0 + dy = min(b1[3], b2[3]) - max(b1[1], b2[1]) + if dy <= 0: + return 0.0, 0.0 + sx = dx if (b1[0] + b1[2]) >= (b2[0] + b2[2]) else -dx + sy = dy if (b1[1] + b1[3]) >= (b2[1] + b2[3]) else -dy + if dx < dy: + first, second = (sx, 0.0), (0.0, sy) + else: + first, second = (0.0, sy), (sx, 0.0) + if _fits(b1, first[0], first[1], bounds) or not _fits( + b1, second[0], second[1], bounds + ): + return first + return second + + +def _point_shift(box, px, py): + """ + Return the shortest translation ``(sx, sy)`` that moves ``box`` off a point. + """ + left = px - box[0] + right = box[2] - px + down = py - box[1] + up = box[3] - py + best = min(left, right, down, up) + if best == left: + return left, 0.0 + if best == right: + return -right, 0.0 + if best == down: + return 0.0, down + return 0.0, -up + + +def _colliding_pairs(boxes, live): + """ + Indices (i, j), i < j, of every pair of live boxes that currently overlap. + + One vectorised sweep. A k-d tree radius query is asymptotically better and is + genuinely faster at this step in isolation, but it does not pay for itself + here: even at 800 labels this sweep is a low single-digit percentage of the + solve, which is dominated by resolving the collisions it finds. It is not worth + a SciPy dependency to speed up something that is not the bottleneck. + """ + x0, y0, x1, y1 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] + hit = ( + (x0[:, None] < x1[None, :]) + & (x0[None, :] < x1[:, None]) + & (y0[:, None] < y1[None, :]) + & (y0[None, :] < y1[:, None]) + ) + hit &= live[:, None] & live[None, :] + return np.argwhere(np.triu(hit, k=1)) + + +def _count_overlaps(boxes) -> int: + """ + Number of label pairs that visibly overlap. Takes the raw boxes, without the + padding cushion: the cushion is a solver knob, but the score must be judged on + the boxes the reader actually sees. + """ + if len(boxes) < 2: + return 0 + x0, y0, x1, y1 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] + inter_x = np.minimum(x1[:, None], x1[None, :]) - np.maximum( + x0[:, None], x0[None, :] + ) + inter_y = np.minimum(y1[:, None], y1[None, :]) - np.maximum( + y0[:, None], y0[None, :] + ) + hits = (inter_x > 0) & (inter_y > 0) + return int(np.count_nonzero(np.triu(hits, k=1))) + + +def _crowd_seed(anchors, obstacles, sizes) -> np.ndarray: + """ + Initial offsets that push every label away from its local crowd. + + Relaxing from the original positions is a purely local search, so a dense + cluster can settle into a knot that no amount of further iteration undoes. + Starting from a pre-exploded configuration puts the solver in a different + basin, which is what the restarts are for. + """ + crowd = anchors if not len(obstacles) else np.vstack([anchors, obstacles]) + seed = np.zeros_like(anchors) + for i, point in enumerate(anchors): + delta = point - crowd # vectors pointing away from each neighbour + dist = np.hypot(delta[:, 0], delta[:, 1]) + near = (dist > 1e-9) & (dist < 3 * sizes[i]) + if not np.any(near): + continue + # Weight nearer neighbours more heavily, then step out along the escape + direction = np.sum(delta[near] / dist[near, None] ** 2, axis=0) + norm = np.hypot(*direction) + if norm > 1e-9: + seed[i] = direction / norm * sizes[i] + return seed + + +def _artist_points(artist) -> Optional[np.ndarray]: + """ + Sample the display-space points an artist occupies, or None if unsupported. + """ + if isinstance(artist, mlines.Line2D): + data = artist.get_xydata() + if data is None or len(data) == 0: + return None + return artist.get_transform().transform(np.asarray(data, dtype=float)) + if isinstance(artist, mcollections.PathCollection): + offsets = artist.get_offsets() + if offsets is None or len(offsets) == 0: + return None + return artist.get_offset_transform().transform(np.asarray(offsets, dtype=float)) + if isinstance(artist, mcollections.LineCollection): + segs = artist.get_segments() + if not segs: + return None + pts = np.concatenate([np.asarray(s, dtype=float) for s in segs]) + return artist.get_transform().transform(pts) + return None + + +def _gather_obstacles(ax, labels, avoid_points: bool): + """ + Collect display-space points (data markers/vertices) that labels avoid. + """ + if not avoid_points: + return np.empty((0, 2)) + points = [] + for artist in (*ax.lines, *ax.collections): + if not artist.get_visible() or artist in labels: + continue + pts = _artist_points(artist) + if pts is not None and len(pts): + points.append(pts) + if not points: + return np.empty((0, 2)) + points = np.concatenate(points) + finite = np.all(np.isfinite(points), axis=1) + return points[finite] + + +def _label_bbox(label, renderer, padx, pady): + try: + if isinstance(label, mtext.Annotation): + # Annotation.get_window_extent() unions the text with its arrow, and + # the arrow only grows as the label moves away. Measure the box alone. + label._renderer = renderer + label.update_positions(renderer) + bbox = mtext.Text.get_window_extent(label) + else: + bbox = label.get_window_extent(renderer=renderer) + except Exception: + return None + if not np.all(np.isfinite(bbox.get_points())): + return None + return _expand_bbox(bbox, padx, pady) + + +def _position(label) -> tuple: + """ + Where the label currently sits, in its own coordinate system. + """ + if isinstance(label, mtext.Annotation): + return tuple(map(float, label.xyann)) + return tuple(map(float, label.get_position())) + + +def _place(label, point) -> None: + """ + Move a label to ``point`` and record that we are the ones who put it there. + """ + if isinstance(label, mtext.Annotation): + label.xyann = tuple(point) + else: + label.set_position(tuple(point)) + label._uplt_align_placed = _position(label) + + +def _reset_label(label) -> np.ndarray: + """ + Restore a label to its user-specified anchor and return that anchor. + + The anchor is cached in the label's own coordinate system the first time we + see it, so re-running the solver on every draw is idempotent rather than + cumulative. A label that is not where we last left it has been repositioned by + the user since the previous solve, and that new position becomes the anchor -- + otherwise ``set_position`` on an aligned label would appear to do nothing, + with the next draw quietly dragging it back to an anchor the user has + abandoned. + """ + current = _position(label) + if current != getattr(label, "_uplt_align_placed", None): + label._uplt_align_anchor = current + anchor = label._uplt_align_anchor + _place(label, anchor) + return np.asarray(anchor, dtype=float) + + +def _text_transform(label, renderer): + """ + Return the transform mapping a label's stored position to display space. + """ + if isinstance(label, mtext.Annotation): + return label._get_xy_transform(renderer, label.anncoords) + return label.get_transform() + + +def _move_label(label, delta_display, anchor, transform) -> None: + """ + Offset a label by ``delta_display`` pixels from its anchor. + """ + if not np.any(delta_display): + return + anchor_display = transform.transform(anchor) + target = anchor_display + delta_display + try: + new = transform.inverted().transform(target) + except Exception: + return + if not np.all(np.isfinite(new)): + return + _place(label, new) + + +def _target_display(label, renderer): + """ + Display-space point the label refers to (the annotated point, or its anchor). + """ + if isinstance(label, mtext.Annotation): + trans = label._get_xy_transform(renderer, label.xycoords) + return np.asarray(trans.transform(label.xy), dtype=float) + return np.asarray( + label.get_transform().transform(label._uplt_align_anchor), dtype=float + ) + + +def align_text( + ax, + labels: Optional[Sequence[mtext.Text]] = None, + *, + renderer=None, + pad: float = 2.0, + avoid_points: bool = True, + avoid: Iterable = (), + only_move: str = "xy", + max_iter: int = 60, + spring: float = 0.05, + step: float = 0.6, + clip: bool = True, + arrows: bool | dict = False, + min_arrow_dist: float = 8.0, +) -> list: + """ + Nudge text objects until they no longer overlap each other or the data. + + Parameters + ---------- + ax : ultraplot.axes.Axes + The axes whose labels are aligned. + labels : sequence of `~matplotlib.text.Text`, optional + The labels to move. Default is every text registered for alignment on + the axes (see `~ultraplot.axes.Axes.text` with ``avoid_overlap=True``). + pad : float, default: 2.0 + Padding in points added around every label bounding box. + avoid_points : bool, default: True + Whether labels also repel the data points of lines and scatter plots. + avoid : sequence of `~matplotlib.artist.Artist`, optional + Additional artists (a legend, an inset, ...) whose bounding boxes the + labels must stay clear of. + only_move : {'xy', 'x', 'y'}, default: 'xy' + Restrict movement to a single axis. Useful when the horizontal position + of a label is meaningful, e.g. labels on a time series. + max_iter : int, default: 60 + Maximum number of relaxation iterations. + spring : float, default: 0.05 + Strength of the pull back towards the original anchor. Larger values + keep labels closer to where they were placed at the cost of overlap. + step : float, default: 0.6 + Damping applied to each iteration's displacement. + clip : bool, default: True + Whether to keep labels inside the axes. + arrows : bool or dict, default: False + Whether to draw a connector from displaced labels back to their anchor. + A dict is passed to `~matplotlib.patches.FancyArrowPatch`. + min_arrow_dist : float, default: 8.0 + Only draw connectors for labels displaced further than this (in points). + + Returns + ------- + list + The connector patches that were drawn, if any. + """ + if only_move not in ("xy", "x", "y"): + raise ValueError(f"Invalid only_move={only_move!r}. Choose 'xy', 'x' or 'y'.") + mask = np.array([float(only_move != "y"), float(only_move != "x")]) + + fig = ax.figure + if fig is None: + return [] + if renderer is None: + renderer = fig._get_renderer() + if labels is None: + labels = list(getattr(ax, "_align_texts", ())) + labels = [t for t in labels if t.get_visible() and t.get_text().strip()] + if not labels: + return [] + + padx = pady = _points_to_pixels(fig, pad) + + # Reset to anchors first so repeated draws converge to the same layout + anchors = [_reset_label(t) for t in labels] + transforms = [_text_transform(t, renderer) for t in labels] + + obstacles = _gather_obstacles(ax, labels, avoid_points) + static_bboxes = [] + for artist in avoid: + try: + static_bboxes.append(artist.get_window_extent(renderer=renderer)) + except Exception: + continue + + axes_bbox = ax.get_window_extent(renderer=renderer) if clip else None + + # Measure each label exactly once. Translating a text does not change the size + # of its bounding box, so every later box is just this one plus an offset -- + # which keeps the whole relaxation in numpy and off the renderer. + raw = [] + for label in labels: + bbox = _label_bbox(label, renderer, 0, 0) + raw.append( + (np.nan, np.nan, np.nan, np.nan) + if bbox is None + else (bbox.x0, bbox.y0, bbox.x1, bbox.y1) + ) + raw = np.array(raw, dtype=float) + live = np.all(np.isfinite(raw), axis=1) + order = [i for i in range(len(labels)) if live[i]] + + statics = np.array( + [(bb.x0, bb.y0, bb.x1, bb.y1) for bb in static_bboxes], dtype=float + ).reshape(-1, 4) + + obs_x = obstacles[:, 0] if len(obstacles) else np.empty(0) + obs_y = obstacles[:, 1] if len(obstacles) else np.empty(0) + + bounds = ( + None + if axes_bbox is None + else np.array([axes_bbox.x0, axes_bbox.y0, axes_bbox.x1, axes_bbox.y1]) + ) + + def _solve(step, spring, seed=None, cushion=1.0): + """One relaxation run; returns (score, offsets).""" + base = raw + cushion * np.array([-padx, -pady, padx, pady]) + offsets = np.zeros((len(labels), 2)) if seed is None else seed * mask + best_offsets = offsets.copy() + best_score = (np.inf, np.inf, np.inf, np.inf) + + for it in range(max_iter): + boxes = base + np.hstack([offsets, offsets]) + + # Score the layout that `offsets` actually produces, on the true boxes, + # and before the sweep below starts mutating `boxes` in place. Scoring + # afterwards would rank the provisional fully-resolved state while + # saving the damped offsets that do not correspond to it. + overlaps = _count_overlaps((raw + np.hstack([offsets, offsets]))[live]) + + # Resolve collisions one at a time, each push applied immediately so that + # later pairs see the updated box (Gauss-Seidel). Accumulating all the + # pushes and applying them together lets a label squeezed between two + # neighbours take equal and opposite shifts that cancel, deadlocking the + # stack instead of expanding it. + # + # This sweep is the hot loop of the whole module, so it runs on plain + # Python floats. Indexing a row of a numpy array and adding a length-two + # array to it costs several times the arithmetic it performs. + work = boxes.tolist() + fx = [0.0] * len(labels) + fy = [0.0] * len(labels) + px_ = [0.0] * len(labels) + py_ = [0.0] * len(labels) + + def _push(i, sx, sy): + if sx or sy: + fx[i] += sx + fy[i] += sy + box = work[i] + box[0] += sx + box[1] += sy + box[2] += sx + box[3] += sy + + # Points and static artists cannot move, so settle the labels against + # them first and give label-on-label collisions the last word: a label + # shoved off a marker must not come to rest on top of a neighbour. + for i in order: + for k in range(len(statics)): + _push(i, *_overlap_shift(work[i], statics[k], bounds)) + if len(obstacles): + b = work[i] + inside = ( + (obs_x > b[0]) + & (obs_x < b[2]) + & (obs_y > b[1]) + & (obs_y < b[3]) + ) + for ox, oy in zip(obs_x[inside], obs_y[inside]): + _push(i, *_point_shift(work[i], ox, oy)) + + # Look for label-on-label collisions only now, so the pairs reflect where + # the pushes above actually left the labels. One vectorised sweep finds + # them; the loop then resolves just those. Pairs that a push creates + # mid-sweep are picked up on the next iteration. + for i, j in _colliding_pairs(np.array(work), live): + sx, sy = _overlap_shift(work[i], work[j], bounds) + if not (sx or sy): + continue # an earlier push in this sweep already separated them + hx, hy = 0.5 * sx, 0.5 * sy + px_[i] += hx + py_[i] += hy + px_[j] -= hx + py_[j] -= hy + _push(i, hx, hy) + _push(j, -hx, -hy) + + forces = np.column_stack([fx, fy]) + pair_forces = np.column_stack([px_, py_]) + # Whatever is left over came from the markers and the `avoid` artists + obstacle_forces = forces - pair_forces + + # Rank on the overlap count taken above: what the eye objects to first is + # labels sitting on each other. Sitting on a marker is the next cost, and + # it has to be scored -- a lone label parked on top of a marker overlaps + # no other label, so without this term it would score a perfect zero on + # the very first iteration and never move. Distance from the anchor only + # breaks ties, so labels stay put when nothing is in their way. + crowding = float(np.sum(np.hypot(pair_forces[:, 0], pair_forces[:, 1]))) + blocked = float( + np.sum(np.hypot(obstacle_forces[:, 0], obstacle_forces[:, 1])) + ) + displacement = float(np.sum(np.hypot(offsets[:, 0], offsets[:, 1]))) + score = (overlaps, crowding, blocked, displacement) + if score < best_score: + best_score = score + best_offsets = offsets.copy() + if not forces.any(): + break + + # Pull labels that are already free back towards their anchor. Springing + # a colliding label would merely balance the repulsion and leave the + # overlap in place. Fade the spring out so late iterations prioritise + # separation over staying close to the anchor. + free = ~np.any(forces, axis=1) + forces[free] -= spring * (1 - it / max_iter) * offsets[free] + forces *= mask + offsets = offsets + step * forces + + if bounds is not None: + boxes = base + np.hstack([offsets, offsets]) + lo = bounds[:2] - boxes[:, :2] + hi = bounds[2:] - boxes[:, 2:] + nudge = np.clip(0.0, np.minimum(lo, hi), np.maximum(lo, hi)) + offsets += np.where(live[:, None], nudge * mask, 0.0) + + return best_score, best_offsets + + # The solve is by far the most expensive part of a draw, and it is a pure + # function of the geometry gathered above. A figure is redrawn far more often + # than it changes -- the tight layout measures it, a legend is toggled, a + # notebook cell re-renders it -- so remember the answer and only pay for it + # again when one of the inputs actually moves. Every input is in display + # space, so a resize, a dpi change or new data limits all invalidate this on + # their own. The obstacles go in as a hash: a dense line contributes tens of + # thousands of points, and the key must not pin a copy of them all in memory. + key = ( + tuple(map(id, labels)), + raw.tobytes(), + statics.tobytes(), + None if bounds is None else bounds.tobytes(), + hash(obstacles.tobytes()), + padx, + pady, + only_move, + max_iter, + spring, + step, + ) + cached = getattr(ax, "_align_cache", None) + if cached is not None and cached[0] == key: + offsets = cached[1] + else: + # No single relaxation schedule wins on every layout: a decisive step clears + # dense stacks but can overshoot into a worse arrangement, a gentle one does + # the reverse, and relaxing from the original positions cannot always escape + # a knotted cluster at all. The solve is cheap now that nothing re-measures, + # so run a handful and keep the best. The schedules are fixed rather than + # random, so the same figure always lands on the same layout. + centres = 0.5 * (raw[:, :2] + raw[:, 2:]) + sizes = 0.5 * np.hypot(raw[:, 2] - raw[:, 0], raw[:, 3] - raw[:, 1]) + seed = _crowd_seed(np.nan_to_num(centres), obstacles, np.nan_to_num(sizes)) + + schedules = [ + (step, spring, None, 1.0), + (1.0, spring, None, 1.0), + (step, spring, seed, 1.0), + (1.0, 0.0, seed, 1.0), + (step, spring, None, 1.5), + (step, spring, None, 0.5), + (0.4, 0.5 * spring, None, 1.0), + ] + # Each restart costs a full relaxation. Past a few hundred labels they cannot + # all fit anyway, so the restarts buy nothing but seconds of redraw time. + if len(order) > _RESTART_LIMIT: + schedules = schedules[:2] + + offsets = None + best = (np.inf, np.inf, np.inf, np.inf) + for step_i, spring_i, seed_i, cushion_i in schedules: + score, candidate = _solve(step_i, spring_i, seed_i, cushion_i) + if score < best: + best, offsets = score, candidate + if best[0] == 0 and best[2] == 0: + break # nothing overlaps and nothing sits on a marker: we are done + ax._align_cache = (key, offsets) + + # Apply the final offsets + for label, anchor, trans, offset in zip(labels, anchors, transforms, offsets): + _move_label(label, offset, anchor, trans) + + # Connectors from displaced labels back to what they describe + patches = [] + for patch in getattr(ax, "_align_arrows", ()): + try: + patch.remove() + except Exception: + pass + ax._align_arrows = patches + if not arrows: + return patches + props = dict(arrowstyle="-", color="0.5", linewidth=0.8, shrinkA=2, shrinkB=2) + if isinstance(arrows, dict): + props.update(arrows) + threshold = _points_to_pixels(fig, min_arrow_dist) + inv = ax.transData.inverted() + for label, offset in zip(labels, offsets): + if np.hypot(*offset) < threshold: + continue + if isinstance(label, mtext.Annotation) and label.arrowprops: + continue # matplotlib already draws this one + target = _target_display(label, renderer) + bbox = _label_bbox(label, renderer, 0, 0) + if bbox is None: + continue + center = np.array([0.5 * (bbox.x0 + bbox.x1), 0.5 * (bbox.y0 + bbox.y1)]) + patch = mpatches.FancyArrowPatch( + posA=inv.transform(center), + posB=inv.transform(target), + transform=ax.transData, + zorder=label.get_zorder() - 0.1, + **props, + ) + # add_artist, not add_patch: a connector is decoration, and letting it + # widen the data limits would move the very labels it is drawn from, + # changing the layout on every redraw. + ax.add_artist(patch) + patch.set_clip_path(ax.patch) + patches.append(patch) + return patches