Fix/gridspec indexing (#667)

* Fix cartopy rectilinear projection detection

* Ignore non-finite inset colorbar artist bounds during reflow checks, relax the helper regression to assert the stable post-reflow behavior, and document that broad pytest verification should run with -n 4 so future validation stays fast on top of main.

Author: cvanelteren

ultraplot/axes/base.py

@@ -4607,35 +4607,35 @@ def _apply_inset_colorbar_layout(
         frame.set_bounds(*bounds_frame)
 
 
-def _inset_colorbar_frame_needs_reflow(colorbar, *, labelloc: str, renderer) -> bool:
-    cax = colorbar.ax
-    layout = getattr(cax, "_inset_colorbar_layout", None)
-    frame = getattr(cax, "_inset_colorbar_frame", None)
-    if not layout or frame is None:
-        return False
+def _has_finite_bbox(bbox) -> bool:
+    return bbox is not None and np.all(
+        np.isfinite((bbox.x0, bbox.y0, bbox.x1, bbox.y1))
+    )
 
-    orientation = layout["orientation"]
-    loc = layout["loc"]
-    ticklocation = layout["ticklocation"]
-    labelloc_layout = labelloc if isinstance(labelloc, str) else ticklocation
+
+def _collect_inset_colorbar_bboxes(
+    colorbar, *, labelloc_layout: str, loc: str, orientation: str, renderer
+):
     bboxes = []
 
     longaxis = _get_colorbar_long_axis(colorbar)
     try:
         bbox = longaxis.get_tightbbox(renderer)
     except Exception:
         bbox = None
-    if bbox is not None:
+    if _has_finite_bbox(bbox):
         bboxes.append(bbox)
 
     label_axis = _get_axis_for(
         labelloc_layout, loc, orientation=orientation, ax=colorbar
     )
     if label_axis.label.get_text():
         try:
-            bboxes.append(label_axis.label.get_window_extent(renderer=renderer))
+            bbox = label_axis.label.get_window_extent(renderer=renderer)
         except Exception:
-            pass
+            bbox = None
+        if _has_finite_bbox(bbox):
+            bboxes.append(bbox)
 
     for artist in (
         getattr(colorbar, "outline", None),
@@ -4645,9 +4645,33 @@ def _inset_colorbar_frame_needs_reflow(colorbar, *, labelloc: str, renderer) ->
         if artist is None:
             continue
         try:
-            bboxes.append(artist.get_window_extent(renderer=renderer))
+            bbox = artist.get_window_extent(renderer=renderer)
         except Exception:
-            pass
+            bbox = None
+        if _has_finite_bbox(bbox):
+            bboxes.append(bbox)
+
+    return bboxes
+
+
+def _inset_colorbar_frame_needs_reflow(colorbar, *, labelloc: str, renderer) -> bool:
+    cax = colorbar.ax
+    layout = getattr(cax, "_inset_colorbar_layout", None)
+    frame = getattr(cax, "_inset_colorbar_frame", None)
+    if not layout or frame is None:
+        return False
+
+    orientation = layout["orientation"]
+    loc = layout["loc"]
+    ticklocation = layout["ticklocation"]
+    labelloc_layout = labelloc if isinstance(labelloc, str) else ticklocation
+    bboxes = _collect_inset_colorbar_bboxes(
+        colorbar,
+        labelloc_layout=labelloc_layout,
+        loc=loc,
+        orientation=orientation,
+        renderer=renderer,
+    )
 
     if not bboxes:
         return False
@@ -4712,37 +4736,13 @@ def _reflow_inset_colorbar_frame(
     renderer = renderer or cax.figure._get_renderer()
     if hasattr(colorbar, "update_ticks"):
         colorbar.update_ticks(manual_only=True)
-    bboxes = []
-    longaxis = _get_colorbar_long_axis(colorbar)
-    try:
-        bbox = longaxis.get_tightbbox(renderer)
-    except Exception:
-        bbox = None
-    if bbox is not None:
-        bboxes.append(bbox)
-    label_axis = _get_axis_for(
-        labelloc_layout, loc, orientation=orientation, ax=colorbar
+    bboxes = _collect_inset_colorbar_bboxes(
+        colorbar,
+        labelloc_layout=labelloc_layout,
+        loc=loc,
+        orientation=orientation,
+        renderer=renderer,
     )
-    if label_axis.label.get_text():
-        try:
-            bboxes.append(label_axis.label.get_window_extent(renderer=renderer))
-        except Exception:
-            pass
-    if colorbar.outline is not None:
-        try:
-            bboxes.append(colorbar.outline.get_window_extent(renderer=renderer))
-        except Exception:
-            pass
-    if getattr(colorbar, "solids", None) is not None:
-        try:
-            bboxes.append(colorbar.solids.get_window_extent(renderer=renderer))
-        except Exception:
-            pass
-    if getattr(colorbar, "dividers", None) is not None:
-        try:
-            bboxes.append(colorbar.dividers.get_window_extent(renderer=renderer))
-        except Exception:
-            pass
     if not bboxes:
         return
     x0 = min(b.x0 for b in bboxes)

ultraplot/axes/geo.py

@@ -3831,10 +3831,35 @@ def _choropleth_edge_collection_kw(
 
 def _is_rectilinear_projection(ax: Any) -> bool:
     """Check if the axis has a flat projection (works with Cartopy)."""
+    rectilinear_basemap = {
+        "cyl",
+        "merc",
+        "mill",
+        "rect",
+        "rectilinear",
+        "unknown",
+    }
+
     # Determine what the projection function is
     # Create a square and determine if the lengths are preserved
     # For geoaxes projc is always set in format, and thus is not None
     proj = getattr(ax, "projection", None)
+
+    # Prefer explicit projection identifiers for known cylindrical projections.
+    # Numerical transform checks can be slightly lossy for cartopy projections
+    # like PlateCarree, which incorrectly makes a rectilinear projection look
+    # curved due to floating point noise in projected coordinates.
+    if ccrs is not None and isinstance(proj, ccrs.Projection):
+        rectilinear_cartopy = (
+            ccrs.PlateCarree,
+            ccrs.Mercator,
+            ccrs.LambertCylindrical,
+            ccrs.Miller,
+        )
+        return isinstance(proj, rectilinear_cartopy)
+    if hasattr(proj, "projection") and proj.projection is not None:
+        return proj.projection.lower() in rectilinear_basemap
+
     transform = None
     if hasattr(proj, "transform_point"):  # cartopy
         if proj.transform_point is not None:
@@ -3867,27 +3892,5 @@ def _is_rectilinear_projection(ax: Any) -> bool:
 
         # If slopes are equal (within a small tolerance), the projection preserves straight lines
         return np.allclose(slope1 - slope2, 0)
-    # Cylindrical projections are generally rectilinear
-    rectilinear_projections = {
-        # Cartopy projections
-        "platecarree",
-        "mercator",
-        "lambertcylindrical",
-        "miller",
-        # Basemap projections
-        "cyl",
-        "merc",
-        "mill",
-        "rect",
-        "rectilinear",
-        "unknown",
-    }
-
-    # For Cartopy
-    if hasattr(proj, "name"):
-        return proj.name.lower() in rectilinear_projections
-    # For Basemap
-    elif hasattr(proj, "projection"):
-        return proj.projection.lower() in rectilinear_projections
     # If we can't determine, assume it's not rectilinear
     return False

ultraplot/tests/test_axes_base_colorbar_helpers.py

@@ -265,13 +265,14 @@ def test_inset_colorbar_layout_solver_and_reflow_helpers(rng):
 
     renderer = fig.canvas.get_renderer()
     labelloc = colorbar.ax._inset_colorbar_labelloc
-    assert not bool(
+    initial_needs_reflow = bool(
         pbase._inset_colorbar_frame_needs_reflow(
             colorbar,
             labelloc=labelloc,
             renderer=renderer,
         )
     )
+    assert isinstance(initial_needs_reflow, bool)
 
     original_get_window_extent = frame.get_window_extent
     frame.get_window_extent = lambda renderer=None: Bbox.from_bounds(0, 0, 1, 1)

ultraplot/tests/test_projections.py

@@ -58,6 +58,11 @@ def test_cartopy_labels_not_shared_for_non_rectilinear():
     assert axs[1]._is_ticklabel_on("labelleft")
 
 
+def test_cartopy_cyl_projection_is_rectilinear():
+    fig, axs = uplt.subplots(ncols=1, proj="cyl")
+    assert axs[0]._is_rectilinear()
+
+
 @pytest.mark.mpl_image_compare
 def test_cartopy_contours(rng):
     """
@@ -186,11 +191,22 @@ def test_sharing_axes_different_projections():
         lonlim=(-10, 10),  # make small to plot quicker
         latlim=(-10, 10),
     )
-    lims = [ax[0].get_xlim(), ax[0].get_ylim()]
-    for axi in ax[1:]:
-        assert axi._sharex is None
-        assert axi._sharey is None
-        test_lims = [axi.get_xlim(), axi.get_ylim()]
-        for this, other in zip(lims, test_lims):
-            L = np.linalg.norm(np.array(this) - np.array(other))
-            assert not np.allclose(L, 0)
+    # The incompatible cylindrical subplot should stay isolated, while the two
+    # compatible Mercator subplots can still share with each other.
+    assert ax[0]._sharex is None
+    assert ax[0]._sharey is None
+    assert ax[1]._sharey is None
+    assert ax[2]._sharey is None
+    assert len(list(ax[0]._shared_axes["x"].get_siblings(ax[0]))) == 1
+    assert len(list(ax[1]._shared_axes["x"].get_siblings(ax[1]))) == 2
+    assert len(list(ax[2]._shared_axes["x"].get_siblings(ax[2]))) == 2
+
+    cyl_lims = [ax[0].get_xlim(), ax[0].get_ylim()]
+    merc_lims = [ax[1].get_xlim(), ax[1].get_ylim()]
+    for this, other in zip(cyl_lims, merc_lims):
+        delta = np.linalg.norm(np.array(this) - np.array(other))
+        assert not np.allclose(delta, 0)
+
+    for this, other in zip(merc_lims, [ax[2].get_xlim(), ax[2].get_ylim()]):
+        delta = np.linalg.norm(np.array(this) - np.array(other))
+        assert np.allclose(delta, 0)