adding function for plotting specific interval lists side by side (#1330)

* adding function for plotting specific interval lists side by side

* incoporating as a class method and adding testing

* update time units and docstring

* add error and warning related to input intervallist

* Fix linting

* Add typing and improve docstring

---------

Co-authored-by: Eric Denovellis <[email protected]>
Co-authored-by: Eric Denovellis <[email protected]>

Author: gshvarts

src/spyglass/common/common_interval.py

@@ -1,8 +1,10 @@
 import itertools
+import warnings
 from functools import reduce
 from typing import Iterable, List, Optional, Tuple, TypeVar, Union
 
 import datajoint as dj
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -79,25 +81,83 @@ def fetch_interval(self):
             raise ValueError(f"Expected one row, got {len(self)}")
         return Interval(self.fetch1())
 
-    def plot_intervals(self, figsize=(20, 5), return_fig=False):
-        """Plot the intervals in the interval list."""
-        interval_list = pd.DataFrame(self)
-        fig, ax = plt.subplots(figsize=figsize)
-        interval_count = 0
-        for row in interval_list.itertuples(index=False):
-            for interval in row.valid_times:
-                ax.plot(interval, [interval_count, interval_count])
-                ax.scatter(
-                    interval,
-                    [interval_count, interval_count],
-                    alpha=0.8,
-                    zorder=2,
-                )
-            interval_count += 1
-        ax.set_yticks(np.arange(interval_list.shape[0]))
-        ax.set_yticklabels(interval_list.interval_list_name)
-        ax.set_xlabel("Time [s]")
-        ax.grid(True)
+    def plot_intervals(
+        self, start_time: float = 0, return_fig: bool = False
+    ) -> Optional[plt.Figure]:
+        """
+        Plot all intervals in the given IntervalList table.
+
+        Parameters
+        ----------
+        start_time : float, optional
+            The reference time (in seconds) for the interval comparison plot.
+            For example, the first timepoint of a session. Defaults to 0.
+        return_fig : bool, optional
+            If True, return the matplotlib Figure object. Defaults to False.
+
+        Returns
+        -------
+        fig : matplotlib.figure.Figure or None
+            The matplotlib Figure object if `return_fig` is True, otherwise None.
+
+        Raises
+        ------
+        ValueError
+            If more than one unique `nwb_file_name` is found in the IntervalList.
+            The intended use is to compare intervals within a single NWB file.
+        UserWarning
+            If more than 100 intervals are being plotted.
+        """
+        interval_lists_df = pd.DataFrame(self)
+
+        if len(interval_lists_df["nwb_file_name"].unique()) > 1:
+            raise ValueError(
+                ">1 nwb_file_name found in IntervalList. the intended use of plot_intervals is to compare intervals within a single nwb_file_name."
+            )
+
+        interval_list_names = interval_lists_df["interval_list_name"].values
+
+        n_compare = len(interval_list_names)
+
+        if n_compare > 100:
+            warnings.warn(
+                f"plot_intervals is plotting {n_compare} intervals. if this is unintended, please pass in a smaller IntervalList.",
+                UserWarning,
+            )
+
+        # plot broken bar horizontals
+        fig, ax = plt.subplots(figsize=(20, 2 / 3 * n_compare))
+
+        # get n colors
+        cmap = plt.get_cmap("turbo", n_compare)
+        custom_palette = [mpl.colors.rgb2hex(cmap(i)) for i in range(cmap.N)]
+
+        def convert_intervals_to_range(intervals, start_time):
+            return [
+                ((i[0] - start_time) / 60, (i[1] - i[0]) / 60)
+                for i in intervals
+            ]  # return time in minutes
+
+        all_intervals = interval_lists_df["valid_times"].values
+        for i, (intervals, color) in enumerate(
+            zip(all_intervals, custom_palette)
+        ):
+            int_range = convert_intervals_to_range(intervals, start_time)
+            ax.broken_barh(
+                int_range, (10 * (i + 1), 6), facecolors=color, alpha=0.7
+            )
+
+        ax.set_ylim(5, 10 * (n_compare + 1) + 5)
+        ax.set_xlabel("time from start (min)", fontsize=16)
+        ax.set_yticks(
+            np.arange(n_compare) * 10 + 15,
+            labels=interval_list_names,
+            fontsize=16,
+        )
+        ax.set_xticks(ax.get_xticks())
+        ax.set_xticklabels(ax.get_xticklabels())
+        ax.tick_params(axis="x", labelsize=16)
+
         if return_fig:
             return fig
 

tests/common/test_interval.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pytest
 from numpy import array_equal
 
@@ -7,17 +8,56 @@ def interval_list(common):
     yield common.IntervalList()
 
 
-def test_plot_intervals(mini_insert, interval_list):
-    fig = (interval_list & 'interval_list_name LIKE "raw%"').plot_intervals(
-        return_fig=True
+def test_plot_intervals(interval_list, start_time=0):
+    fig = (interval_list).plot_intervals(return_fig=True, start_time=start_time)
+
+    ax = fig.axes[0]
+
+    # extract interval list names from the figure
+    fig_interval_list_names = np.array(
+        [label.get_text() for label in ax.get_yticklabels()]
+    )
+
+    # extract interval list names from the IntervalList table
+    fetch_interval_list_names = np.array(
+        interval_list.fetch("interval_list_name")
     )
-    interval_list_name = fig.get_axes()[0].get_yticklabels()[0].get_text()
-    times_fetch = (
-        interval_list & {"interval_list_name": interval_list_name}
-    ).fetch1("valid_times")[0]
-    times_plot = fig.get_axes()[0].lines[0].get_xdata()
 
-    assert array_equal(times_fetch, times_plot), "plot_intervals failed"
+    # check that the interval list names match between the two methods
+    assert array_equal(
+        fig_interval_list_names, fetch_interval_list_names
+    ), "plot_intervals failed: plotted interval list names do not match"
+
+    # extract the interval times from the figure
+    intervals_fig = []
+    for collection in ax.collections:
+        collection_data = []
+        for patch in collection.get_paths():
+            # Extract patch vertices to get x data
+            vertices = patch.vertices
+            x_start = vertices[0, 0]
+            x_end = vertices[2, 0]
+            collection_data.append(
+                [x_start * 60 + start_time, x_end * 60 + start_time]
+            )
+        intervals_fig.append(np.array(collection_data))
+    intervals_fig = np.array(intervals_fig, dtype="object")
+
+    # extract interval times from the IntervalList table
+    intervals_fetch = interval_list.fetch("valid_times")
+
+    all_equal = True
+    for i_fig, i_fetch in zip(intervals_fig, intervals_fetch):
+        # permit rounding errors up to the 4th decimal place, as a result of inaccuracies during unit conversions
+        i_fig = np.round(i_fig.astype("float"), 4)
+        i_fetch = np.round(i_fetch.astype("float"), 4)
+        if not array_equal(i_fig, i_fetch):
+            all_equal = False
+
+    # check that the interval list times match between the two methods
+    assert (
+        all_equal
+    ), "plot_intervals failed: plotted interval times do not match"
 
 
 def test_plot_epoch(mini_insert, interval_list):