diff --git a/cortex/quickflat/composite.py b/cortex/quickflat/composite.py index dba9b7a32..364e471e8 100644 --- a/cortex/quickflat/composite.py +++ b/cortex/quickflat/composite.py @@ -323,8 +323,39 @@ 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 @@ -332,39 +363,47 @@ def add_colorbar(fig, cimg, colorbar_ticks=None, colorbar_location=(0.4, 0.07, 0 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 @@ -372,7 +411,8 @@ def add_colorbar_2d(fig, cmap_name, colorbar_ticks, 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)) diff --git a/cortex/quickflat/view.py b/cortex/quickflat/view.py index 6fab06173..cb2c4e260 100644 --- a/cortex/quickflat/view.py +++ b/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) diff --git a/cortex/tests/test_quickflat.py b/cortex/tests/test_quickflat.py index 11782f1cc..7c97eaf3b 100644 --- a/cortex/tests/test_quickflat.py +++ b/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])