FIX colorbar placement when passing axis to quickshow

cortex/quickflat/composite.py

@@ -323,56 +323,96 @@ def add_hatch(fig, hatch_data, extents=None, height=None, hatch_space=4,
     return img
 
 
-def add_colorbar(fig, cimg, colorbar_ticks=None, colorbar_location=(0.4, 0.07, 0.2, 0.04), 
-                 orientation='horizontal'):
+def _transform_colorbar_location(colorbar_location, ax_bounds):
+    """Transform colorbar location from axis-relative to figure-relative coordinates.
+
+    Parameters
+    ----------
+    colorbar_location : tuple
+        (left, bottom, width, height) in normalized coordinates (0-1)
+    ax_bounds : tuple or None
+        (left, bottom, width, height) of the axis in figure coordinates.
+        If None, returns colorbar_location unchanged.
+
+    Returns
+    -------
+    tuple
+        Transformed colorbar location in figure coordinates
+    """
+    if ax_bounds is None:
+        return colorbar_location
+
+    cb_left, cb_bottom, cb_width, cb_height = colorbar_location
+    ax_left, ax_bottom, ax_width, ax_height = ax_bounds
+
+    # Transform from axis-relative (0-1) to figure-relative coordinates
+    new_left = ax_left + cb_left * ax_width
+    new_bottom = ax_bottom + cb_bottom * ax_height
+    new_width = cb_width * ax_width
+    new_height = cb_height * ax_height
+
+    return (new_left, new_bottom, new_width, new_height)
+
+
+def add_colorbar(fig, cimg, colorbar_ticks=None, colorbar_location=(0.4, 0.07, 0.2, 0.04),
+                 orientation='horizontal', ax_bounds=None):
     """Add a colorbar to a flatmap plot
 
     Parameters
     ----------
     fig : matplotlib Figure object
         Figure into which to insert colormap
     cimg : matplotlib.image.AxesImage object
-        Image for which to create colorbar. For reference, matplotlib.image.AxesImage 
+        Image for which to create colorbar. For reference, matplotlib.image.AxesImage
         is the output of imshow()
     colorbar_ticks : array-like
         values for colorbar ticks
     colorbar_location : array-like
-        Four-long list, tuple, or array that specifies location for colorbar axes 
-        [left, top, width, height] (?)
+        Four-long list, tuple, or array that specifies location for colorbar axes
+        [left, bottom, width, height] in normalized figure coordinates (0-1)
     orientation : string
         'vertical' or 'horizontal'
+    ax_bounds : tuple or None
+        If provided, colorbar_location is interpreted as relative to these bounds
+        (left, bottom, width, height) instead of the full figure
     """
     fig, _ = _get_fig_and_ax(fig)
-    cbar = fig.add_axes(colorbar_location)
+    transformed_location = _transform_colorbar_location(colorbar_location, ax_bounds)
+    cbar = fig.add_axes(transformed_location)
     fig.colorbar(cimg, cax=cbar, orientation=orientation, ticks=colorbar_ticks)
     return cbar
 
 
 def add_colorbar_2d(fig, cmap_name, colorbar_ticks,
-                    colorbar_location=(0.425, 0.02, 0.15, 0.15), fontsize=12):
+                    colorbar_location=(0.425, 0.02, 0.15, 0.15), fontsize=12,
+                    ax_bounds=None):
     """Add a 2D colorbar to a flatmap plot
 
     Parameters
     ----------
     fig : matplotlib Figure object
-    cimg : matplotlib.image.AxesImage object
-        Image for which to create colorbar. For reference, matplotlib.image.AxesImage 
-        is the output of imshow()
+        Figure into which to insert colormap
+    cmap_name : str
+        Name of the 2D colormap to display
     colorbar_ticks : array-like
-        values for colorbar ticks
+        Values for colorbar ticks [vmin, vmax, vmin2, vmax2]
     colorbar_location : array-like
-        Four-long list, tuple, or array that specifies location for colorbar axes 
-        [left, top, width, height] (?)
-    orientation : string
-        'vertical' or 'horizontal'
+        Four-long list, tuple, or array that specifies location for colorbar axes
+        [left, bottom, width, height] in normalized figure coordinates (0-1)
+    fontsize : int
+        Font size for tick labels
+    ax_bounds : tuple or None
+        If provided, colorbar_location is interpreted as relative to these bounds
+        (left, bottom, width, height) instead of the full figure
     """
     # a bit sketchy - lazy imports
     import matplotlib.pyplot as plt
     import os
     cmap_dir = config.get('webgl', 'colormaps')
     cim = plt.imread(os.path.join(cmap_dir, cmap_name + '.png'))
     fig, _ = _get_fig_and_ax(fig)
-    fig.add_axes(colorbar_location)
+    transformed_location = _transform_colorbar_location(colorbar_location, ax_bounds)
+    fig.add_axes(transformed_location)
     cbar = plt.imshow(cim, extent=colorbar_ticks, interpolation='bilinear')
     cbar.axes.set_xticks(colorbar_ticks[:2])
     cbar.axes.set_xticklabels(colorbar_ticks[:2], fontdict=dict(size=fontsize))

cortex/quickflat/view.py

@@ -124,6 +124,8 @@ def make_figure(braindata, recache=False, pixelwise=True, thick=32, sampler='nea
     dataview = dataset.normalize(braindata)
     if not isinstance(dataview, dataset.Dataview):
         raise TypeError('Please provide a Dataview (e.g. an instance of cortex.Volume, cortex.Vertex, etc), not a Dataset')
+    # Track if user provided their own axis (for colorbar placement)
+    ax_bounds = None
     if fig is None:
         fig_resize = True
         fig = plt.figure()
@@ -136,6 +138,8 @@ def make_figure(braindata, recache=False, pixelwise=True, thick=32, sampler='nea
         fig_resize = False
         ax = fig
         fig = ax.figure
+        # Get axis bounds for colorbar placement within subplot
+        ax_bounds = ax.get_position().bounds
 
     # Add data
     data_im, extents = composite.add_data(ax, dataview, pixelwise=pixelwise, thick=thick, sampler=sampler,
@@ -220,12 +224,14 @@ def make_figure(braindata, recache=False, pixelwise=True, thick=32, sampler='nea
                 ], 2)
             colorbar = composite.add_colorbar_2d(
                 ax, dataview.cmap, colorbar_ticks,
-                colorbar_location=colorbar_location)
+                colorbar_location=colorbar_location,
+                ax_bounds=ax_bounds)
         else:
             colorbar = composite.add_colorbar(
                 ax, data_im,
                 colorbar_location=colorbar_location,
-                colorbar_ticks=colorbar_ticks
+                colorbar_ticks=colorbar_ticks,
+                ax_bounds=ax_bounds
             )
         # Reset axis to main figure axis
         plt.sca(ax)

cortex/tests/test_quickflat.py

@@ -27,6 +27,47 @@ def test_colorbar_location():
                                      colorbar_location='unknown_location')
 
 
+def test_colorbar_placement_with_subplot():
+    """Test that colorbar is placed within subplot bounds when axis is passed."""
+    import matplotlib.pyplot as plt
+
+    # Create a figure with 1x2 subplots
+    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+
+    vol1 = cortex.Volume.random("S1", "fullhead", cmap="hot", vmin=0, vmax=1)
+    vol2 = cortex.Volume.random("S1", "fullhead", cmap="viridis", vmin=0, vmax=1)
+
+    # Plot to both subplots (with_rois=False to avoid inkscape dependency)
+    cortex.quickshow(vol1, fig=axes[0], with_colorbar=True,
+                     colorbar_location='center', with_rois=False)
+    cortex.quickshow(vol2, fig=axes[1], with_colorbar=True,
+                     colorbar_location='center', with_rois=False)
+
+    # Get bounds of subplots and colorbars
+    all_axes = fig.get_axes()
+    # Should have 4 axes: 2 subplots + 2 colorbars
+    assert len(all_axes) == 4, f"Expected 4 axes, got {len(all_axes)}"
+
+    left_subplot_bounds = axes[0].get_position().bounds
+    right_subplot_bounds = axes[1].get_position().bounds
+    cb1_bounds = all_axes[2].get_position().bounds
+    cb2_bounds = all_axes[3].get_position().bounds
+
+    # Verify colorbar 1 is within left subplot horizontal bounds
+    assert cb1_bounds[0] >= left_subplot_bounds[0] - 0.001, \
+        "Colorbar 1 left edge is outside left subplot"
+    assert cb1_bounds[0] + cb1_bounds[2] <= left_subplot_bounds[0] + left_subplot_bounds[2] + 0.001, \
+        "Colorbar 1 right edge is outside left subplot"
+
+    # Verify colorbar 2 is within right subplot horizontal bounds
+    assert cb2_bounds[0] >= right_subplot_bounds[0] - 0.001, \
+        "Colorbar 2 left edge is outside right subplot"
+    assert cb2_bounds[0] + cb2_bounds[2] <= right_subplot_bounds[0] + right_subplot_bounds[2] + 0.001, \
+        "Colorbar 2 right edge is outside right subplot"
+
+    plt.close(fig)
+
+
 @pytest.mark.skipif(no_inkscape, reason='Inkscape required')
 @pytest.mark.parametrize("type_", ["thick", "thin"])
 @pytest.mark.parametrize("nanmean", [True, False])