bsplot 0.0.2__py3-none-any.whl

bsplot/data/surface.py

@@ -0,0 +1,112 @@
+from os.path import abspath, dirname, join
+
+import nibabel as nib
+import templateflow
+
+root = dirname(abspath(__file__))
+
+
+def get_flat_surface_geometry(template="fsaverage", hemi="lh"):
+    """Load a cortical flatmap surface (cut and unfolded mesh).
+
+    Currently available for:
+    - ``fsaverage`` (164k vertices) via nilearn
+
+    Parameters
+    ----------
+    template : str
+        Template name.  Only ``"fsaverage"`` is supported.
+    hemi : str
+        Hemisphere: ``"lh"`` / ``"left"`` / ``"L"`` or
+        ``"rh"`` / ``"right"`` / ``"R"``.
+
+    Returns
+    -------
+    vertices : ndarray, shape (n_vertices, 3)
+    faces : ndarray, shape (n_faces, 3)
+    """
+    from nilearn import datasets
+
+    if hemi.lower() in ("lh", "left", "l"):
+        side = "left"
+    elif hemi.lower() in ("rh", "right", "r"):
+        side = "right"
+    else:
+        raise ValueError(f"Unknown hemisphere '{hemi}'")
+
+    tpl = template.lower()
+    import contextlib, io
+    with contextlib.redirect_stdout(io.StringIO()):
+        if tpl == "fsaverage":
+            fsavg = datasets.fetch_surf_fsaverage(mesh="fsaverage")
+        elif tpl in ("fsaverage5", "fsaverage_5"):
+            fsavg = datasets.fetch_surf_fsaverage(mesh="fsaverage5")
+        else:
+            raise ValueError(
+                f"Flat surface not available for template '{template}'. "
+                "Supported: 'fsaverage', 'fsaverage5'."
+            )
+
+    flat_path = fsavg[f"flat_{side}"]
+    gii = nib.load(flat_path)
+    vertices, faces = gii.agg_data()
+    return vertices, faces
+
+
+def get_surface_geometry(
+    template="fsLR",
+    suffix="midthickness",
+    hemi="lh",
+    density="32k",
+    desc=None,
+    extension="surf.gii",
+):
+
+    if template == "fsLR":
+        if suffix not in [
+            "inflated",
+            "midthickness",
+            "veryinflated",
+            "sphere",
+        ]:
+            raise ValueError(
+                f"Invalid suffix '{suffix}' for template 'fsLR'. Must be one of 'inflated', 'midthickness', 'veryinflated', or 'sphere'."
+            )
+
+    elif template == "fsaverage":
+        if suffix not in [
+            "inflated",
+            "midthickness",
+            "pial",
+            "smoothwm",
+            "sphere",
+            "white",
+        ]:
+            raise ValueError(
+                f"Invalid suffix '{suffix}' for template 'fsaverage'. Must be one of 'inflated', 'midthickness', 'pial', 'smoothwm', 'sphere', or 'white'."
+            )
+        if density not in ["164k", "41k", "10k", "3k"]:
+            raise ValueError(
+                f"Invalid density '{density}' for template 'fsaverage'. Must be one of '164k', '41k', '10k', or '3k'."
+            )
+    if hemi in ["lh", "left"]:
+        hemi = "L"
+    elif hemi in ["rh", "right"]:
+        hemi = "R"
+
+    fsLR = templateflow.api.get(
+        template=template,
+        density=density,
+        suffix=suffix,
+        hemi=hemi,
+        desc=desc,
+        extension=extension,
+    )
+    # print(f"Loading surface geometry from {fsLR}")
+    if isinstance(fsLR, list):
+        raise ValueError(f"Expected a single file path, but got a list: {fsLR}")
+    gii = nib.load(fsLR)
+
+    vertices, triangles = gii.agg_data()
+
+    return vertices, triangles

bsplot/figure.py

@@ -0,0 +1,177 @@
+"""
+Figure layout wizard for bsplot.
+
+Provides a convenient API for creating matplotlib figures with labeled axes
+using ``subplot_mosaic`` under the hood.
+
+Examples
+--------
+Simple single-panel figure:
+
+>>> fig, ax = bsplot.figure.subplots(layout="a")
+
+Row of three panels:
+
+>>> fig, ax = bsplot.figure.subplots(layout="abc")
+
+2×3 grid with mosaic layout:
+
+>>> fig, ax = bsplot.figure.subplots(layout="abc\\ndef")
+
+Spanning cells (matplotlib mosaic syntax):
+
+>>> fig, ax = bsplot.figure.subplots(layout="aab\\nccd")
+
+Access axes by label:
+
+>>> ax["a"].plot([1, 2, 3])
+"""
+
+from __future__ import annotations
+
+import string
+from typing import Any
+
+import matplotlib.pyplot as plt
+import matplotlib.axes
+import matplotlib.figure
+import numpy as np
+
+
+def _parse_layout(layout):
+    """Convert a compact layout string into the list-of-lists expected by
+    ``subplot_mosaic``.
+
+    Parameters
+    ----------
+    layout : str or list[list[str]]
+        If a string, each character becomes an axes label.  Rows are
+        separated by ``\\n``.  Spaces and ``|`` are treated as column
+        separators for readability but are otherwise ignored.
+        A ``.`` marks an empty cell.
+        If already a list-of-lists, it is returned as-is.
+
+    Returns
+    -------
+    list[list[str]]
+        Mosaic specification suitable for ``plt.subplot_mosaic``.
+    """
+    if isinstance(layout, list):
+        return layout
+
+    rows = layout.split("\n")
+    mosaic = []
+    for row in rows:
+        # strip leading/trailing whitespace per row
+        row = row.strip()
+        if not row:
+            continue
+        # allow pipe separators for readability (e.g. "a|b|c")
+        row = row.replace("|", "")
+        # each remaining character is a cell
+        mosaic.append(list(row))
+    return mosaic
+
+
+def subplots(
+    layout: str | list[list[str]] = "a",
+    *,
+    figsize: tuple[float, float] | None = None,
+    dpi: float | None = None,
+    width_ratios: list[float] | None = None,
+    height_ratios: list[float] | None = None,
+    gridspec_kw: dict | None = None,
+    subplot_kw: dict | None = None,
+    per_subplot_kw: dict[str, dict[str, Any]] | None = None,
+    constrained_layout: bool = True,
+    **fig_kw,
+) -> tuple[matplotlib.figure.Figure, dict[str, matplotlib.axes.Axes]]:
+    """Create a figure with labeled subplots using mosaic layout.
+
+    This is bsplot's main entry-point for building multi-panel figures.
+    Under the hood it delegates to :func:`matplotlib.pyplot.subplot_mosaic`,
+    so every feature of the mosaic API is available.
+
+    Parameters
+    ----------
+    layout : str or list[list[str]], default ``"a"``
+        Compact layout specification.
+
+        * A plain string like ``"abc"`` creates a single row of panels
+          labeled *a*, *b*, *c*.
+        * Use ``\\n`` to separate rows: ``"ab\\ncd"`` gives a 2×2 grid.
+        * A ``.`` marks an empty cell.
+        * Repeating a letter spans that panel across cells:
+          ``"aab\\nccd"`` makes *a* span two columns in the first row.
+        * You can also pass a nested list directly for full control.
+
+    figsize : tuple of float, optional
+        Figure size ``(width, height)`` in inches.  When *None* a
+        sensible default is chosen based on the number of rows and
+        columns.
+    dpi : float, optional
+        Resolution.  Defaults to matplotlib's ``figure.dpi`` rcParam.
+    width_ratios : list of float, optional
+        Relative widths of each column.
+    height_ratios : list of float, optional
+        Relative heights of each row.
+    gridspec_kw : dict, optional
+        Extra keyword arguments forwarded to the underlying
+        :class:`~matplotlib.gridspec.GridSpec`.
+    subplot_kw : dict, optional
+        Keyword arguments applied to **every** subplot
+        (e.g. ``{"projection": "polar"}``).
+    per_subplot_kw : dict, optional
+        Per-label keyword arguments, e.g.
+        ``{"a": {"projection": "polar"}}``.
+    constrained_layout : bool, default True
+        Use matplotlib's constrained layout engine for automatic spacing.
+    **fig_kw
+        Any remaining keyword arguments are forwarded to
+        :func:`matplotlib.pyplot.figure`.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    ax : dict[str, matplotlib.axes.Axes]
+        A dict keyed by the panel label strings.
+
+    Examples
+    --------
+    >>> fig, ax = bsplot.figure.subplots()                # single panel
+    >>> fig, ax = bsplot.figure.subplots("abc")           # 1×3 row
+    >>> fig, ax = bsplot.figure.subplots("ab\\ncd")        # 2×2 grid
+    >>> fig, ax = bsplot.figure.subplots("aab\\n.cd")      # mosaic
+    >>> ax["a"].plot([0, 1], [0, 1])
+    """
+    mosaic = _parse_layout(layout)
+
+    # --- derive a default figsize from the grid dimensions ---------------
+    nrows = len(mosaic)
+    ncols = max(len(row) for row in mosaic) if mosaic else 1
+    if figsize is None:
+        col_width = 4.0
+        row_height = 3.5
+        figsize = (ncols * col_width, nrows * row_height)
+
+    # --- assemble gridspec_kw --------------------------------------------
+    if gridspec_kw is None:
+        gridspec_kw = {}
+    if width_ratios is not None:
+        gridspec_kw["width_ratios"] = width_ratios
+    if height_ratios is not None:
+        gridspec_kw["height_ratios"] = height_ratios
+
+    # --- create the figure -----------------------------------------------
+    fig, ax_dict = plt.subplot_mosaic(
+        mosaic,  # type: ignore[arg-type]
+        figsize=figsize,
+        dpi=dpi,
+        gridspec_kw=gridspec_kw or None,
+        subplot_kw=subplot_kw,
+        per_subplot_kw=per_subplot_kw,  # type: ignore[arg-type]
+        layout="constrained" if constrained_layout else None,
+        **fig_kw,
+    )
+
+    return fig, ax_dict

bsplot/graph/__init__.py

@@ -0,0 +1,31 @@
+from . import edges, layout, nodes
+from .flowchart import (
+    build_dot,
+    flowchart_html,
+    prune_nodes_to_edges,
+    render_html,
+    save_html,
+    show_matplotlib,
+    show_matplotlib_bioicons,
+)
+from .network import (
+    get_centers_from_surface_parc,
+    create_network,
+    create_node_mesh,
+    create_edge_mesh,
+    plot_network_on_surface,
+)
+from .nodes import (
+    draw_custom_nodes,
+    draw_nodes_with_icons,
+    draw_icon_nodes,
+)
+
+# Re-export from parent for convenience
+from ..bioicons import (
+    BioIcons,
+    plot_icon,
+    plot_icons_grid,
+    IconAnnotation,
+    search_icons,
+)

bsplot/graph/edges.py

@@ -0,0 +1,312 @@
+import numpy as np
+from matplotlib.colors import Normalize
+from bsplot.colors import Colormap as _resolve_cmap
+
+import numpy as np
+from matplotlib.colors import Normalize
+from matplotlib.patches import FancyArrowPatch
+
+
+def get_edge_color_mapping(G, colormap="viridis", color_by="type"):
+    """
+    Generate a color mapping for edges in graph G based on their specified attribute.
+
+    Parameters:
+    ----------
+    G : networkx.Graph
+        The graph containing the edges with the specified attribute.
+    colormap : str, optional
+        The name of the colormap to use (default is 'viridis').
+    color_by : str, optional
+        The edge attribute to color by (default is 'type').
+
+    Returns:
+    -------
+    edge_colors : dict
+        A dictionary mapping edge attributes to colors.
+    """
+    # Extract all edge attributes for the specified attribute
+    edge_attributes = [
+        data[color_by] if color_by in data.keys() else "n.a."
+        for u, v, data in G.edges(data=True)
+    ]
+    unique_edge_attributes = list(set(edge_attributes))
+
+    # Ensure "n.a." is first in the list
+    if "n.a." in unique_edge_attributes:
+        unique_edge_attributes.remove("n.a.")
+    unique_edge_attributes.insert(0, "n.a.")
+
+    # Create a colormap
+    cmap = _resolve_cmap(colormap)
+    norm = Normalize(vmin=0, vmax=len(unique_edge_attributes) - 1)
+
+    # Map edge attributes to colors
+    edge_colors = {
+        attr: cmap(norm(i)) if attr != "n.a." else (0.6, 0.6, 0.6, 1)
+        for i, attr in enumerate(unique_edge_attributes)
+    }
+
+    return edge_colors
+
+
+def get_unique_node_pairs(G):
+    unique_pairs = set()  # Use a set to store unique node pairs
+    for n1, n2 in G.edges():
+        # Sort using the 'name' attribute of the nodes (assuming each node has a 'name' attribute)
+        pair = tuple(
+            sorted(
+                [n1, n2], key=lambda node: node if isinstance(node, str) else node.name
+            )
+        )
+        unique_pairs.add(pair)  # Add the sorted pair to the set
+    return unique_pairs
+
+
+def compute_midpoint(start, end, rad):
+    """
+    Compute the midpoint of a quadratic Bézier curve.
+
+    Parameters:
+    - start: tuple of (x, y) coordinates of the start point.
+    - end: tuple of (x, y) coordinates of the end point.
+    - rad: curvature radius as a percentage.
+
+    Returns:
+    - tuple of (x, y) coordinates of the midpoint.
+    """
+    x1, y1 = start
+    x2, y2 = end
+    y12 = (y1 + y2) / 2
+    dy = y2 - y1
+    cy = y12 + rad * dy
+
+    # Ensure cy is a single value
+    if isinstance(cy, np.ndarray):
+        cy = cy.item()
+
+    # Midpoint of the quadratic Bézier curve
+    t = 0.5
+    midpoint_x = (1 - t) ** 2 * x1 + 2 * (1 - t) * t * (x1 + (x2 - x1) / 2) + t**2 * x2
+    midpoint_y = (1 - t) ** 2 * y1 + 2 * (1 - t) * t * cy + t**2 * y2
+    return (midpoint_x, midpoint_y)
+
+
+def _bezier_point(x1, y1, x2, y2, cy, t):
+    """Quadratic Bézier point at parameter t. Control point is ((x1+x2)/2, cy)."""
+    cx = (x1 + x2) / 2.0
+    bx = (1 - t)**2 * x1 + 2 * (1 - t) * t * cx + t**2 * x2
+    by = (1 - t)**2 * y1 + 2 * (1 - t) * t * cy + t**2 * y2
+    return float(bx), float(by)
+
+
+def _bezier_tangent(x1, y1, x2, y2, cy, t):
+    """Tangent vector of quadratic Bézier at t (not normalised)."""
+    cx = (x1 + x2) / 2.0
+    tx = 2 * (1 - t) * (cx - x1) + 2 * t * (x2 - cx)
+    ty = 2 * (1 - t) * (cy - y1) + 2 * t * (y2 - cy)
+    return float(tx), float(ty)
+
+
+def plot_curve(
+    ax,
+    start,
+    end,
+    rad,
+    color="black",
+    shrinkA=0,
+    shrinkB=0.03,
+    arrow_size=0.1,
+    arrow_style="->",
+    annotate=None,
+    label_t=0.5,
+    **kwargs,
+):
+    font_size = kwargs.pop("font_size") if "font_size" in kwargs else 12
+
+    x1, y1 = start
+    x2, y2 = end
+    y12 = (y1 + y2) / 2
+    dy = y2 - y1
+    cy = y12 + rad * dy
+    midpoint = compute_midpoint(start, end, rad)
+
+    tau = np.linspace(0, 1, 100)
+    xsupport = np.linspace(x1, x2, 100)
+    ysupport = [(1 - i) ** 2 * y1 + 2 * (1 - i) * i * cy + i**2 * y2 for i in tau]
+
+    # Apply shrink to the start and end
+    shrink_index_start = int(len(xsupport) * shrinkA)
+    shrink_index_end = int(len(xsupport) * (1 - shrinkB))
+
+    line = ax.plot(
+        xsupport[shrink_index_start:shrink_index_end],
+        ysupport[shrink_index_start:shrink_index_end],
+        color=color,
+        **kwargs,
+    )
+    line = line[0]
+    if isinstance(annotate, str):
+        # Label position: Bézier point at label_t
+        lp_x, lp_y = _bezier_point(x1, y1, x2, y2, cy, label_t)
+
+        # True tangent at label_t — varies with both rad and t
+        tx, ty = _bezier_tangent(x1, y1, x2, y2, cy, label_t)
+        length = np.sqrt(tx**2 + ty**2)
+        if length > 0:
+            tx /= length
+            ty /= length
+        else:
+            tx, ty = 1.0, 0.0
+
+        # Right-perpendicular of tangent (rotate 90° clockwise)
+        right_x = ty
+        right_y = -tx
+
+        chord_len = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
+        offset = max(0.08 * chord_len, 0.04)
+        ax.annotate(
+            annotate,
+            xy=(lp_x, lp_y),
+            xytext=(lp_x + offset * right_x, lp_y + offset * right_y),
+            xycoords="data",
+            ha="left",
+            va="center",
+            color="k",
+            fontsize=font_size,
+        )
+
+
+def draw_custom_edges(
+    G,
+    pos,
+    ax=None,
+    edge_labels=False,
+    color_by="type",
+    edge_colors="#606060",
+    edge_radius=0,
+    **kwargs,
+):
+    # if "shrinkA" not in kwargs.keys():
+    #     kwargs["shrinkA"] = .1
+    # if "shrinkB" not in kwargs.keys():
+    #     kwargs["shrinkB"] = .3
+    # if "arrowstyle" not in kwargs.keys():
+    #     kwargs["arrowstyle"] = "-|>,head_length=0.4,head_width=0.2"
+    # if "mutation_scale" not in kwargs.keys():
+    #     kwargs["mutation_scale"] = 10.0
+
+    if edge_colors in colormaps.keys():
+        cmap = (
+            edge_colors
+            if isinstance(edge_colors, str) and edge_colors in colormaps.keys()
+            else "viridis"
+        )
+        colmap = get_edge_color_mapping(G, colormap=cmap, color_by=color_by)
+        edge_colors = [
+            colmap[data[color_by] if color_by in data else "n.a."]
+            for (u, v, data) in G.edges(data=True)
+        ]
+    elif isinstance(edge_colors, str):
+        edge_colors = np.repeat(edge_colors, len(G.edges()))
+
+    for n1, n2 in get_unique_node_pairs(G):
+        point1 = pos[n1]
+        point2 = pos[n2]
+        types12 = {e["type"] for e in G[n1][n2].values()} if n2 in G[n1] else set()
+        types21 = {e["type"] for e in G[n2][n1].values()} if n1 in G[n2] else set()
+
+        num_edges = len(types12) + len(types21)
+
+        if num_edges == 1:
+            rad = [edge_radius]
+        else:
+            rad = np.linspace(-0.5, 0.5, num_edges)
+
+        # Spread label positions along the curve so labels don't stack
+        if num_edges == 1:
+            label_ts = [0.5]
+        else:
+            label_ts = np.linspace(0.3, 0.7, num_edges).tolist()
+
+        i = 0
+        for t in types12:
+            plot_curve(
+                ax,
+                point1,
+                point2,
+                rad[i],
+                color=edge_colors[i],
+                annotate=t if edge_labels else None,
+                label_t=label_ts[i],
+                **kwargs,
+            )
+            i += 1
+        for t in types21:
+            plot_curve(
+                ax,
+                point2,
+                point1,
+                rad[i],
+                color=edge_colors[i],
+                annotate=t if edge_labels else None,
+                label_t=label_ts[i],
+                **kwargs,
+            )
+            i += 1
+
+
+def adjust_arrow_end(start, end, bbox_end):
+    x_start, y_start = start
+    x_end, y_end = end
+
+    x_min_end, y_min_end, width_end, height_end = bbox_end.bounds
+    x_max_end, y_max_end = x_min_end + width_end, y_min_end + height_end
+
+    # x_min_end, y_min_end = bbox_end[0]
+    # x_max_end, y_max_end = bbox_end[1]
+
+    # Determine which edge of the bounding box to intersect with
+    if y_start > y_end:  # Arrow goes downward; intersect with top edge
+        y_bbox = y_max_end
+    else:  # Arrow goes upward; intersect with bottom edge
+        y_bbox = y_min_end
+
+    # Calculate the intersection point
+    if x_start == x_end:  # Vertical line case
+        x_intersect = x_start
+    else:
+        slope = (y_end - y_start) / (x_end - x_start)
+        x_intersect = (y_bbox - y_start) / slope + x_start
+
+    # Clip x_intersect to within the bounding box's horizontal bounds
+    x_intersect = max(x_min_end, min(x_intersect, x_max_end))
+
+    return (x_intersect, y_bbox)
+
+
+def plot_edge(edge, pos, bbox_positions_data, ax, **kwargs):
+    start, end = pos[edge[0]], pos[edge[1]]
+    xstart, ystart = start
+    xend, yend = end
+
+    # bbox_start = bbox_positions_data[edge[0]].bounds
+    bbox_end = bbox_positions_data[edge[1]]
+
+    # x_min_start, y_min_start = bbox_start[0]
+    # x_max_start, y_max_start = bbox_start[1]
+
+    # x_min_end, y_min_end = bbox_end[0]
+    # x_max_end, y_max_end = bbox_end[1]
+
+    start = pos[edge[0]]
+    end = pos[edge[1]]
+    adjusted_end = adjust_arrow_end(start, end, bbox_end)
+
+    # Plot the edge with adjusted endpoint
+    arrow = FancyArrowPatch(
+        start,
+        adjusted_end,
+        **kwargs,
+    )
+    ax.add_patch(arrow)