Fix SubplotGrid indexing and enhance legend placement with 'ref' argument

[cvanelteren]

Closes #459, #460

This PR addresses issues with SubplotGrid indexing and significantly enhances the flexibility and intuitiveness of legend placement.

1. SubplotGrid 1D List/Array Indexing

• Fix: SubplotGrid (returned by subplots) now supports 1D list or numpy array indexing (e.g., axs[[0, 5]]). This returns a new SubplotGrid containing only the selected axes, whereas previously it was misinterpreted as 2D indexing or raised errors.

2. Decoupled Legend Placement (ref argument)

• Feature: Added a ref keyword argument to fig.legend (and fig.colorbar).
  • ax: Specifies the content source (handles/labels).
  • ref: Specifies the location anchor.
  • This allows patterns like fig.legend(ax=axs[1, :], ref=axs[0, :], loc='bottom') to place a legend using data from the second row below the first row.
• Intelligent Placement & Inference:
  • Implicit Span: If ref (or ax if ref is missing) is a list of axes (e.g., axs[[0, 1]]) and no explicit span (rows/cols) is provided, UltraPlot now automatically infers the span based on the legend location. For example, fig.legend(ax=axs[[0, 1]], loc='bottom') will create a single legend spanning the width of both axes.
  • Edge Anchoring: The legend is automatically anchored to the appropriate outer edge of the ref group (e.g., for loc='right', it anchors to the rightmost axis in the group; for loc='bottom', the bottommost).

Documentation & Tests

• Added a new tutorial section "Decoupling legend content and location" to docs/colorbars_legends.py.
• Added comprehensive tests in ultraplot/tests/test_gridspec.py (slicing) and ultraplot/tests/test_legend.py (ref argument, span inference).

API Changes

• No breaking API changes. The ref argument is optional, and standard fig.legend usage remains compatible.

[codecov]

Codecov Report

❌ Patch coverage is 82.02765% with 39 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ultraplot/figure.py	73.13%	18 Missing and 18 partials ⚠️
ultraplot/gridspec.py	77.77%	2 Missing ⚠️
ultraplot/tests/test_legend.py	97.77%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[gepcel]

I've had some simple tests, very brilliant work. Now the logic is very simple and very clear.

It already satisfied for pretty much everything. Bringing up those weird edge cases is kind of like looking for faults, unless you're really into making it perfect and still care about that stuff.

[cvanelteren]

Failing image is shape difference fyi

[cvanelteren]

I've had some simple tests, very brilliant work. Now the logic is very simple and very clear.

It already satisfied for pretty much everything. Bringing up those weird edge cases is kind of like looking for faults, unless you're really into making it perfect and still care about that stuff.

Any cases we should be aware of?

[gepcel]

Any cases we should be aware of?

Example 1

# %% test ultraplot legend
import ultraplot as pplt
fig, axs = pplt.subplots(ncols=4, nrows=4, share='all', width='10cm')
axs.format(yticklabels=[], xticklabels=[])
cycle = pplt.Cycle("default", N=16)
for n, ax in enumerate(axs):
    ax.text(.5, .5, f'ax:{n}\ncols:{n%4+1}', transform=ax.transAxes, ha='center',va='center')
    ax.scatter(.5, .25, transform=ax.transAxes, color=cycle.get_next()['color'], label=f"ax{n}")
fig.legend(ax=axs[1, 1:], loc="b", ref=axs[2,:-1])
fig.legend(ax=axs[[3,5,8,12]], loc="b", ref=axs[0,1:-1])
fig.legend(ax=axs[-2:], loc='l', ref=axs[1:3, 3], ncols=1)

Example 2

# %% test ultraplot legend
import ultraplot as pplt
fig, axs = pplt.subplots(ncols=4, nrows=4, share='all', width='10cm')
axs.format(yticklabels=[], xticklabels=[])
cycle = pplt.Cycle("default", N=16)
for n, ax in enumerate(axs):
    ax.text(.5, .5, f'ax:{n}\ncols:{n%4+1}', transform=ax.transAxes, ha='center',va='center')
    ax.scatter(.5, .25, transform=ax.transAxes, color=cycle.get_next()['color'], label=f"ax{n}")
fig.legend(ax=axs[-2:], loc='l', ref=axs[1:3, 3], ncols=1)
fig.legend(ax=axs[1, 1:], loc="b", ref=axs[2,:-1])
fig.legend(ax=axs[[3,5,8,12]], loc="b", ref=axs[0,1:-1])

The only difference is the last line moves two lines up. It not something I need, just play around for fun.

[cvanelteren]

I see. The legend does take up space to prevent overlap. For outside legends it is more a feature than a bug.

[gepcel]

And there's also overlaps between legend and axs when space of axs is small. Take fig, axs = pplt.subplots(ncols=2, nrows=2, space='1mm') for an example.