Merge pull request #1 from khanlab/axes

Axes

hippunfold_plot/plotting.py

@@ -1,8 +1,28 @@
 import matplotlib.pyplot as plt
 from nilearn.plotting import plot_surf 
 from hippunfold_plot.utils import get_surf_limits, get_data_limits, get_resource_path, check_surf_map_is_label_gii, get_legend_elements_from_label_gii
-
-def plot_hipp_surf(surf_map, density='0p5mm', hemi='left', space=None, figsize=(12, 8), dpi=300, vmin=None, vmax=None, colorbar=False, colorbar_shrink=0.25, cmap=None, view='dorsal', avg_method='median', bg_on_data=True, alpha=0.1, darkness=2, **kwargs):
+from typing import Union, Tuple, Optional, List
+import numpy as np
+
+def plot_hipp_surf(surf_map: Union[str, list], 
+                   density: str = '0p5mm', 
+                   hemi: str = 'left', 
+                   space: Optional[str] = None, 
+                   figsize: Tuple[int, int] = (12, 8), 
+                   dpi: int = 300, 
+                   vmin: Optional[float] = None, 
+                   vmax: Optional[float] = None, 
+                   colorbar: bool = False, 
+                   colorbar_shrink: float = 0.25, 
+                   cmap: Optional[Union[str, plt.cm.ScalarMappable]] = None, 
+                   view: str = 'dorsal', 
+                   avg_method: str = 'median', 
+                   bg_on_data: bool = True, 
+                   alpha: float = 0.1, 
+                   darkness: float = 2, 
+                   axes: Optional[Union[plt.Axes, List[plt.Axes]]] = None, 
+                   figure: Optional[plt.Figure] = None,
+                   **kwargs) -> plt.Figure:
     """Plot hippocampal surface map.
 
     This function plots a surface map of the hippocampus, which can be a label-hippdentate shape.gii, func.gii, or a Vx1 array 
@@ -43,24 +63,24 @@ def plot_hipp_surf(surf_map, density='0p5mm', hemi='left', space=None, figsize=(
         The alpha transparency level. Default is 0.1.
     darkness : float, optional
         The darkness level of the background. Default is 2.
+    axes : matplotlib.axes.Axes or list of matplotlib.axes.Axes, optional
+        Axes to plot on. If None, new axes will be created. If a single axis is provided, it will be used for a single plot.
+        If multiple plots are to be made, a list of axes should be provided.
+    figure : matplotlib.figure.Figure, optional
+        The figure to plot on. If None, a new figure will be created.
     **kwargs : dict
         Additional arguments to pass to nilearn's plot_surf().
 
     Returns
     -------
-    fig : matplotlib.figure.Figure
+    figure : matplotlib.figure.Figure
         The figure object.
-    mappable : matplotlib.cm.ScalarMappable, optional
-        The mappable object, if return_mappable is True.
 
     Notes
     -----
     By default, this function will plot one hemisphere (left by default) in both canonical and unfolded space.
     Both surfaces can be plotted with hemi=None, but the same surf_map will be plotted on both.
 
-    Use return_mappable=True if you want to make a colorbar afterwards, e.g.:
-    fig, mappable = plot_hipp_surf(..., return_mappable=True)
-    plt.colorbar(mappable, shrink=0.5)  # shrink makes it smaller which is recommended
     """
     # Validate inputs
     valid_densities = ['unfoldiso', '0p5mm', '1mm', '2mm']
@@ -87,9 +107,33 @@ def plot_hipp_surf(surf_map, density='0p5mm', hemi='left', space=None, figsize=(
 
     #add any user arguments
     plot_kwargs.update(kwargs)
-
-    # Create a figure
-    fig = plt.figure(figsize=figsize,dpi=dpi)  # Adjust figure size for tall axes
+
+    # Define the plotting order for each hemisphere
+    hemi_space_map = {
+        'left': ['unfold', 'canonical'],
+        'right': ['canonical', 'unfold']
+    }
+
+    # Determine the number of plots to be made
+    hemis_to_plot = [hemi] if hemi else hemi_space_map.keys()
+    num_plots = sum(len([space] if space else hemi_space_map[h]) for h in hemis_to_plot)
+
+    # Validate axes input
+    if axes is not None:
+        if isinstance(axes, plt.Axes):
+            if num_plots > 1:
+                raise ValueError("Multiple plots requested, but only one axis provided.")
+            axes = [axes]
+        elif isinstance(axes, np.ndarray):
+            if len(axes) != num_plots:
+                raise ValueError(f"Expected {num_plots} axes, but got {len(axes)}.")
+        else:
+            raise ValueError("Invalid type for 'axes'. Expected matplotlib.axes.Axes or array of matplotlib.axes.Axes.")
+
+
+    # Create a figure if not provided
+    if figure is None:
+        figure = plt.figure(figsize=figsize, dpi=dpi)
 
     # Define positions for 4 tall side-by-side axes
     positions = [
@@ -98,29 +142,22 @@ def plot_hipp_surf(surf_map, density='0p5mm', hemi='left', space=None, figsize=(
     [0.30, 0.1, 0.2, 0.8],
     [0.43, 0.1, 0.2, 0.8],
     [0.55, 0.1, 0.2, 0.8],
-
     ]
-
-    # Define the plotting order for each hemisphere
-    hemi_space_map = {
-        'left': ['unfold', 'canonical'],
-        'right': ['canonical', 'unfold']
-    }
-
 
+    # Create axes if not provided
+    if axes is None:
+        axes = [figure.add_axes(positions[i], projection='3d') for i in range(num_plots)]
+
     pos=0
 
     # Build the composite plot
-    hemis_to_plot = [hemi] if hemi else hemi_space_map.keys()
     for h in hemis_to_plot:
         spaces_to_plot = [space] if space else hemi_space_map[h]
         for s in spaces_to_plot:
-
-            ax = fig.add_axes(positions[pos], projection='3d')  # Add 3D axes
-
+            ax = axes[pos] 
             plot_surf(surf_mesh=surf_gii.format(hemi=h,space=s,density=density),
                           axes=ax,
-                          figure=fig,
+                          figure=figure,
                           **plot_kwargs)
             (xlim_kwargs,ylim_kwargs) = get_surf_limits(surf_mesh=surf_gii.format(hemi=h,space=s,density=density))
             ax.set_xlim(**xlim_kwargs)
@@ -135,7 +172,7 @@ def plot_hipp_surf(surf_map, density='0p5mm', hemi='left', space=None, figsize=(
         norm = mpl.colors.Normalize(vmin=vmin if vmin else datamin, vmax=vmax if vmax else datamax)  # Match your data range
         sm = mpl.cm.ScalarMappable(cmap=cmap, norm=norm)
         sm.set_array([])  # Dummy array for ScalarMappable
-        plt.colorbar(sm,ax=fig.axes,shrink=colorbar_shrink)
+        plt.colorbar(sm,ax=figure.axes,shrink=colorbar_shrink)
 
-    return fig
+    return figure
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "hippunfold_plot"
 version = "0.1.0"
-description = "Add your description here"
+description = "Plotting functions for hippocampal surface maps from HippUnfold"
 readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [