compare_channels

Author: alexpiet

src/aind_dynamic_foraging_basic_analysis/plot/plot_fip.py

@@ -3,6 +3,7 @@
 """
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import numpy as np
 
 from aind_dynamic_foraging_data_utils import alignment as an
@@ -116,47 +117,76 @@ def plot_fip_psth_compare_channels(
     data_column="data",
 ):
     """
-    nwb, the nwb object for the session of interest
+    nwb, the nwb object for the session of interest, or a list of nwb objects
     align should either be a string of the name of an event type in nwb.df_events,
-        or a list of timepoints
+        or a list of timepoints. if nwb is a list, then align should be a list containing
+        either the string of the name of an event type, or a list of timepoints.
     channels should be a list of channel names (strings)
     censor, censor important timepoints before and after aligned timepoints
     data_column (string), name of data column in nwb.df_fip
 
     EXAMPLE
     ********************
     plot_fip_psth(nwb, 'goCue_start_time')
+    plot_fip_psth(nwb_list, 'goCue_start_time')
+    plot_fip_psth(nwb_list, ['goCue_start_time','goCue_start_time'])
     """
-    if not hasattr(nwb, "df_fip"):
-        print("You need to compute the df_fip first")
-        print("running `nwb.df_fip = create_fib_df(nwb,tidy=True)`")
-        nwb.df_fip = nu.create_fib_df(nwb, tidy=True)
-    if not hasattr(nwb, "df_events"):
-        print("You need to compute the df_events first")
-        print("run `nwb.df_events = create_events_df(nwb)`")
-        nwb.df_events = nu.create_events_df(nwb)
+    # Check if nwb is a list, otherwise put it in a list to check
+    nwb_list = nwb if isinstance(nwb, list) else [nwb]
 
-    if isinstance(align, str):
-        if align not in nwb.df_events["event"].values:
-            print("{} not found in the events table".format(align))
-            return
-        align_timepoints = nwb.df_events.query("event == @align")["timestamps"].values
-        align_label = "Time from {} (s)".format(align)
-    else:
-        align_timepoints = align
-        align_label = "Time (s)"
+    if isinstance(nwb, list) and isinstance(align, list) and (len(nwb) != len(align)):
+        raise Exception("NWB list and align list must match")
+    if isinstance(nwb, list) and isinstance(align, str):
+        align = [align] * len(nwb)
+    if not isinstance(nwb, list):
+        align = [align]
 
+    # First check that each session has an events table and fip table
+    for nwb_i in nwb_list:
+        if not hasattr(nwb_i, "df_fip"):
+            print("You need to compute the df_fip first")
+            print("running `nwb.df_fip = create_fib_df(nwb,tidy=True)`")
+            nwb_i.df_fip = nu.create_fib_df(nwb_i, tidy=True)
+        if not hasattr(nwb_i, "df_events"):
+            print("You need to compute the df_events first")
+            print("run `nwb.df_events = create_events_df(nwb)`")
+            nwb_i.df_events = nu.create_events_df(nwb_i)
+
+    align_timepoints_list = []
+    # Generate the alignment timepoints for each session
+    for i, nwb_i in enumerate(nwb_list):
+        align_i = align[i]
+        if isinstance(align_i, str):
+            if align_i not in nwb_i.df_events["event"].values:
+                print("{} not found in the events table, {}".format(align_i, nwb_i.session_id))
+                return
+
+            align_timepoints_list.append(
+                nwb_i.df_events.query("event == @align")["timestamps"].values
+            )
+            align_label = "Time from {} (s)".format(align_i)
+        else:
+            align_timepoints_list.append(align_i)
+            align_label = "Time (s)"
+
+    # Make figure if not supplied
     if fig is None and ax is None:
         fig, ax = plt.subplots()
 
+    # Iterate through channels and plot
     colors = [FIP_COLORS.get(c, "") for c in channels]
     for dex, c in enumerate(channels):
-        if c in nwb.df_fip["event"].values:
-            etr = fip_psth_inner_compute(nwb, align_timepoints, c, True, tw,
-                                         censor, data_column=data_column)
-            fip_psth_inner_plot(ax, etr, colors[dex], c, data_column)
-        else:
-            print("No data for channel: {}".format(c))
+        include = [c in nwb.df_fip["event"].values for nwb in nwb_list]
+        etr = fip_psth_multiple_inner_compute(
+            [x for dex, x in enumerate(nwb_list) if include[dex]],
+            [x for dex, x in enumerate(align_timepoints_list) if include[dex]],
+            c,
+            True,
+            tw,
+            censor,
+            data_column=data_column,
+        )
+        fip_psth_inner_plot(ax, etr, colors[dex], c, data_column)
 
     plt.legend()
     ax.set_xlabel(align_label, fontsize=STYLE["axis_fontsize"])
@@ -166,7 +196,10 @@ def plot_fip_psth_compare_channels(
     ax.set_xlim(tw)
     ax.axvline(0, color="k", alpha=0.2)
     ax.tick_params(axis="both", labelsize=STYLE["axis_ticks_fontsize"])
-    ax.set_title(nwb.session_id)
+    if len(nwb_list) == 1:
+        ax.set_title(nwb_list[0].session_id)
+    else:
+        ax.set_title("{} sessions".format(len(nwb_list)))
     plt.tight_layout()
     return fig, ax
 
@@ -183,11 +216,69 @@ def fip_psth_inner_plot(ax, etr, color, label, data_column):
     if color == "":
         cmap = plt.get_cmap("tab20")
         color = cmap(np.random.randint(20))
-    ax.fill_between(etr.index, etr[data_column] - etr["sem"],
-                    etr[data_column] + etr["sem"], color=color, alpha=0.2)
+    ax.fill_between(
+        etr.index,
+        etr[data_column] - etr["sem"],
+        etr[data_column] + etr["sem"],
+        color=color,
+        alpha=0.2,
+    )
     ax.plot(etr.index, etr[data_column], color=color, label=label)
 
 
+def fip_psth_multiple_inner_compute(
+    nwb_list,
+    align_timepoints_list,
+    channel,
+    average,
+    tw=[-1, 1],
+    censor=True,
+    censor_times=None,
+    data_column="data",
+):
+    """ """
+    # Check that len(nwb_list) = len(align_timepoints_list) = len(censor_times)
+    if len(nwb_list) != len(align_timepoints_list):
+        raise Exception("length of nwb list and alignments list must match")
+    if censor and censor_times is None:
+        censor_times = [None] * len(nwb_list)
+    if censor and (len(nwb_list) != len(censor_times)):
+        raise Exception("length of nwb list and censor times must match")
+
+    etr_list = []
+    # Iterate through list of sessions, computing the etr for each
+    for i, nwb_i in enumerate(nwb_list):
+        etr_i = fip_psth_inner_compute(
+            nwb_i,
+            align_timepoints_list[i],
+            channel,
+            average=False,
+            tw=tw,
+            censor=censor,
+            censor_times=censor_times[i],
+            data_column=data_column,
+        )
+        etr_i["ses_idx"] = nwb_i.session_id
+        etr_list.append(etr_i)
+
+    # Concat etrs from each session into one dataframe
+    etr_all = pd.concat(etr_list, axis=0).reset_index(drop=True)
+
+    if average:
+        # Average within each ses_idx for each time point
+        mean_per_ses = etr_all.groupby(["ses_idx", "time"])[data_column].mean().unstack("ses_idx")
+        # Grand mean: average across ses_idx for each time point
+        grand_mean = mean_per_ses.mean(axis=1)
+        # SEM over ses_idx for each time point
+        grand_sem = mean_per_ses.sem(axis=1)
+        # Combine into a DataFrame
+        result = grand_mean.to_frame(name=data_column)
+        result["sem"] = grand_sem
+        return result
+    else:
+        return etr_all
+
+
 def fip_psth_inner_compute(
     nwb,
     align_timepoints,