bsplot 0.0.2__py3-none-any.whl

bsplot/streamlines.py

@@ -0,0 +1,314 @@
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import nibabel as nib
+import numpy as np
+from matplotlib import colormaps
+from matplotlib.collections import LineCollection
+from matplotlib.colors import Normalize, TwoSlopeNorm
+from bsplot.colors import Colormap as _resolve_cmap
+
+
+def sample_streamlines(tractogram, N=100000, values=None):
+    total_streamlines = len(tractogram.streamlines)
+    subsample_indices = np.linspace(0, total_streamlines - 1, N, dtype=int)
+    subsampled_streamlines = [tractogram.streamlines[i] for i in subsample_indices]
+    if values is not None:
+        subsampled_values = [values[i] for i in subsample_indices]
+        return subsampled_streamlines, subsampled_values
+    return subsampled_streamlines
+
+
+def find_optimal_slice(streamlines, affine, plane="coronal"):
+    """
+    Find the optimal slice position for visualizing streamlines on a 2D plane.
+
+    Parameters:
+    streamlines (list of ndarray): List of streamlines, where each streamline is an ndarray of shape (n_points, 3).
+    affine (ndarray): Affine transformation matrix for converting voxel coordinates to real-world coordinates.
+    plane (str, optional): The plane to plot ("sagittal", "coronal", "horizontal").
+
+    Returns:
+    optimal_slice (float): The optimal slice position in millimeters.
+    """
+    all_coords = []
+
+    for s_coords in streamlines:
+        s_coords = nib.affines.apply_affine(affine, s_coords)
+        all_coords.append(s_coords)
+
+    all_coords = np.concatenate(all_coords, axis=0)
+
+    if plane == "sagittal":
+        axis_coords = all_coords[:, 0]
+    elif plane == "coronal":
+        axis_coords = all_coords[:, 1]
+    elif plane == "horizontal":
+        axis_coords = all_coords[:, 2]
+
+    # Find the slice with the maximum number of streamlines crossing it
+    hist, bin_edges = np.histogram(axis_coords, bins=50)
+    optimal_slice = bin_edges[np.argmax(hist)]
+
+    return optimal_slice
+
+
+def get_streamline_color(s_coords, cmap=None):
+    """
+    Calculate the color of a streamline based on its principal direction,
+    where x, y, and z components correspond to red, green, and blue channels respectively.
+
+    Parameters:
+    s_coords (ndarray): Streamline coordinates of shape (n_points, 3).
+
+    Returns:
+    color (tuple): RGB color tuple for the streamline.
+    """
+    # Compute the mean direction vector from the start to the end of the streamline
+    mean_direction = s_coords[-1] - s_coords[0]
+    mean_direction = mean_direction / np.linalg.norm(
+        mean_direction
+    )  # Normalize the direction vector
+
+    if cmap:
+        if isinstance(cmap, str):
+            colormap = _resolve_cmap(cmap)
+        else:
+            colormap = cmap
+        scalar_value = np.abs(mean_direction[0])
+
+        color = colormap(scalar_value)  # Returns an RGBA tuple
+    else:
+        # Convert the direction vector into an RGB color
+        # Each component of the mean direction corresponds to a color channel
+        # Ensure the direction is within the range [0, 1] for color mapping
+        r = np.abs(mean_direction[0])  # Red corresponds to x-direction
+        g = np.abs(mean_direction[1])  # Green corresponds to y-direction
+        b = np.abs(mean_direction[2])  # Blue corresponds to z-direction
+
+        # Combine the components into an RGB color
+        color = (r, g, b)
+
+    return color
+
+
+def load_leaddbs_mat(path):
+    """Load a Lead-DBS / DBS-Tractography-Atlas ``.mat`` fiber file.
+
+    The format stores all fiber coordinates in a ``fibers`` matrix of shape
+    ``(N, 4)`` — columns 0-2 are x/y/z in MNI mm space, column 3 is the
+    1-based fiber index — and a ``idx`` vector giving the number of points in
+    each fiber.
+
+    Returns a ``nibabel.streamlines.Tractogram`` with coordinates in RAS mm.
+    """
+    import scipy.io as sio
+
+    d = sio.loadmat(str(path))
+    fibers = d["fibers"].astype(np.float32)
+    lengths = d["idx"].squeeze().astype(int)
+    streamlines = np.split(fibers[:, :3], np.cumsum(lengths)[:-1])
+    return nib.streamlines.Tractogram(
+        streamlines=streamlines,
+        affine_to_rasmm=np.eye(4),
+    )
+
+
+def plot_tractogram(
+    tractogram,
+    ax,
+    cmap=None,
+    view="sagittal",
+    norm=None,
+    vmin=None,
+    vmax=None,
+    alpha=0.8,
+    adjustable="box",
+    values=None,
+    subsample=None,
+    threshold=None,
+    datatype="data_per_streamline",
+    datakey=None,
+    flip_lr=False,  # flip x-coordinate for radiological convention
+    rotate=0,  # rotation in degrees (0, 90, 180, 270) for 2D projection
+    **kwargs,
+):
+    from tqdm import tqdm
+    import importlib
+
+    try:
+        dipy = importlib.import_module("dipy")
+    except ImportError as e:
+        raise ImportError(
+            "plot_tractogram requires 'dipy'. Please install it: pip install dipy"
+        ) from e
+    from matplotlib.cm import ScalarMappable
+
+    # Load tractogram from file if path is provided
+    if isinstance(tractogram, (str, os.PathLike, Path)):
+        path = Path(tractogram)
+        if path.suffix.lower() == ".mat":
+            tractogram = load_leaddbs_mat(path)
+        else:
+            tractogram = nib.streamlines.load(str(path))
+
+    # Prepare streamlines
+    streamlines = tractogram.streamlines.copy()
+
+    # Prepare values if requested via datakey
+    if values is None and datakey is not None:
+        data = getattr(tractogram, datatype)
+        if datakey not in data:
+            raise ValueError(f"Data key '{datakey}' not found in streamlines data.")
+        values = data[datakey]
+
+    def _is_pointwise_values(values_in, streamlines_in):
+        if values_in is None:
+            return False
+        if not isinstance(values_in, (list, tuple, np.ndarray)):
+            return False
+        if len(values_in) != len(streamlines_in):
+            return False
+        try:
+            return all(
+                hasattr(v, "__len__") and len(v) == len(streamlines_in[i])
+                for i, v in enumerate(values_in)
+            )
+        except Exception:
+            return False
+
+    pointwise_values = _is_pointwise_values(values, streamlines)
+
+    # Threshold by values if provided
+    if values is not None and threshold is not None:
+        if pointwise_values:
+            mask = [
+                np.nanmax(np.abs(np.asarray(v))) >= threshold for v in values
+            ]
+            streamlines = [s for s, keep in zip(streamlines, mask) if keep]
+            values = [v for v, keep in zip(values, mask) if keep]
+        else:
+            values = np.asarray(values).ravel()
+            mask = np.abs(values) >= threshold
+            streamlines = streamlines[mask.ravel()]
+            values = values[mask]
+
+    # Subsample consistently from the current set
+    if subsample:
+        n = len(streamlines)
+        N = min(int(subsample), n)
+        idx = np.linspace(0, n - 1, N, dtype=int)
+        # Keep ordering stable; slice both streamlines and values if present
+        streamlines = [streamlines[i] for i in idx]
+        if values is not None:
+            if pointwise_values:
+                values = [values[i] for i in idx]
+            else:
+                values = np.asarray(values).ravel()[idx]
+
+    # Resolve colormap
+    if isinstance(cmap, str):
+        cmap = _resolve_cmap(cmap)
+
+    mappable = None
+    # Build normalization only when we have values
+    if values is not None and norm is None:
+        if pointwise_values:
+            data_min = np.inf
+            data_max = -np.inf
+            for v in values:
+                v_arr = np.asarray(v)
+                if v_arr.size == 0:
+                    continue
+                data_min = min(data_min, np.nanmin(v_arr))
+                data_max = max(data_max, np.nanmax(v_arr))
+            if not np.isfinite(data_min) or not np.isfinite(data_max):
+                data_min, data_max = 0.0, 1.0
+        else:
+            vals_array = np.asarray(values).ravel()
+            data_min = np.nanmin(vals_array)
+            data_max = np.nanmax(vals_array)
+        vmin_eff = data_min if vmin is None else vmin
+        vmax_eff = data_max if vmax is None else vmax
+        # Auto-center on zero when the effective range crosses zero
+        if vmin_eff < 0 and vmax_eff > 0:
+            norm = TwoSlopeNorm(vmin=vmin_eff, vcenter=0.0, vmax=vmax_eff)
+        else:
+            norm = Normalize(vmin=vmin_eff, vmax=vmax_eff)
+
+    # Prepare a ScalarMappable for colorbars when we have values
+    if values is not None:
+        if cmap is None:
+            cmap = _resolve_cmap("viridis")
+        mappable = ScalarMappable(norm=norm, cmap=cmap)
+        # Optional: set_array helps some colorbar behaviors
+        if pointwise_values:
+            mappable.set_array(np.array([vmin_eff, vmax_eff]))
+        else:
+            mappable.set_array(np.asarray(values).ravel())
+
+    for i, streamline in tqdm(enumerate(streamlines), total=len(streamlines)):
+        # Coordinates
+        x, y, z = streamline[:, :3].T
+
+        # Flip x-coordinate for radiological convention if requested
+        if flip_lr:
+            x = -x
+
+        # 2D projection
+        if view.lower() in ["sagittal", "lateral"]:
+            x2d, y2d = y, z
+        elif view.lower() in ["coronal", "frontal"]:
+            x2d, y2d = x, z
+        elif view.lower() in ["horizontal", "axial"]:
+            x2d, y2d = x, y
+
+        # Apply rotation if specified
+        if rotate != 0:
+            angle_rad = np.radians(rotate)
+            cos_a, sin_a = np.cos(angle_rad), np.sin(angle_rad)
+            x2d_rot = x2d * cos_a - y2d * sin_a
+            y2d_rot = x2d * sin_a + y2d * cos_a
+            x2d, y2d = x2d_rot, y2d_rot
+
+        if values is None:
+            color = get_streamline_color(streamline, cmap=cmap)
+            ax.plot(x2d, y2d, color=color, alpha=alpha, **kwargs)
+        elif pointwise_values:
+            v_arr = np.asarray(values[i]).astype(float)
+            if v_arr.ndim != 1 or len(v_arr) != len(x2d):
+                raise ValueError(
+                    "Pointwise values must match streamline point count."
+                )
+            v_seg = 0.5 * (v_arr[:-1] + v_arr[1:])
+            segments = np.stack(
+                [
+                    np.column_stack([x2d[:-1], y2d[:-1]]),
+                    np.column_stack([x2d[1:], y2d[1:]]),
+                ],
+                axis=1,
+            )
+            colors = cmap(norm(v_seg))
+            if alpha == "relative":
+                colors[:, 3] = colors[:, 3] * norm(v_seg)
+            elif alpha is not None:
+                colors[:, 3] = alpha
+
+            lc_kwargs = dict(kwargs)
+            if "linewidth" in lc_kwargs and "linewidths" not in lc_kwargs:
+                lc_kwargs["linewidths"] = lc_kwargs.pop("linewidth")
+            if "lw" in lc_kwargs and "linewidths" not in lc_kwargs:
+                lc_kwargs["linewidths"] = lc_kwargs.pop("lw")
+            lc_kwargs.pop("alpha", None)
+            lc = LineCollection(segments, colors=colors, **lc_kwargs)
+            ax.add_collection(lc)
+        else:
+            v = float(values[i])
+            color = cmap(norm(v))
+            alpha_i = float(norm(v)) if alpha == "relative" else alpha
+            ax.plot(x2d, y2d, color=color, alpha=alpha_i, **kwargs)
+
+        ax.set_aspect("equal", adjustable=adjustable)
+
+    return mappable