bsplot 0.0.2__py3-none-any.whl

bsplot/brain.py

@@ -0,0 +1,188 @@
+from os.path import join
+import inspect
+
+import matplotlib.pyplot as plt
+import numpy as np
+from nilearn import surface
+
+from bsplot import anat, data
+
+# imshow kwargs that are not in plot_slice's signature but should still be
+# routed to plot_slice / imshow rather than to Patch schematics
+_IMSHOW_PARAMS = frozenset({
+    "norm", "vmin", "vmax", "interpolation", "interpolation_stage",
+    "filternorm", "filterrad", "resample", "url", "origin", "extent",
+})
+
+
+def _split_glass_brain_kwargs(kwargs):
+    """Split **kwargs into (slice_kwargs, patch_kwargs).
+
+    Kwargs whose names match a ``plot_slice`` parameter or a known
+    ``imshow`` parameter are routed to *slice_kwargs* (for the image).
+    Everything else is routed to *patch_kwargs* (for the brain outline
+    Patch objects drawn by ``plot_brain_schematics``).
+    """
+    from bsplot.volume import plot_slice
+
+    sig = inspect.signature(plot_slice)
+    slice_param_names = {
+        name
+        for name, p in sig.parameters.items()
+        if p.kind != inspect.Parameter.VAR_KEYWORD
+    }
+    slice_param_names |= _IMSHOW_PARAMS
+
+    slice_kw = {}
+    patch_kw = {}
+    for key, value in kwargs.items():
+        if key in slice_param_names:
+            slice_kw[key] = value
+        else:
+            patch_kw[key] = value
+    return slice_kw, patch_kw
+
+
+def glass_brain(
+    img,
+    ax=None,
+    view="sagittal",
+    cmap="viridis",
+    colorbar=False,
+    intensity_projection="absmax",
+    linewidth=None,
+    abs_linewidth=None,
+    **kwargs,
+):
+    # Split kwargs: plot_slice / imshow kwargs  vs  Patch / schematics kwargs
+    slice_kwargs, patch_kwargs = _split_glass_brain_kwargs(kwargs)
+
+    if linewidth is not None:
+        patch_kwargs["linewidth"] = linewidth
+    if abs_linewidth is not None:
+        patch_kwargs["abs_linewidth"] = abs_linewidth
+
+    if view in ["lateral", "sagittal"]:
+        view = "sagittal"
+
+    if ax is None and view.lower() == "all":
+        fig, axs = plt.subplots(ncols=3, layout="compressed", figsize=(9, 4))
+        views = ["coronal", "sagittal", "horizontal"]
+        for i, view in enumerate(views):
+            ax = axs[i]
+            mappable = anat.plot_img(
+                img,
+                intensity_projection=intensity_projection,
+                view=view,
+                cmap=cmap,
+                ax=ax,
+                **slice_kwargs,
+            )
+            anat.plot_brain_schematics(ax, view=view, **patch_kwargs)
+            ax.axis("off")
+        return_fig = True
+    else:
+        if ax is None:
+            fig, ax = plt.subplots()
+            return_fig = True
+        else:
+            return_fig = False
+
+        mappable = anat.plot_img(
+            img, intensity_projection=intensity_projection, view=view, cmap=cmap, ax=ax, **slice_kwargs
+        )
+        # Get image extent before schematics overwrite it
+        img_xlim = ax.get_xlim()
+        img_ylim = ax.get_ylim()
+
+        schematic_bounds = anat.plot_brain_schematics(ax=ax, view=view, **patch_kwargs)
+
+        # Use union of image and schematic bounds
+        xlim = (min(img_xlim[0], schematic_bounds[0]), max(img_xlim[1], schematic_bounds[1]))
+        ylim = (min(img_ylim[0], schematic_bounds[2]), max(img_ylim[1], schematic_bounds[3]))
+        ax.set_xlim(xlim)
+        ax.set_ylim(ylim)
+
+    ax.set_aspect("equal")
+    if return_fig:
+        if colorbar:
+            fig.colorbar(mappable, ax=ax, shrink=0.8)
+        plt.close()
+        return fig
+    else:
+        return mappable
+
+
+def glass_brain_3d(img, ax=None, surf="fsaverage", cmap="viridis", brain_color="k"):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1)
+        return_fig = True
+    else:
+        return_fig = False
+
+    img_data = img.get_fdata()
+
+    if surf == "fsaverage":
+        # fsaverage = datasets.fetch_surf_fsaverage()
+        # mesh = surface.load_surf_mesh(fsaverage.pial_right)
+        mesh = surface.load_surf_mesh(
+            data.surface.get_surface_geometry(
+                template="fsaverage", hemi="rh", suffix="pial", density="10k"
+            )
+        )
+        linewidth = 0.1
+    elif surf == "mni09c":
+        mesh = surface.load_surf_mesh(join(data.surface.root, "MNI152.rh.pial"))
+        linewidth = 0.05
+
+    coords, faces = mesh[0], mesh[1]
+    x, y, z = coords.T
+
+    if surf != "nilearn":
+        ax.triplot(
+            y, z, faces, color=brain_color, alpha=0.8, linewidth=linewidth, zorder=100
+        )
+
+    xlim = ax.get_xlim()
+    ylim = ax.get_ylim()
+
+    affine = img.affine
+
+    voxel_sizes = np.linalg.norm(affine[:3, :3], axis=0)  # Get voxel sizes
+    origin = affine[:3, 3]  # Get translation offset
+    x_coords = np.arange(img_data.shape[1]) * voxel_sizes[1] + origin[1]
+    y_coords = np.arange(img_data.shape[2]) * voxel_sizes[2] + origin[2]
+
+    mip_sagittal = np.max(img_data, axis=0)  # Maximum Intensity Projection
+    mip_sagittal[mip_sagittal == 0] = np.nan
+    mip_sagittal = mip_sagittal.T
+    valid_x = np.any(~np.isnan(mip_sagittal), axis=0)  # Check columns
+    valid_y = np.any(~np.isnan(mip_sagittal), axis=1)  # Check rows
+
+    first_valid_x_index, last_valid_x_index = np.flatnonzero(valid_x)[[0, -1]]
+    first_valid_y_index, last_valid_y_index = np.flatnonzero(valid_y)[[0, -1]]
+
+    ax.imshow(
+        mip_sagittal,
+        cmap=cmap,
+        origin="lower",
+        interpolation="none",
+        alpha=1,
+        extent=[
+            x_coords[0],
+            x_coords[-1],
+            y_coords[0],
+            y_coords[-1],
+        ],
+        zorder=0,
+    )
+    ax.set_xlim(x_coords[first_valid_x_index], x_coords[last_valid_x_index])
+    ax.set_ylim(y_coords[first_valid_y_index], y_coords[last_valid_y_index])
+    ax.set_ylim(ylim)
+    ax.set_xlim(xlim)
+    ax.set_aspect("equal")
+
+    if return_fig:
+        ax.axis("off")
+        plt.close()
+        return fig