Add UltraLayout: Advanced constraint-based positioning for non-orthogonal subplot arrangements

ultraplot/figure.py

@@ -1832,7 +1832,7 @@ def _axes_dict(naxs, input, kw=False, default=None):
         # Create or update the gridspec and add subplots with subplotspecs
         # NOTE: The gridspec is added to the figure when we pass the subplotspec
         if gs is None:
-            gs = pgridspec.GridSpec(*array.shape, **gridspec_kw)
+            gs = pgridspec.GridSpec(*array.shape, layout_array=array, **gridspec_kw)
         else:
             gs.update(**gridspec_kw)
         axs = naxs * [None]  # list of axes

ultraplot/gridspec.py

@@ -6,21 +6,31 @@
 import itertools
 import re
 from collections.abc import MutableSequence
+from functools import wraps
 from numbers import Integral
+from typing import List, Optional, Tuple, Union
 
 import matplotlib.axes as maxes
 import matplotlib.gridspec as mgridspec
 import matplotlib.transforms as mtransforms
 import numpy as np
-from typing import List, Optional, Union, Tuple
-from functools import wraps
 
 from . import axes as paxes
 from .config import rc
-from .internals import ic  # noqa: F401
-from .internals import _not_none, docstring, warnings
+from .internals import (
+    _not_none,
+    docstring,
+    ic,  # noqa: F401
+    warnings,
+)
 from .utils import _fontsize_to_pt, units
-from .internals import warnings
+
+try:
+    from . import ultralayout
+    ULTRA_AVAILABLE = True
+except ImportError:
+    ultralayout = None
+    ULTRA_AVAILABLE = False
 
 __all__ = ["GridSpec", "SubplotGrid"]
 
@@ -225,6 +235,18 @@ def get_position(self, figure, return_all=False):
             nrows, ncols = gs.get_total_geometry()
         else:
             nrows, ncols = gs.get_geometry()
+
+        # Check if we should use UltraLayout for this subplot
+        if isinstance(gs, GridSpec) and gs._use_ultra_layout:
+            bbox = gs._get_ultra_position(self.num1, figure)
+            if bbox is not None:
+                if return_all:
+                    rows, cols = np.unravel_index([self.num1, self.num2], (nrows, ncols))
+                    return bbox, rows[0], cols[0], nrows, ncols
+                else:
+                    return bbox
+
+        # Default behavior: use grid positions
         rows, cols = np.unravel_index([self.num1, self.num2], (nrows, ncols))
         bottoms, tops, lefts, rights = gs.get_grid_positions(figure)
         bottom = bottoms[rows].min()
@@ -264,14 +286,19 @@ def __getattr__(self, attr):
         super().__getattribute__(attr)  # native error message
 
     @docstring._snippet_manager
-    def __init__(self, nrows=1, ncols=1, **kwargs):
+    def __init__(self, nrows=1, ncols=1, layout_array=None, **kwargs):
         """
         Parameters
         ----------
         nrows : int, optional
             The number of rows in the subplot grid.
         ncols : int, optional
             The number of columns in the subplot grid.
+        layout_array : array-like, optional
+            2D array specifying the subplot layout, where each unique integer
+            represents a subplot and 0 represents empty space. When provided,
+            enables UltraLayout constraint-based positioning for non-orthogonal
+            arrangements (requires kiwisolver package).
 
         Other parameters
         ----------------
@@ -301,6 +328,16 @@ def __init__(self, nrows=1, ncols=1, **kwargs):
         manually and want the same geometry for multiple figures, you must create
         a copy with `GridSpec.copy` before working on the subsequent figure).
         """
+        # Layout array for non-orthogonal layouts with UltraLayout
+        self._layout_array = np.array(layout_array) if layout_array is not None else None
+        self._ultra_positions = None  # Cache for UltraLayout-computed positions
+        self._use_ultra_layout = False  # Flag to enable UltraLayout
+
+        # Check if we should use UltraLayout
+        if self._layout_array is not None and ULTRA_AVAILABLE:
+            if not ultralayout.is_orthogonal_layout(self._layout_array):
+                self._use_ultra_layout = True
+
         # Fundamental GridSpec properties
         self._nrows_total = nrows
         self._ncols_total = ncols
@@ -363,6 +400,119 @@ def __init__(self, nrows=1, ncols=1, **kwargs):
         }
         self._update_params(pad=pad, **kwargs)
 
+    def _get_ultra_position(self, subplot_num, figure):
+        """
+        Get the position of a subplot using UltraLayout constraint-based positioning.
+
+        Parameters
+        ----------
+        subplot_num : int
+            The subplot number (in total geometry indexing)
+        figure : Figure
+            The matplotlib figure instance
+
+        Returns
+        -------
+        bbox : Bbox or None
+            The bounding box for the subplot, or None if kiwi layout fails
+        """
+        if not self._use_ultra_layout or self._layout_array is None:
+            return None
+
+        # Ensure figure is set
+        if not self.figure:
+            self._figure = figure
+        if not self.figure:
+            return None
+
+        # Compute or retrieve cached UltraLayout positions
+        if self._ultra_positions is None:
+            self._compute_ultra_positions()
+
+        # Find which subplot number in the layout array corresponds to this subplot_num
+        # We need to map from the gridspec cell index to the layout array subplot number
+        nrows, ncols = self._layout_array.shape
+
+        # Decode the subplot_num to find which layout number it corresponds to
+        # This is a bit tricky because subplot_num is in total geometry space
+        # We need to find which unique number in the layout_array this corresponds to
+
+        # Get the cell position from subplot_num
+        row, col = divmod(subplot_num, self.ncols_total)
+
+        # Check if this is within the layout array bounds
+        if row >= nrows or col >= ncols:
+            return None
+
+        # Get the layout number at this position
+        layout_num = self._layout_array[row, col]
+
+        if layout_num == 0 or layout_num not in self._ultra_positions:
+            return None
+
+        # Return the cached position
+        left, bottom, width, height = self._ultra_positions[layout_num]
+        bbox = mtransforms.Bbox.from_bounds(left, bottom, width, height)
+        return bbox
+
+    def _compute_ultra_positions(self):
+        """
+        Compute subplot positions using UltraLayout and cache them.
+        """
+        if not ULTRA_AVAILABLE or self._layout_array is None:
+            return
+
+        # Get figure size
+        if not self.figure:
+            return
+
+        figwidth, figheight = self.figure.get_size_inches()
+
+        # Convert spacing to inches
+        wspace_inches = []
+        for i, ws in enumerate(self._wspace_total):
+            if ws is not None:
+                wspace_inches.append(ws)
+            else:
+                # Use default spacing
+                wspace_inches.append(0.2)  # Default spacing in inches
+
+        hspace_inches = []
+        for i, hs in enumerate(self._hspace_total):
+            if hs is not None:
+                hspace_inches.append(hs)
+            else:
+                hspace_inches.append(0.2)
+
+        # Get margins
+        left = self.left if self.left is not None else self._left_default if self._left_default is not None else 0.125 * figwidth
+        right = self.right if self.right is not None else self._right_default if self._right_default is not None else 0.125 * figwidth
+        top = self.top if self.top is not None else self._top_default if self._top_default is not None else 0.125 * figheight
+        bottom = self.bottom if self.bottom is not None else self._bottom_default if self._bottom_default is not None else 0.125 * figheight
+
+        # Compute positions using UltraLayout
+        try:
+            self._ultra_positions = ultralayout.compute_ultra_positions(
+                self._layout_array,
+                figwidth=figwidth,
+                figheight=figheight,
+                wspace=wspace_inches,
+                hspace=hspace_inches,
+                left=left,
+                right=right,
+                top=top,
+                bottom=bottom,
+                wratios=self._wratios_total,
+                hratios=self._hratios_total
+            )
+        except Exception as e:
+            warnings._warn_ultraplot(
+                f"Failed to compute UltraLayout: {e}. "
+                "Falling back to default grid layout."
+            )
+            self._use_ultra_layout = False
+            self._ultra_positions = None
+
     def __getitem__(self, key):
         """
         Get a `~matplotlib.gridspec.SubplotSpec`. "Hidden" slots allocated for axes