Add Sankey diagram visualization functions

docs/api/plot_index.md

@@ -26,6 +26,8 @@ For most tools and for some preprocessing functions, you will find a plotting fu
     plot.ranking
     plot.dendrogram
     plot.catplot
+    plot.sankey_diagram
+    plot.sankey_diagram_time
 ```
 
 ## Quality Control and missing values

ehrapy/_compat.py

@@ -7,6 +7,7 @@
 from subprocess import PIPE, Popen
 from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast
 
+import holoviews as hv
 import numpy as np
 import scipy.sparse as sp
 
@@ -256,3 +257,20 @@ def as_dense_dask_array(a, chunk_size=1000):
     import dask.array as da
 
     return da.from_array(a, chunks=chunk_size)
+
+
+def choose_hv_backend():
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            if hv.Store.current_backend is None:
+                raise RuntimeError(
+                    "No holoviews backend selected. "
+                    "Call holoviews.extension('matplotlib') or "
+                    "holoviews.extension('bokeh') before using this function."
+                )
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

ehrapy/plot/__init__.py

@@ -8,6 +8,7 @@
     missing_values_heatmap,
     missing_values_matrix,
 )
+from ehrapy.plot._sankey import sankey_diagram, sankey_diagram_time
 from ehrapy.plot._scanpy_pl_api import (
     clustermap,
     dendrogram,

ehrapy/plot/_sankey.py

@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, Any
+
+import holoviews as hv
+import numpy as np
+import pandas as pd
+from holoviews import opts
+
+from ehrapy._compat import _raise_array_type_not_implemented, choose_hv_backend
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from ehrdata import EHRData
+
+
+@choose_hv_backend()
+def sankey_diagram(
+    edata: EHRData,
+    *,
+    columns: Sequence[str],
+    node_width: int | float = 20,
+    node_padding: int | float = 10,
+    node_color: str = None,
+    label_position: str | None = "right",
+    show_values: bool = True,
+    title: str | None = None,
+    width: int | None = 600,
+    height: int | None = 400,
+    **kwargs,
+) -> hv.Sankey:
+    """Create a Sankey diagram of relationships across the flat observation table.
+
+    Args:
+        edata : Central data object.
+        columns : Column names from `edata.obs` to visualize
+        node_width : Width of the nodes in the Sankey diagram.
+        node_padding : Padding between nodes in the Sankey diagram.
+        node_color : Color of the nodes. If None, default coloring is used.
+        edge_color : Color of the edges. If None, default coloring is used.
+        label_position : Position of the labels on the nodes. Options are 'left', 'right', 'top', 'bottom', or 'center'.
+        show_values : Whether to display the values on the edges.
+        title : Title of the Sankey diagram.
+        width : Width of the Sankey diagram.
+        height : Height of the Sankey diagram.
+        **kwargs: Additional styling options passed to :class:`holoviews.element.sankey.Sankey`.
+
+    Examples:
+        >>> import ehrdata as ed
+        >>> edata = ed.dt.diabetes_130_fairlearn(columns_obs_only=["gender", "race"])
+        >>> ep.pl.sankey_diagram(edata, columns=["gender", "race"])
+    """
+    if hv.Store.current_backend is None:
+        raise RuntimeError(
+            "No holoviews backend selected. "
+            ":func:`holoviews.extension` with ``matplotlib`` or ``bokeh`` must be called before using this function."
+        )
+    df = edata.obs[columns]
+
+    # Build links between consecutive columns
+    sources, targets, values = [], [], []
+    source_levels, target_levels = [], []
+    for i in range(len(columns) - 1):
+        col_from, col_to = columns[i], columns[i + 1]
+        flows = df.groupby([col_from, col_to]).size().reset_index(name="count")
+        sources.extend(col_from + ": " + flows[col_from].astype("string"))
+        targets.extend(col_to + ": " + flows[col_to].astype("string"))
+        values.extend(flows["count"].to_numpy())
+        source_levels.extend([col_from] * len(flows))
+        target_levels.extend([col_to] * len(flows))
+
+    sankey_df = pd.DataFrame(
+        {
+            "source": sources,
+            "target": targets,
+            "value": values,
+            "source_level": source_levels,
+            "target_level": target_levels,
+        }
+    )
+
+    sankey = hv.Sankey(sankey_df, kdims=["source", "target"], vdims=["value"])
+
+    opts_dict: dict[str, Any] = {}
+
+    if hv.Store.current_backend == "bokeh":
+        if width is not None:
+            opts_dict["width"] = width
+        if height is not None:
+            opts_dict["height"] = height
+
+    if node_width is not None:
+        opts_dict["node_width"] = node_width
+    if node_padding is not None:
+        opts_dict["node_padding"] = node_padding
+    if title is not None:
+        opts_dict["title"] = title
+    if node_color is not None:
+        opts_dict["node_color"] = node_color
+    if label_position is not None:
+        opts_dict["label_position"] = label_position
+    if show_values is not None:
+        opts_dict["show_values"] = show_values
+
+    opts_dict.update(kwargs)
+
+    sankey = sankey.opts(**opts_dict)
+    return sankey
+
+
+@singledispatch
+def _generate_sankey(mtx, time: list[Any], state_labels: dict[int, str] | None = None):
+    _raise_array_type_not_implemented(mtx, type(mtx))
+
+
+@_generate_sankey.register(np.ndarray)
+def _(mtx: np.ndarray, time: list[Any], state_labels: dict[int, str] | None = None) -> pd.DataFrame:
+    if state_labels is None:
+        unique_states = np.unique(mtx)
+        if np.issubdtype(unique_states.dtype, np.floating):
+            unique_states = unique_states[~np.isnan(unique_states)]
+
+        state_labels = {int(state): str(state) for state in unique_states}
+
+    state_values = sorted(state_labels.keys())
+    state_names = [state_labels[val] for val in state_values]
+
+    sources, targets, values = [], [], []
+    for t in range(len(time) - 1):
+        for s_from_idx, s_from_val in enumerate(state_values):
+            for s_to_idx, s_to_val in enumerate(state_values):
+                count = np.sum((mtx[:, t] == s_from_val) & (mtx[:, t + 1] == s_to_val))
+                if count > 0:
+                    source_label = f"{state_names[s_from_idx]} ({time[t]})"
+                    target_label = f"{state_names[s_to_idx]} ({time[t + 1]})"
+                    sources.append(source_label)
+                    targets.append(target_label)
+                    values.append(int(count))
+
+    return pd.DataFrame({"source": sources, "target": targets, "value": values})
+
+
+@choose_hv_backend()
+def sankey_diagram_time(
+    edata: EHRData,
+    *,
+    columns: Sequence[str],
+    layer: str,
+    state_labels: dict[int, str] | None = None,
+    node_width: int | float = 20,
+    node_padding: int | float = 10,
+    node_color: str = None,
+    label_position: str | None = "right",
+    show_values: bool = True,
+    title: str | None = None,
+    width: int | None = 600,
+    height: int | None = 400,
+    **kwargs,
+) -> hv.Sankey:
+    """Create a Sankey diagram showing patient state transitions over time.
+
+    Each node represents a state at a specific time point, and flows show the
+    number of patients transitioning between states.
+    Visualizes how patients transition between different states
+    (e.g. disease severity, treatment status) across consecutive time points.
+
+    Args:
+        edata: Central data object.
+        columns: Variable name from  `edata.var_names` to visualize
+        layer: Name of the layer in `edata.layers` containing the feature data to visualize.
+        state_labels: Mapping from numeric state values to readable labels.
+                    If None, state values will be displayed as strings of their numeric codes (e.g., "0", "1", "2").
+        node_width : Width of the nodes in the Sankey diagram.
+        node_padding : Padding between nodes in the Sankey diagram.
+        node_color : Color of the nodes. If None, default coloring is used.
+        edge_color : Color of the edges. If None, default coloring is used.
+        label_position : Position of the labels on the nodes. Options are 'left', 'right', 'outer', or 'inner'.
+        show_values : Whether to display the values on the edges.
+        title : Title of the Sankey diagram.
+        width : Width of the Sankey diagram.
+        height : Height of the Sankey diagram.
+        **kwargs: Additional styling options passed to :class:`holoviews.element.sankey.Sankey`.
+
+    Examples:
+    >>> import ehrdata as ed
+    >>> edata = ed.dt.ehrdata_blobs(base_timepoints=5, n_variables=1, n_observations=5, random_state=59)
+    >>> edata.layers["tem_data"] = edata.layers["tem_data"].astype(int)
+    >>> state_labels = {-2: "no", -3: "mild", -4: "moderate", -5: "severe", -6: "critical"}
+    >>> plot = sankey_diagram_time(
+    ...     edata,
+    ...     columns=["feature_0"],
+    ...     layer="tem_data",
+    ...     state_labels=state_labels,
+    ... )
+
+    .. image:: /_static/docstring_previews/sankey_time.png
+    """
+    flare_data = edata[:, edata.var_names.isin(columns), :].layers[layer][:, 0, :]
+    time_steps = edata.tem.index.tolist()
+
+    sankey_df = _generate_sankey(flare_data, time=time_steps, state_labels=state_labels)
+
+    sankey = hv.Sankey(sankey_df, kdims=["source", "target"], vdims=["value"])
+
+    opts_dict: dict[str, Any] = {}
+
+    if hv.Store.current_backend == "bokeh":
+        if width is not None:
+            opts_dict["width"] = width
+        if height is not None:
+            opts_dict["height"] = height
+
+    if node_width is not None:
+        opts_dict["node_width"] = node_width
+    if node_padding is not None:
+        opts_dict["node_padding"] = node_padding
+    if title is not None:
+        opts_dict["title"] = title
+    if node_color is not None:
+        opts_dict["node_color"] = node_color
+    if label_position is not None:
+        opts_dict["label_position"] = label_position
+    if show_values is not None:
+        opts_dict["show_values"] = show_values
+
+    opts_dict.update(kwargs)
+
+    sankey = sankey.opts(**opts_dict)
+
+    return sankey