Hotfix get_border_axes (#236)

* simplified logic and basing reference on original grid

* rm print statements

* spelling

* Update ultraplot/figure.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* do more comprehensive checking

* rm debug

* add some reasoning

* fixed

* add a stale check

* compound tests and check for the failure

* update calls

* update tests

* formatting restored to defaults

* don't adjust labels when not sharing

* tests pass -- some expected failures

* rm debug statements

* fix test

* Update ultraplot/tests/test_geographic.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ultraplot/figure.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ultraplot/axes/geo.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Revert changes

* rm duplicate

* fixes

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: cvanelteren

ultraplot/axes/geo.py

@@ -585,6 +585,9 @@ def __share_axis_setup(
         if level > 1 and limits:
             self._share_limits_with(other, which=which)
 
+        if level >= 1 and labels:
+            self._share_labels_with_others()
+
     @override
     def _sharey_setup(self, sharey, *, labels=True, limits=True):
         """
@@ -666,6 +669,17 @@ def _apply_axis_sharing(self):
                 target_axis=self._lataxis,
             )
 
+        # This block is apart of the draw sequence as the
+        # gridliner object is created late in the
+        # build chain.
+        if not self.stale:
+            return
+        if self.figure._get_sharing_level() == 0:
+            return
+        # Share labels with all levels higher or equal
+        # to 1.
+        self._share_labels_with_others()
+
     def _get_gridliner_labels(
         self,
         bottom=None,
@@ -719,9 +733,11 @@ def _handle_axis_sharing(
             target_axis.set_view_interval(*source_axis.get_view_interval())
             target_axis.set_minor_locator(source_axis.get_minor_locator())
 
-        if not self.stale:
-            return
-
+    def _share_labels_with_others(self):
+        """
+        Helpers function to ensure the labels
+        are shared for rectilinear GeoAxes.
+        """
         # Turn all labels off
         # Note: this action performs it for all the axes in
         # the figure. We use the stale here to only perform
@@ -735,10 +751,7 @@ def _handle_axis_sharing(
 
         # We turn off the tick labels when the scale and
         # ticks are shared (level >= 3)
-        are_ticks_on = True
-        if self.figure._get_sharing_level() >= 3:
-            are_ticks_on = False
-
+        are_ticks_on = False
         default = dict(
             left=are_ticks_on,
             right=are_ticks_on,
@@ -751,11 +764,14 @@ def _handle_axis_sharing(
             # sharing that is specific for the GeoAxes.
             if not isinstance(axi, GeoAxes):
                 continue
-
+            gridlabels = self._get_gridliner_labels(
+                bottom=True, top=True, left=True, right=True
+            )
             sides = recoded.get(axi, [])
             tmp = default.copy()
             for side in sides:
-                tmp[side] = True
+                if side in gridlabels and gridlabels[side]:
+                    tmp[side] = True
             axi._toggle_gridliner_labels(**tmp)
         self.stale = False
 
@@ -1430,6 +1446,8 @@ def _get_gridliner_labels(
         ):
             if side != True:
                 continue
+            if self.gridlines_major is None:
+                continue
             sides[dir] = getattr(self.gridlines_major, f"{dir}_label_artists")
         return sides
 

ultraplot/figure.py

@@ -914,66 +914,84 @@ def _get_border_axes(self) -> dict[str, list[paxes.Axes]]:
         """
 
         gs = self.gridspec
-        all_axes = self.axes
+
+        # Skip colorbars or panels etc
+        all_axes = [axi for axi in self.axes if axi.number is not None]
 
         # Handle empty cases
         nrows, ncols = gs.nrows, gs.ncols
+        border_axes = dict(top=[], bottom=[], left=[], right=[])
         if nrows == 0 or ncols == 0 or not all_axes:
-            return dict(top=[], bottom=[], left=[], right=[])
+            return border_axes
+        # We cannot use the gridspec on the axes as it
+        # is modified when a colorbar is added. Use self.gridspec
+        # as a reference.
+        # Reconstruct the grid based on axis locations. Note that
+        # spanning axes will fit into one of the boxes. Check
+        # this with unittest to see how empty axes are handles
+        grid = np.zeros((gs.nrows, gs.ncols))
+        for axi in all_axes:
+            # Infer coordinate from grdispec
+            spec = axi.get_subplotspec()
+            spans = spec._get_rows_columns()
+            rowspans = spans[:2]
+            colspans = spans[-2:]
+
+            grid[
+                rowspans[0] : rowspans[1] + 1,
+                colspans[0] : colspans[1] + 1,
+            ] = axi.number
+        directions = {
+            "left": (0, -1),
+            "right": (0, 1),
+            "top": (-1, 0),
+            "bottom": (1, 0),
+        }
 
-        # Find occupied grid cells and valid axes
-        occupied_cells = set()
-        axes_with_spec = []
+        def is_border(pos, grid, target, direction):
+            x, y = pos
+            # Check if we are at an edge of the grid (out-of-bounds).
+            if x < 0:
+                return True
+            elif x > grid.shape[0] - 1:
+                return True
+
+            if y < 0:
+                return True
+            elif y > grid.shape[1] - 1:
+                return True
+
+            # Check if we reached a plot or an internal edge
+            if grid[x, y] != target and grid[x, y] > 0:
+                return False
+            if grid[x, y] == 0:
+                return True
+            dx, dy = direction
+            new_pos = (x + dx, y + dy)
+            return is_border(new_pos, grid, target, direction)
+
+        from itertools import product
 
         for axi in all_axes:
             spec = axi.get_subplotspec()
-            if spec is not None:
-                axes_with_spec.append((axi, spec))
-                r0, r1 = spec.rowspan.start, spec.rowspan.stop
-                c0, c1 = spec.colspan.start, spec.colspan.stop
-                for r in range(r0, r1):
-                    for c in range(c0, c1):
-                        occupied_cells.add((r, c))
-
-        if not axes_with_spec:
-            return dict(top=[], bottom=[], left=[], right=[])
-
-        # Initialize border axes sets
-        border_axes_sets = dict(top=set(), bottom=set(), left=set(), right=set())
-
-        # Check each axis against border criteria
-        for axi, spec in axes_with_spec:
-            r0, r1 = spec.rowspan.start, spec.rowspan.stop
-            c0, c1 = spec.colspan.start, spec.colspan.stop
-
-            # Check top border
-            if r0 == 0 or (
-                r0 == 1 and any((0, c) not in occupied_cells for c in range(c0, c1))
-            ):
-                border_axes_sets["top"].add(axi)
-
-            # Check bottom border
-            if r1 == nrows or (
-                r1 == nrows - 1
-                and any((nrows - 1, c) not in occupied_cells for c in range(c0, c1))
-            ):
-                border_axes_sets["bottom"].add(axi)
-
-            # Check left border
-            if c0 == 0 or (
-                c0 == 1 and any((r, 0) not in occupied_cells for r in range(r0, r1))
-            ):
-                border_axes_sets["left"].add(axi)
-
-            # Check right border
-            if c1 == ncols or (
-                c1 == ncols - 1
-                and any((r, ncols - 1) not in occupied_cells for r in range(r0, r1))
-            ):
-                border_axes_sets["right"].add(axi)
-
-        # Convert sets to lists
-        return {key: list(val) for key, val in border_axes_sets.items()}
+            spans = spec._get_rows_columns()
+            rowspan = spans[:2]
+            colspan = spans[-2:]
+            # Check all cardinal directions. When we find a
+            #  border for any starting conditions we break and
+            # consider it a border. This could mean that for some
+            # partial overlaps we consider borders that should
+            # not be borders -- we are conservative in this
+            # regard
+            for direction, d in directions.items():
+                xs = range(rowspan[0], rowspan[1] + 1)
+                ys = range(colspan[0], colspan[1] + 1)
+                for x, y in product(xs, ys):
+                    pos = (x, y)
+                    if is_border(pos=pos, grid=grid, target=axi.number, direction=d):
+                        border_axes[direction].append(axi)
+                        break
+        return border_axes
 
     def _get_align_coord(self, side, axs, includepanels=False):
         """

ultraplot/tests/test_axes.py

@@ -260,6 +260,14 @@ def test_sharing_labels_top_right():
 
 
 def test_sharing_labels_top_right_odd_layout():
+
+    # Helper function to check if the labels
+    # on an axis direction is visible
+    def check_state(numbers: list, state: bool, which: str):
+        for number in numbers:
+            for label in getattr(ax[number], f"get_{which}ticklabels")():
+                assert label.get_visible() == state
+
     layout = [
         [1, 2, 0],
         [1, 2, 5],
@@ -272,11 +280,6 @@ def test_sharing_labels_top_right_odd_layout():
         yticklabelloc="r",
     )
 
-    def check_state(numbers: list, state: bool, which: str):
-        for number in numbers:
-            for label in getattr(ax[number], f"get_{which}ticklabels")():
-                assert label.get_visible() == state
-
     # these correspond to the indices of the axis
     # in the axes array (so the grid number minus 1)
     check_state([0, 2], False, which="y")
@@ -291,15 +294,15 @@ def check_state(numbers: list, state: bool, which: str):
         [4, 0, 5],
     ]
 
-    fig, ax = uplt.subplots(layout)
+    fig, ax = uplt.subplots(layout, hspace=0.2, wspace=0.2, share=1)
     ax.format(
         xticklabelloc="t",
         yticklabelloc="r",
     )
     # these correspond to the indices of the axis
     # in the axes array (so the grid number minus 1)
-    check_state([0, 3], False, which="y")
+    check_state([0, 3], True, which="y")
     check_state([1, 2, 4], True, which="y")
     check_state([0, 1, 2], True, which="x")
-    check_state([3, 4], False, which="x")
+    check_state([3, 4], True, which="x")
     uplt.close(fig)

ultraplot/tests/test_format.py

@@ -58,7 +58,7 @@ def test_patch_format():
     """
     Test application of patch args on initialization.
     """
-    fig = uplt.figure(suptitle="Super title")
+    fig = uplt.figure(suptitle="Super title", share=0)
     fig.subplot(
         121, proj="cyl", labels=True, land=True, latlines=20, abcloc="l", abc="[A]"
     )
@@ -82,7 +82,7 @@ def test_multi_formatting():
     Support formatting in multiple projections.
     """
     # Mix Cartesian with a projection
-    fig, axs = uplt.subplots(ncols=2, proj=("cart", "cyl"))
+    fig, axs = uplt.subplots(ncols=2, proj=("cart", "cyl"), share=0)
     axs[0].pcolormesh(np.random.rand(5, 5))
 
     # Warning is raised based on projection. Cart does not have lonlim, latllim or labels

ultraplot/tests/test_geographic.py

@@ -253,6 +253,12 @@ def are_labels_on(ax, which=["top", "bottom", "right", "left"]) -> tuple[bool]:
     n = 3
     settings = dict(land=True, ocean=True, labels="both")
     fig, ax = uplt.subplots(ncols=n, nrows=n, share="all", proj="cyl")
+    # Add data and ensure the tests still hold
+    # Adding a colorbar will change the underlying gridspec, the
+    # labels should still be correctly treated.
+    data = np.random.rand(10, 10)
+    h = ax.imshow(data)[0]
+    fig.colorbar(h, loc="r")
     ax.format(**settings)
     fig.canvas.draw()  # need a draw to trigger ax.draw for  sharing
 
@@ -574,7 +580,7 @@ def assert_views_are_sharing(ax):
             assert_views_are_sharing(axi)
             # When we share the labels but not the limits,
             # we expect all ticks to be on
-            if level < 3:
+            if level == 0:
                 assert s == 4
             else:
                 assert s == 2
@@ -605,7 +611,7 @@ def test_cartesian_and_geo():
         ax[0].pcolormesh(np.random.rand(10, 10))
         ax[1].scatter(*np.random.rand(2, 100))
         ax[0]._apply_axis_sharing()
-        assert mocked.call_count == 1
+        assert mocked.call_count == 2
     return fig