Add Hypnogram.plot_hypnodensity method (#233)

Author: raphaelvallat

HYPNOGRAM_ROADMAP.md

@@ -28,6 +28,7 @@ Goal: make `yasa.Hypnogram` the industry-standard Python object for handling sle
 
 ### Visualization & export
 - `plot_hypnogram()` — standard hypnogram figure
+- `plot_hypnodensity()` — per-epoch stage probabilities as a stacked area chart (requires `proba`); supports 2/3/4/5-stage hypnograms and datetime x-axis when `start` is set
 - `as_int()` — integer-encoded `pandas.Series`
 - `as_events()` — BIDS-compatible events `DataFrame` (onset, duration, stage)
 - `upsample(new_freq)` — change epoch resolution
@@ -42,8 +43,5 @@ Goal: make `yasa.Hypnogram` the industry-standard Python object for handling sle
 ### I/O
 - **`from_edf_annotations(raw)`** — load hypnogram from EDF+ annotations.
 
-### Analysis
-- **`plot_hypnodensity()`** — when `proba` is available, plot the per-epoch stage probability as a color-map (signature visualization of modern auto-staging papers).
-
 ### Multi-scorer support
 - **`HypnogramSet`** — new container class for multiple scorers of the same night (alignment, pairwise agreement, consensus scoring). See [HYPNOGRAM_MULTIPLE_SCORERS.md](HYPNOGRAM_MULTIPLE_SCORERS.md) for the full design plan.

src/yasa/hypno.py

@@ -93,6 +93,8 @@ class Hypnogram:
          - Compare two hypnograms epoch-by-epoch (kappa, F1, MCC, ...).
        * - :py:meth:`plot_hypnogram`
          - Plot the hypnogram as a standard hypnogram figure.
+       * - :py:meth:`plot_hypnodensity`
+         - Plot per-epoch stage probabilities as a stacked area chart (requires ``proba``).
        * - :py:meth:`simulate_similar`
          - Simulate a new hypnogram with the same transition probabilities as this one.
 
@@ -1666,7 +1668,7 @@ def sleep_statistics(self):
 
         >>> from yasa import simulate_hypnogram
         >>> # Generate a 8 hr (= 480 minutes) 5-stage hypnogram with a 30-seconds resolution
-        >>> hyp = simulate_hypnogram(tib=480, seed=42)
+        >>> hyp = simulate_hypnogram(tib=300, seed=42)
         >>> pd.Series(hyp.sleep_statistics())
         TIB        480.0000
         SPT        477.5000
@@ -1791,7 +1793,7 @@ def transition_matrix(self):
         --------
         >>> from yasa import Hypnogram, simulate_hypnogram
         >>> # Generate a 8 hr (= 480 minutes) 5-stage hypnogram with a 30-seconds resolution
-        >>> hyp = simulate_hypnogram(tib=480, seed=42)
+        >>> hyp = simulate_hypnogram(tib=300, seed=42)
         >>> counts, probs = hyp.transition_matrix()
         >>> counts
         To Stage    WAKE  N1   N2  N3  REM
@@ -1963,10 +1965,167 @@ def plot_hypnogram(self, **kwargs):
         .. plot::
 
             >>> from yasa import simulate_hypnogram
-            >>> ax = simulate_hypnogram(tib=480, seed=88).plot_hypnogram(highlight="REM")
+            >>> ax = simulate_hypnogram(tib=300, seed=88).plot_hypnogram(highlight="REM")
         """
         return plot_hypnogram(self, **kwargs)
 
+    def plot_hypnodensity(self, palette=None, ax=None):
+        """Plot the hypnodensity: per-epoch stage probabilities as a stacked area chart.
+
+        Requires that the :py:attr:`proba` attribute is set (i.e. the hypnogram was created by
+        :py:meth:`yasa.SleepStaging.predict`).
+
+        Parameters
+        ----------
+        palette : dict or None
+            A dictionary mapping stage names to matplotlib colors, e.g.
+            ``{"WAKE": "#99d7f1", "REM": "xkcd:sunflower"}``. When ``None`` (default), a
+            built-in palette is used. Missing stage keys fall back to ``"gray"``.
+        ax : :py:class:`matplotlib.axes.Axes` or None
+            Axis on which to draw the plot. If ``None`` (default), the current axis is used.
+
+        Returns
+        -------
+        ax : :py:class:`matplotlib.axes.Axes`
+            Matplotlib Axes
+
+        Raises
+        ------
+        ValueError
+            If :py:attr:`proba` is ``None``.
+
+        Examples
+        --------
+        5-stage hypnogram:
+
+        .. plot::
+
+            >>> import numpy as np
+            >>> import pandas as pd
+            >>> from yasa import Hypnogram, simulate_hypnogram
+            >>> import matplotlib.pyplot as plt
+            >>> hyp = simulate_hypnogram(tib=300, n_stages=5, seed=42)
+            >>> stages = ["WAKE", "N1", "N2", "N3", "REM"]
+            >>> rng = np.random.default_rng(42)
+            >>> one_hot = (
+            ...     pd.get_dummies(hyp.hypno)
+            ...     .reindex(columns=stages, fill_value=0)
+            ...     .to_numpy(dtype=float)
+            ... )
+            >>> noise = rng.dirichlet(np.ones(5) * 0.5, size=hyp.n_epochs)
+            >>> raw = 0.75 * one_hot + 0.25 * noise
+            >>> proba = pd.DataFrame(raw / raw.sum(axis=1, keepdims=True), columns=stages)
+            >>> ax = Hypnogram(hyp.hypno, n_stages=5, proba=proba).plot_hypnodensity()
+            >>> plt.tight_layout()
+
+        4-stage hypnogram:
+
+        .. plot::
+
+            >>> import numpy as np
+            >>> import pandas as pd
+            >>> from yasa import Hypnogram, simulate_hypnogram
+            >>> import matplotlib.pyplot as plt
+            >>> hyp = simulate_hypnogram(tib=300, n_stages=4, seed=42)
+            >>> stages = ["WAKE", "LIGHT", "DEEP", "REM"]
+            >>> rng = np.random.default_rng(42)
+            >>> one_hot = (
+            ...     pd.get_dummies(hyp.hypno)
+            ...     .reindex(columns=stages, fill_value=0)
+            ...     .to_numpy(dtype=float)
+            ... )
+            >>> noise = rng.dirichlet(np.ones(4) * 0.5, size=hyp.n_epochs)
+            >>> raw = 0.75 * one_hot + 0.25 * noise
+            >>> proba = pd.DataFrame(raw / raw.sum(axis=1, keepdims=True), columns=stages)
+            >>> ax = Hypnogram(hyp.hypno, n_stages=4, proba=proba).plot_hypnodensity()
+            >>> plt.tight_layout()
+
+        2-stage hypnogram:
+
+        .. plot::
+
+            >>> import numpy as np
+            >>> import pandas as pd
+            >>> from yasa import Hypnogram, simulate_hypnogram
+            >>> import matplotlib.pyplot as plt
+            >>> hyp = simulate_hypnogram(tib=300, n_stages=2, seed=42)
+            >>> stages = ["WAKE", "SLEEP"]
+            >>> rng = np.random.default_rng(42)
+            >>> one_hot = (
+            ...     pd.get_dummies(hyp.hypno)
+            ...     .reindex(columns=stages, fill_value=0)
+            ...     .to_numpy(dtype=float)
+            ... )
+            >>> noise = rng.dirichlet(np.ones(2) * 0.5, size=hyp.n_epochs)
+            >>> raw = 0.75 * one_hot + 0.25 * noise
+            >>> proba = pd.DataFrame(raw / raw.sum(axis=1, keepdims=True), columns=stages)
+            >>> ax = Hypnogram(hyp.hypno, n_stages=2, proba=proba).plot_hypnodensity()
+            >>> plt.tight_layout()
+        """
+        import matplotlib.dates as mdates
+        import matplotlib.pyplot as plt
+
+        if self._proba is None:
+            raise ValueError(
+                "No probability data found. `proba` is only available when the Hypnogram "
+                "was created by `yasa.SleepStaging.predict()`."
+            )
+
+        # Default color palette covering all possible stage names.
+        # Base 5-stage colors: WAKE=#99d7f1, N1=#009ddc, N2=#0a437a, N3=#720058, REM=#ffc512
+        # Derived colors: LIGHT=avg(N1,N2), NREM=avg(N1,N2,N3), DEEP=N3, SLEEP=dark navy
+        _default_palette = {
+            "WAKE": "#99d7f1",
+            "N1": "#009ddc",
+            "N2": "#0a437a",
+            "N3": "#720058",
+            "REM": "#ffc512",
+            "LIGHT": "#0570ab",  # avg(N1, N2)
+            "DEEP": "#720058",  # = N3
+            "NREM": "#294b8f",  # avg(N1, N2, N3)
+            "SLEEP": "#003366",  # dark navy, pairs with light-blue WAKE
+            "ART": "#999999",
+            "UNS": "#cccccc",
+        }
+        if palette is None:
+            palette = _default_palette
+
+        proba = self._proba.copy()
+        stages = proba.columns.tolist()
+        colors = [palette.get(s, "gray") for s in stages]
+
+        # Increase font size while preserving original
+        old_fontsize = plt.rcParams["font.size"]
+        plt.rcParams.update({"font.size": 18})
+
+        if ax is None:
+            _, ax = plt.subplots(figsize=(12, 4))
+
+        # Build x-axis values
+        if self._start is not None:
+            times = pd.date_range(start=self._start, freq=self._freq, periods=self._n_epochs)
+            x = mdates.date2num(times)
+            xlabel = "Time"
+        else:
+            x = self._timedelta.total_seconds() / 60  # minutes
+            xlabel = "Time [mins]" if self._duration <= 90 else "Time [hrs]"
+            if self._duration > 90:
+                x = x / 60  # convert to hours
+
+        ax.stackplot(x, proba.to_numpy().T, labels=stages, colors=colors, alpha=0.85)
+        ax.set_xlim(x[0], x[-1])
+        ax.set_ylim(0, 1)
+        ax.set_ylabel("Probability")
+        ax.set_xlabel(xlabel)
+        ax.legend(frameon=False, bbox_to_anchor=(1, 1), loc="upper left")
+        ax.spines[["right", "top"]].set_visible(False)
+        if self._start is not None:
+            ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
+            ax.xaxis.set_major_locator(mdates.AutoDateLocator())
+
+        plt.rcParams.update({"font.size": old_fontsize})
+        return ax
+
     #######################################################################
     # SIMULATION
     #######################################################################
@@ -2566,7 +2725,7 @@ def hypno_find_periods(hypno, sf_hypno, threshold="5min", equal_length=False):
 
 
 def simulate_hypnogram(
-    tib=480,
+    tib=300,
     trans_probas=None,
     init_probas=None,
     seed=None,

tests/test_hypnoclass.py

@@ -750,3 +750,111 @@ def test_pad_negative_after_raises():
 def test_pad_bad_type_raises():
     with pytest.raises(TypeError):
         Hypnogram(_PAD_STAGES).pad(before=3.5)
+
+
+###############################################################################
+# plot_hypnodensity
+###############################################################################
+
+# Shared fixtures for hypnodensity tests
+_N = 100
+_rng = np.random.default_rng(0)
+
+
+def _make_proba(stages):
+    """Return a valid probability DataFrame for the given stage list."""
+    raw = _rng.dirichlet(np.ones(len(stages)), size=_N)
+    return pd.DataFrame(raw, columns=stages)
+
+
+def test_plot_hypnodensity_5stage_returns_axes():
+    proba = _make_proba(["WAKE", "N1", "N2", "N3", "REM"])
+    hyp = Hypnogram(
+        ["WAKE"] * 20 + ["N1"] * 10 + ["N2"] * 40 + ["N3"] * 20 + ["REM"] * 10, proba=proba
+    )
+    ax = hyp.plot_hypnodensity()
+    assert isinstance(ax, plt.Axes)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_4stage_returns_axes():
+    stages_seq = ["WAKE"] * 25 + ["LIGHT"] * 25 + ["DEEP"] * 25 + ["REM"] * 25
+    proba = _make_proba(["WAKE", "LIGHT", "DEEP", "REM"])
+    hyp = Hypnogram(stages_seq, n_stages=4, proba=proba)
+    ax = hyp.plot_hypnodensity()
+    assert isinstance(ax, plt.Axes)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_3stage_returns_axes():
+    stages_seq = ["WAKE"] * 34 + ["NREM"] * 33 + ["REM"] * 33
+    proba = _make_proba(["WAKE", "NREM", "REM"])
+    hyp = Hypnogram(stages_seq, n_stages=3, proba=proba)
+    ax = hyp.plot_hypnodensity()
+    assert isinstance(ax, plt.Axes)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_2stage_returns_axes():
+    stages_seq = ["WAKE"] * 50 + ["SLEEP"] * 50
+    proba = _make_proba(["WAKE", "SLEEP"])
+    hyp = Hypnogram(stages_seq, n_stages=2, proba=proba)
+    ax = hyp.plot_hypnodensity()
+    assert isinstance(ax, plt.Axes)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_no_proba_raises():
+    hyp = Hypnogram(["WAKE"] * 10 + ["N2"] * 90)
+    with pytest.raises(ValueError, match="proba"):
+        hyp.plot_hypnodensity()
+
+
+def test_plot_hypnodensity_with_start_uses_datetime_axis():
+    proba = _make_proba(["WAKE", "N1", "N2", "N3", "REM"])
+    hyp = Hypnogram(
+        ["WAKE"] * 20 + ["N1"] * 10 + ["N2"] * 40 + ["N3"] * 20 + ["REM"] * 10,
+        proba=proba,
+        start="2022-12-15 22:30:00",
+    )
+    ax = hyp.plot_hypnodensity()
+    assert isinstance(ax, plt.Axes)
+    # x-axis should use a DateFormatter when start is set
+    import matplotlib.dates as mdates
+
+    assert isinstance(ax.xaxis.get_major_formatter(), mdates.DateFormatter)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_accepts_ax_argument():
+    proba = _make_proba(["WAKE", "N1", "N2", "N3", "REM"])
+    hyp = Hypnogram(
+        ["WAKE"] * 20 + ["N1"] * 10 + ["N2"] * 40 + ["N3"] * 20 + ["REM"] * 10, proba=proba
+    )
+    fig, ax = plt.subplots()
+    returned_ax = hyp.plot_hypnodensity(ax=ax)
+    assert returned_ax is ax
+    plt.close("all")
+
+
+def test_plot_hypnodensity_custom_palette():
+    proba = _make_proba(["WAKE", "N1", "N2", "N3", "REM"])
+    hyp = Hypnogram(
+        ["WAKE"] * 20 + ["N1"] * 10 + ["N2"] * 40 + ["N3"] * 20 + ["REM"] * 10, proba=proba
+    )
+    custom = {"WAKE": "red", "N1": "green", "N2": "blue", "N3": "purple", "REM": "orange"}
+    ax = hyp.plot_hypnodensity(palette=custom)
+    assert isinstance(ax, plt.Axes)
+    plt.close("all")
+
+
+def test_plot_hypnodensity_ylim_and_legend():
+    proba = _make_proba(["WAKE", "N1", "N2", "N3", "REM"])
+    hyp = Hypnogram(
+        ["WAKE"] * 20 + ["N1"] * 10 + ["N2"] * 40 + ["N3"] * 20 + ["REM"] * 10, proba=proba
+    )
+    ax = hyp.plot_hypnodensity()
+    assert ax.get_ylim() == (0, 1)
+    legend_labels = [t.get_text() for t in ax.get_legend().get_texts()]
+    assert set(legend_labels) == {"WAKE", "N1", "N2", "N3", "REM"}
+    plt.close("all")