implementation of joyplot

src/cfp/plotting/__init__.py

@@ -1,3 +1,3 @@
-from cfp.plotting._plotting import plot_condition_embedding
+from cfp.plotting._plotting import plot_condition_embedding, plot_densities
 
-__all__ = ["plot_condition_embedding"]
+__all__ = ["plot_condition_embedding", "plot_densities"]

src/cfp/plotting/_plotting.py

@@ -1,22 +1,30 @@
 import types
+from collections.abc import Sequence
 from typing import Any, Literal
 
 import anndata as ad
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import seaborn as sns
 from adjustText import adjust_text
 
-from cfp import _constants
+from cfp import _constants, _logging
 from cfp.plotting._utils import (
     _compute_kernel_pca_from_df,
     _compute_pca_from_df,
     _compute_umap_from_df,
     _get_colors,
+    _grouped_df_to_standard,
+    _is_numeric,
+    _joyplot,
+    _remove_na,
     get_plotting_vars,
 )
 
+__all__ = ["plot_condition_embedding", "plot_densities"]
+
 
 def plot_condition_embedding(
     adata: ad.AnnData,
@@ -175,3 +183,175 @@ def plot_condition_embedding(
     ax.yaxis.set_tick_params(labelsize=fontsize)
 
     return fig if return_fig else None
+
+
+def plot_densities(
+    data: pd.DataFrame,
+    feature: str,
+    group_by: str | None = None,
+    density_fit: Literal["log1p", "raw"] = "raw",
+    ax: mpl.axes.Axes | None = None,
+    figsize: tuple[float, float] | None = None,
+    dpi: int | None = None,
+    xlabels: bool = False,
+    ylabels: bool = True,
+    xlabelsize: float | None = None,
+    xrot: float | None = None,
+    labels: Sequence[Any] = None,
+    ylabelsize: float | None = None,
+    yrot: float | None = None,
+    hist: bool = False,
+    bins: int = 10,
+    fade: bool = False,
+    ylim: Literal["max"] | tuple[float, float] | None = "max",
+    fill: bool = True,
+    linecolor: Any = None,
+    overlap: float = 1.0,
+    background: Any = None,
+    range_style: Literal["all", "individual", "group"] | list[float] = "all",
+    x_range: tuple[float, float] = None,
+    title: str | None = None,
+    colormap: str | mpl.colors.Colormap | None = None,
+    color: Any = None,
+    grid: bool = False,
+    return_fig: bool = True,
+    **kwargs,
+):
+    """Plot kernel density estimations of expressions.
+
+    This function is adapted from https://github.com/leotac/joypy/blob/master/joypy/joyplot.py
+
+    Parameters
+    ----------
+    data
+        :class:`pandas.DataFrame` object containing (predicted) expression values.
+    feature
+        Column in ``'data'`` to plot.
+    group_by
+        Column in ``'data'`` to group by.
+    density_fit
+        Type of density fit to use. If "raw", the kernel density estimation is plotted. If "log1p", the log1p
+        transformed values of the densities are plotted.
+    ax
+        :class:`matplotlib.axes.Axes` used for plotting. If :obj:`None`, create a new one.
+    figsize
+        Size of the figure.
+    dpi
+        Dots per inch.
+    xlabels
+        Whether to show x-axis labels.
+    ylabels
+        Whether to show y-axis labels.
+    xlabelsize
+        Size of the x-axis labels.
+    xrot
+        Rotation (in degrees) of the x-axis labels.
+    labels
+        Sequence of labels for each density plot.
+    ylabelsize
+        Size of the y-axis labels.
+    yrot
+        Rotation (in degrees) of the y-axis labels.
+    hist
+        If :obj:`True`, plot a histogram, otherwise a density plot.
+    bins
+        Number of bins to use, only applicable if ``hist`` is :obj:`True`.
+    fade
+        If :obj:`True`, automatically sets different values of transparency of the density plots.
+    ylim
+        Limits of the y-axis.
+    fill
+        Whether to fill the density plots. If :obj:`False`, only the lines are plotted.
+    linecolor: :mpltype:`color`
+        Color of the contour lines.
+    overlap
+        Overlap between the density plots. The higher the value, the more overlap between densities.
+    background: :mpltype:`color`
+        Background color of the plot.
+    range_style
+        Style of the range. Options are
+
+        - "all" - all density plots have the same range, autmoatically determined.
+        - "individual" - every density plot has its own range, automatically determined.
+        - "group" - each plot has a range that covers the whole group
+        - type :obj:`list` - custom ranges for each density plot.
+
+    x_range
+        Custom range for the x-axis, shared across all density plots. If :obj:`None`, set via ``'range_style'``.
+    title
+        Title of the plot.
+    colormap
+        Colormap to use.
+    color: :mpltype:`color`
+        Color of the density plots.
+    grid
+        Whether to show the grid.
+    return_fig
+        Whether to return the figure.
+    kwargs
+        Additional keyword arguments for the plot.
+
+    Returns
+    -------
+    :class:`matplotlib.figure.Figure` if ``'return_fig'`` is :obj:`True`, else :obj:`None`.
+    """
+    if group_by is not None and isinstance(data, pd.DataFrame):
+        grouped = data.groupby(group_by)
+        if feature is None:
+            feature = list(data.columns)
+            feature.remove(group_by)
+        converted, _labels, sublabels = _grouped_df_to_standard(grouped, feature)  # type: ignore[arg-type]
+        if labels is None:
+            labels = _labels
+    elif isinstance(data, pd.DataFrame):
+        if feature is not None:
+            data = data[feature]
+        converted = [
+            [_remove_na(data[col])] for col in data.columns if _is_numeric(data[col])
+        ]
+        labels = [col for col in data.columns if _is_numeric(data[col])]
+        sublabels = None
+    else:
+        raise TypeError(f"Unknown type for 'data': {type(data)!r}")
+
+    if ylabels is False:
+        labels = None
+
+    if all(len(subg) == 0 for g in converted for subg in g):
+        raise ValueError(
+            "No numeric values found. Joyplot requires at least a numeric column/group."
+        )
+
+    if any(len(subg) == 0 for g in converted for subg in g):
+        _logging.logger.warning("At least a column/group has no numeric values.")
+
+    fig, axes = _joyplot(
+        converted,
+        labels=labels,
+        sublabels=sublabels,
+        grid=grid,
+        xlabelsize=xlabelsize,
+        xrot=xrot,
+        ylabelsize=ylabelsize,
+        yrot=yrot,
+        ax=ax,
+        dpi=dpi,
+        figsize=figsize,
+        hist=hist,
+        bins=bins,
+        fade=fade,
+        ylim=ylim,
+        fill=fill,
+        linecolor=linecolor,
+        overlap=overlap,
+        background=background,
+        xlabels=xlabels,
+        range_style=range_style,
+        x_range=x_range,
+        title=title,
+        colormap=colormap,
+        color=color,
+        density_fit=density_fit,
+        **kwargs,
+    )
+    return fig if return_fig else None