From bdb1efaa36dd08a6094cdd87e90c5f35ec23f280 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Mon, 29 Jun 2026 23:42:14 +0000
Subject: [PATCH 01/14] small fixes to plotly code

---
 .../plot/plot_foraging_session_plotly.py      | 79 ++++++++++++++-----
 1 file changed, 58 insertions(+), 21 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index d3ce910..7e54197 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -92,6 +92,7 @@ def _broken(x, y, segments):
     return xs, ys
 
 
+
 def plot_foraging_session_plotly(  # noqa: C901 pragma: no cover
     choice_history,
     reward_history,
@@ -417,9 +418,41 @@ def _smoothed_trace(num, den):
     )
     return fig
 
+def plot_foraging_session_nwb_plotly(nwb, **kwargs):
+    """
+    Wrapper function that extracts fields
+    """
+
+    if not hasattr(nwb, "df_trials"):
+        print("You need to compute df_trials: nwb_utils.create_trials_df(nwb)")
+        return
+
+    if "side_bias" not in nwb.df_trials:
+        fig, axes = plot_foraging_session_plotly(
+            [np.nan if x == 2 else x for x in nwb.df_trials["animal_response"].values],
+            nwb.df_trials["earned_reward"].values,
+            [nwb.df_trials["reward_probabilityL"], nwb.df_trials["reward_probabilityR"]],
+            **kwargs,
+        )
+    else:
+        if "plot_list" not in kwargs:
+            kwargs["plot_list"] = ["choice", "reward_prob", "bias"]
+        fig, axes = plot_foraging_session_plotly(
+            [np.nan if x == 2 else x for x in nwb.df_trials["animal_response"].values],
+            nwb.df_trials["earned_reward"].values,
+            [nwb.df_trials["reward_probabilityL"], nwb.df_trials["reward_probabilityR"]],
+            bias=nwb.df_trials["side_bias"].values,
+            bias_lower=[x[0] for x in nwb.df_trials["side_bias_confidence_interval"].values],
+            bias_upper=[x[1] for x in nwb.df_trials["side_bias_confidence_interval"].values],
+            autowater_offered=nwb.df_trials[["auto_waterL", "auto_waterR"]].any(axis=1),
+            **kwargs,
+        )
+
+
+
 
 def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
-    df_events, df_trials=None, fip_df=None, adjust_time=True, title=None, smooth_factor=5
+    nwb, fip = [], adjust_time=True, title=None, smooth_factor=5
 ):
     """Plotly version of :func:`plot_session_scroller.plot_session_scroller` (time-based).
 
@@ -436,18 +469,9 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
 
     Parameters
     ----------
-    df_events : pandas.DataFrame
-        Tidy dataframe of session events (``event`` + ``timestamps``; optional ``session_id``).
-        Recognised events: ``left_lick_time``, ``right_lick_time``, ``left_reward_delivery_time``,
-        ``right_reward_delivery_time`` and ``goCue_start_time``.
-    df_trials : pandas.DataFrame, optional
-        Per-trial dataframe for the reward-probability band / overlays / red ignored go cues
-        (and a fallback source of go-cue times). Uses ``goCue_start_time``,
-        ``reward_probabilityL/R`` and ``animal_response``; go-cue times must share the
-        ``df_events`` time base. Matched per session via ``session_id`` when present.
-    fip_df : pandas.DataFrame, optional
-        Tidy FIP measurements (single-session only); each present channel is normalised and
-        stacked above the behavior panel.
+    nwb, an nwb like object that contains attributes: df_events, session_id
+        and optionally contains attributes fip_df, df_licks
+    fip (list), FIP channels to plot. Must be present in nwb.df_fip
     adjust_time : bool, optional
         If True (default), shift time so the first event is at t = 0 (always shifted when
         concatenating multiple sessions).
@@ -460,11 +484,25 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     -------
     plotly.graph_objects.Figure
     """
-    df_events = df_events.copy()
-    if df_trials is not None:
-        df_trials = df_trials.copy()
-    if fip_df is not None:
-        fip_df = fip_df.copy()
+    if not hasattr(nwb, "df_events"):
+        print("computing df_events first")
+        nwb.df_events = nu.create_df_events(nwb)
+        df_events = nwb.df_events
+    else:
+        df_events = nwb.df_events
+    if hasattr(nwb, "df_fip"):
+        fip_df = nwb.df_fip
+    if hasattr(nwb, "df_licks"):
+        df_licks = nwb.df_licks
+    else:
+        df_licks = None
+    if not hasattr(nwb, "df_trials"):
+        print("computing df_trials")
+        nwb.df_trials = nu.create_df_trials(nwb)
+        df_trials = nwb.df_trials
+    else:
+        df_trials = nwb.df_trials
+
 
     # y-layout, bottom -> top:
     #   * event rows in [0, 1]: rewards at the outer edges, licks inside (right pair near the
@@ -657,9 +695,8 @@ def _trial_of(times, _gc=gc, _n=n_tr):
     y_main_top = params["curve_top"]
 
     # FIP channels (single-session only), normalised and stacked above the behavior panel
-    if fip_df is not None and len(sessions) == 1:
-        fip_channels = ["G_1_preprocessed", "G_2_preprocessed",
-                        "R_1_preprocessed", "R_2_preprocessed"]
+    if fip_df is not None and len(sessions) == 1 and len(fip) > 0:
+        fip_channels = fip
         fip_colors = {"G_1": "green", "G_2": "darkgreen", "R_1": "red", "R_2": "darkred"}
         present = set(fip_df["event"].unique())
         band = 0

From 253c4b5f5a920e5f47d58732feca5acd1406e99b Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Mon, 29 Jun 2026 23:43:23 +0000
Subject: [PATCH 02/14] linting

---
 .../plot/plot_foraging_session_plotly.py      | 434 +++++++++++++-----
 1 file changed, 315 insertions(+), 119 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index 7e54197..ddae2b6 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -92,7 +92,6 @@ def _broken(x, y, segments):
     return xs, ys
 
 
-
 def plot_foraging_session_plotly(  # noqa: C901 pragma: no cover
     choice_history,
     reward_history,
@@ -236,53 +235,95 @@ def _markers(mask):
     # Rewarded (real foraging, autowater excluded): tall black ticks just outside [0, 1]
     xs, ys, cd = _raster(rewarded_excluding_autowater, (1.05, 1.15), (-0.15, -0.05))
     fig.add_trace(
-        go.Scattergl(x=xs, y=ys, customdata=cd, mode="lines", line=dict(color="black", width=1),
-                     name="Rewarded choices", hovertemplate=hovertemplate % "Rewarded choices"),
-        row=1, col=1,
+        go.Scattergl(
+            x=xs,
+            y=ys,
+            customdata=cd,
+            mode="lines",
+            line=dict(color="black", width=1),
+            name="Rewarded choices",
+            hovertemplate=hovertemplate % "Rewarded choices",
+        ),
+        row=1,
+        col=1,
     )
 
     # Unrewarded (real foraging): short gray ticks
     xs, ys, cd = _raster(unrewarded_trials, (1.05, 1.10), (-0.10, -0.05))
     fig.add_trace(
-        go.Scattergl(x=xs, y=ys, customdata=cd, mode="lines", line=dict(color="gray", width=1),
-                     name="Unrewarded choices", hovertemplate=hovertemplate % "Unrewarded choices"),
-        row=1, col=1,
+        go.Scattergl(
+            x=xs,
+            y=ys,
+            customdata=cd,
+            mode="lines",
+            line=dict(color="gray", width=1),
+            name="Unrewarded choices",
+            hovertemplate=hovertemplate % "Unrewarded choices",
+        ),
+        row=1,
+        col=1,
     )
 
     # Ignored trials: red x at the top
     xx, cd = _markers(ignored & ~autowater_ignored)
     fig.add_trace(
-        go.Scattergl(x=xx, y=[1.2] * len(xx), customdata=cd, mode="markers",
-                     marker=dict(symbol="x", color="red", size=4), name="Ignored",
-                     hovertemplate=hovertemplate % "Ignored"),
-        row=1, col=1,
+        go.Scattergl(
+            x=xx,
+            y=[1.2] * len(xx),
+            customdata=cd,
+            mode="markers",
+            marker=dict(symbol="x", color="red", size=4),
+            name="Ignored",
+            hovertemplate=hovertemplate % "Ignored",
+        ),
+        row=1,
+        col=1,
     )
 
     # Autowater collected / ignored
     if autowater_offered is not None:
         xs, ys, cd = _raster(autowater_collected, (1.05, 1.15), (-0.15, -0.05))
         fig.add_trace(
-            go.Scattergl(x=xs, y=ys, customdata=cd, mode="lines",
-                         line=dict(color="royalblue", width=1), name="Autowater collected",
-                         hovertemplate=hovertemplate % "Autowater collected"),
-            row=1, col=1,
+            go.Scattergl(
+                x=xs,
+                y=ys,
+                customdata=cd,
+                mode="lines",
+                line=dict(color="royalblue", width=1),
+                name="Autowater collected",
+                hovertemplate=hovertemplate % "Autowater collected",
+            ),
+            row=1,
+            col=1,
         )
         xx, cd = _markers(autowater_ignored)
         fig.add_trace(
-            go.Scattergl(x=xx, y=[1.2] * len(xx), customdata=cd, mode="markers",
-                         marker=dict(symbol="x", color="royalblue", size=4),
-                         name="Autowater ignored",
-                         hovertemplate=hovertemplate % "Autowater ignored"),
-            row=1, col=1,
+            go.Scattergl(
+                x=xx,
+                y=[1.2] * len(xx),
+                customdata=cd,
+                mode="markers",
+                marker=dict(symbol="x", color="royalblue", size=4),
+                name="Autowater ignored",
+                hovertemplate=hovertemplate % "Autowater ignored",
+            ),
+            row=1,
+            col=1,
         )
 
     # Base reward probability (broken at session boundaries)
     if "reward_prob" in plot_list:
         xs, ys = _broken(np.arange(n_trials) + 1, p_reward_fraction, segments)
         fig.add_trace(
-            go.Scattergl(x=xs, y=ys, mode="lines",
-                         line=dict(color=_color(base_color), width=1.5), name="Base rew. prob."),
-            row=1, col=1,
+            go.Scattergl(
+                x=xs,
+                y=ys,
+                mode="lines",
+                line=dict(color=_color(base_color), width=1.5),
+                name="Base rew. prob.",
+            ),
+            row=1,
+            col=1,
         )
 
     def _smoothed_trace(num, den):
@@ -301,18 +342,30 @@ def _smoothed_trace(num, den):
     if "choice" in plot_list:
         xs, ys = _smoothed_trace(choice_history, ~np.isnan(choice_history))
         fig.add_trace(
-            go.Scattergl(x=xs, y=ys, mode="lines", line=dict(color="black", width=1.5),
-                         name=f"Choice (smooth = {smooth_factor})"),
-            row=1, col=1,
+            go.Scattergl(
+                x=xs,
+                y=ys,
+                mode="lines",
+                line=dict(color="black", width=1.5),
+                name=f"Choice (smooth = {smooth_factor})",
+            ),
+            row=1,
+            col=1,
         )
 
     # Finished ratio (only meaningful if there are ignored trials)
     if "finished" in plot_list and np.sum(np.isnan(choice_history)):
         xs, ys = _smoothed_trace(~np.isnan(choice_history), None)
         fig.add_trace(
-            go.Scattergl(x=xs, y=ys, mode="lines", line=dict(color="magenta", width=0.8),
-                         name=f"Finished (smooth = {smooth_factor})"),
-            row=1, col=1,
+            go.Scattergl(
+                x=xs,
+                y=ys,
+                mode="lines",
+                line=dict(color="magenta", width=0.8),
+                name=f"Finished (smooth = {smooth_factor})",
+            ),
+            row=1,
+            col=1,
         )
 
     # Bias trace + confidence band (broken at session boundaries)
@@ -326,20 +379,37 @@ def _smoothed_trace(num, den):
         xb, y_bias = _broken(xx, bias, segments)
         # go.Scatter (not Scattergl) for the filled band -- Scattergl ignores fill.
         fig.add_trace(
-            go.Scatter(x=xb_up, y=y_up, mode="lines", line=dict(width=0),
-                       showlegend=False, hoverinfo="skip"),
-            row=1, col=1,
+            go.Scatter(
+                x=xb_up,
+                y=y_up,
+                mode="lines",
+                line=dict(width=0),
+                showlegend=False,
+                hoverinfo="skip",
+            ),
+            row=1,
+            col=1,
         )
         fig.add_trace(
-            go.Scatter(x=xb_lo, y=y_lo, mode="lines", line=dict(width=0),
-                       fill="tonexty", fillcolor="rgba(0,128,0,0.25)",
-                       showlegend=False, hoverinfo="skip"),
-            row=1, col=1,
+            go.Scatter(
+                x=xb_lo,
+                y=y_lo,
+                mode="lines",
+                line=dict(width=0),
+                fill="tonexty",
+                fillcolor="rgba(0,128,0,0.25)",
+                showlegend=False,
+                hoverinfo="skip",
+            ),
+            row=1,
+            col=1,
         )
         fig.add_trace(
-            go.Scattergl(x=xb, y=y_bias, mode="lines", line=dict(color="green", width=1.5),
-                         name="bias"),
-            row=1, col=1,
+            go.Scattergl(
+                x=xb, y=y_bias, mode="lines", line=dict(color="green", width=1.5), name="bias"
+            ),
+            row=1,
+            col=1,
         )
 
     # Valid (engaged) range
@@ -350,9 +420,15 @@ def _smoothed_trace(num, den):
     # Fitted model overlay
     if fitted_data is not None:
         fig.add_trace(
-            go.Scattergl(x=np.arange(n_trials), y=fitted_data, mode="lines",
-                         line=dict(width=1.5), name="model"),
-            row=1, col=1,
+            go.Scattergl(
+                x=np.arange(n_trials),
+                y=fitted_data,
+                mode="lines",
+                line=dict(width=1.5),
+                name="model",
+            ),
+            row=1,
+            col=1,
         )
 
     # Photostim markers
@@ -364,27 +440,35 @@ def _smoothed_trace(num, den):
         else:
             colors = "darkcyan"
         fig.add_trace(
-            go.Scattergl(x=trial, y=np.ones_like(trial, dtype=float) + 0.4, mode="markers",
-                         marker=dict(symbol="triangle-down", size=power * 2,
-                                     color="rgba(0,0,0,0)",
-                                     line=dict(color=colors, width=0.5)),
-                         name="photostim"),
-            row=1, col=1,
+            go.Scattergl(
+                x=trial,
+                y=np.ones_like(trial, dtype=float) + 0.4,
+                mode="markers",
+                marker=dict(
+                    symbol="triangle-down",
+                    size=power * 2,
+                    color="rgba(0,0,0,0)",
+                    line=dict(color=colors, width=0.5),
+                ),
+                name="photostim",
+            ),
+            row=1,
+            col=1,
         )
 
     # == Reward schedule (bottom panel; broken at session boundaries) ==
     xx = np.arange(n_trials) + 1
     xr, y_pr = _broken(xx, p_reward[1, :], segments)
     fig.add_trace(
-        go.Scattergl(x=xr, y=y_pr, mode="lines", line=dict(color="blue", width=1),
-                     name="p_right"),
-        row=2, col=1,
+        go.Scattergl(x=xr, y=y_pr, mode="lines", line=dict(color="blue", width=1), name="p_right"),
+        row=2,
+        col=1,
     )
     xl, y_pl = _broken(xx, p_reward[0, :], segments)
     fig.add_trace(
-        go.Scattergl(x=xl, y=y_pl, mode="lines", line=dict(color="red", width=1),
-                     name="p_left"),
-        row=2, col=1,
+        go.Scattergl(x=xl, y=y_pl, mode="lines", line=dict(color="red", width=1), name="p_left"),
+        row=2,
+        col=1,
     )
 
     # Thick vertical lines marking session boundaries (between trials b and b+1)
@@ -394,14 +478,19 @@ def _smoothed_trace(num, den):
 
     # Axes styling to match the matplotlib version
     fig.update_yaxes(
-        tickvals=[0, 1, 1.2], ticktext=["Left", "Right", "Ignored"],
-        range=[-0.15, 1.25], fixedrange=True, row=1, col=1,
+        tickvals=[0, 1, 1.2],
+        ticktext=["Left", "Right", "Ignored"],
+        range=[-0.15, 1.25],
+        fixedrange=True,
+        row=1,
+        col=1,
     )
     fig.update_yaxes(title_text="p_reward", range=[0, 1], fixedrange=True, row=2, col=1)
     # Bottom x-axis: a rangeslider scroller (drag to pan/zoom), and -- for multiple sessions --
     # tick labels that restart at 0 each session.
-    fig.update_xaxes(title_text="Trial number", row=2, col=1,
-                     rangeslider=dict(visible=True, thickness=0.08))
+    fig.update_xaxes(
+        title_text="Trial number", row=2, col=1, rangeslider=dict(visible=True, thickness=0.08)
+    )
     if len(segments) > 1:
         step = 250
         tickvals, ticktext = [], []
@@ -411,13 +500,17 @@ def _smoothed_trace(num, den):
                 ticktext.append(str(w))
         fig.update_xaxes(tickvals=tickvals, ticktext=ticktext, row=2, col=1)
     fig.update_layout(
-        width=1000, height=460, template="simple_white",
-        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0,
-                    font=dict(size=9)),
+        width=1000,
+        height=460,
+        template="simple_white",
+        legend=dict(
+            orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0, font=dict(size=9)
+        ),
         margin=dict(l=60, r=20, t=90, b=50),
     )
     return fig
 
+
 def plot_foraging_session_nwb_plotly(nwb, **kwargs):
     """
     Wrapper function that extracts fields
@@ -449,10 +542,8 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
         )
 
 
-
-
 def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
-    nwb, fip = [], adjust_time=True, title=None, smooth_factor=5
+    nwb, fip=[], adjust_time=True, title=None, smooth_factor=5
 ):
     """Plotly version of :func:`plot_session_scroller.plot_session_scroller` (time-based).
 
@@ -503,18 +594,18 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     else:
         df_trials = nwb.df_trials
 
-
     # y-layout, bottom -> top:
     #   * event rows in [0, 1]: rewards at the outer edges, licks inside (right pair near the
     #     top, left pair near the bottom), like the trial-based figure.
     #   * smoothed overlays in their own band [curve_bottom, curve_top] above the events.
     #   * the reward-probability band sits higher still -- it only shows in the rangeslider.
     params = {
-        "behavior_bottom": 0.0, "behavior_top": 1.0,  # event rows (row 1)
-        "curve_bottom": 1.1, "curve_top": 2.1,        # smoothed overlays (row 1)
+        "behavior_bottom": 0.0,
+        "behavior_top": 1.0,  # event rows (row 1)
+        "curve_bottom": 1.1,
+        "curve_top": 2.1,  # smoothed overlays (row 1)
     }
-    row_centers = {"right_reward": 0.92, "right_lick": 0.78,
-                   "left_lick": 0.22, "left_reward": 0.08}
+    row_centers = {"right_reward": 0.92, "right_lick": 0.78, "left_lick": 0.22, "left_reward": 0.08}
     tick_half = 0.25 * 0.30 / 2.0
 
     def _to_curve(v):
@@ -522,20 +613,27 @@ def _to_curve(v):
         span = params["curve_top"] - params["curve_bottom"]
         return params["curve_bottom"] + np.asarray(v, dtype=float) * span
 
-    yticks = [0.92, 0.78, 0.22, 0.08,
-              params["curve_bottom"], (params["curve_bottom"] + params["curve_top"]) / 2,
-              params["curve_top"]]
+    yticks = [
+        0.92,
+        0.78,
+        0.22,
+        0.08,
+        params["curve_bottom"],
+        (params["curve_bottom"] + params["curve_top"]) / 2,
+        params["curve_top"],
+    ]
     ylabels = ["right reward", "right lick", "left lick", "left reward", "0", "0.5", "1"]
 
     # Sessions in order of appearance; concatenate end-to-end along time when more than one.
     has_sess = "session_id" in df_events.columns
-    sessions = (list(dict.fromkeys(df_events["session_id"].tolist())) if has_sess else [None])
+    sessions = list(dict.fromkeys(df_events["session_id"].tolist())) if has_sess else [None]
     shift_each = adjust_time or len(sessions) > 1
 
     # Same two-panel layout as the trial-based figure: the raster/curves on top (row 1) over a
     # reward-schedule panel (row 2), with a rangeslider scroller under row 2.
-    fig = make_subplots(rows=2, cols=1, shared_xaxes=True,
-                        row_heights=[0.85, 0.15], vertical_spacing=0.04)
+    fig = make_subplots(
+        rows=2, cols=1, shared_xaxes=True, row_heights=[0.85, 0.15], vertical_spacing=0.04
+    )
 
     ev_meta = {
         "left_lick": ("left_lick_time", "gray", 1.5, "left lick"),
@@ -544,8 +642,10 @@ def _to_curve(v):
         "right_reward": ("right_reward_delivery_time", "black", 2, "right reward"),
     }
     ev_acc = {k: {"x": [], "y": [], "cd": []} for k in ev_meta}
-    gocue_acc = {"go cue": {"x": [], "y": [], "cd": []},
-                 "go cue (ignored)": {"x": [], "y": [], "cd": []}}
+    gocue_acc = {
+        "go cue": {"x": [], "y": [], "cd": []},
+        "go cue (ignored)": {"x": [], "y": [], "cd": []},
+    }
     frac_x, frac_y, choice_x, choice_y, lick_x, lick_y = [], [], [], [], [], []
     probL_x, probL_y, probR_x, probR_y = [], [], [], []  # reward-prob lines (scroller)
     boundaries, sess_spans, has_prob = [], [], False
@@ -568,8 +668,11 @@ def _ev(name, _ev=ev_s, _off=off):
 
         tr_s = None
         if df_trials is not None:
-            tr_s = (df_trials[df_trials["session_id"] == sess]
-                    if (sess is not None and "session_id" in df_trials.columns) else df_trials)
+            tr_s = (
+                df_trials[df_trials["session_id"] == sess]
+                if (sess is not None and "session_id" in df_trials.columns)
+                else df_trials
+            )
 
         gc = _ev("goCue_start_time")
         if len(gc) == 0 and tr_s is not None and "goCue_start_time" in tr_s.columns:
@@ -608,8 +711,13 @@ def _trial_of(times, _gc=gc, _n=n_tr):
             for gname, mask in [("go cue", ~ign), ("go cue (ignored)", ign)]:
                 if mask.any():
                     hov = [(int(tr), sess_disp) for tr in trial_no[mask]]
-                    xs, ys, cd = _vline_hover(gc[mask], params["behavior_bottom"],
-                                              params["behavior_top"], hov, gap=(None, None))
+                    xs, ys, cd = _vline_hover(
+                        gc[mask],
+                        params["behavior_bottom"],
+                        params["behavior_top"],
+                        hov,
+                        gap=(None, None),
+                    )
                     gocue_acc[gname]["x"] += xs
                     gocue_acc[gname]["y"] += ys
                     gocue_acc[gname]["cd"] += cd
@@ -624,26 +732,31 @@ def _trial_of(times, _gc=gc, _n=n_tr):
                 frac_y += [*_to_curve(frac), None]
             if choice is not None:
                 sm = moving_average(choice, smooth_factor) / (
-                    moving_average(~np.isnan(choice), smooth_factor) + 1e-6)
+                    moving_average(~np.isnan(choice), smooth_factor) + 1e-6
+                )
                 sm[sm > 100] = np.nan
-                xsm = gc[off_s: off_s + len(sm)]
+                xsm = gc[off_s : off_s + len(sm)]
                 choice_x += [*xsm, None]
                 choice_y += [*_to_curve(sm[: len(xsm)]), None]
             lt = np.concatenate([_ev("left_lick_time"), _ev("right_lick_time")])
             if len(lt):
-                counts = np.bincount(_trial_of(lt), minlength=n_tr + 1)[1:n_tr + 1]
+                counts = np.bincount(_trial_of(lt), minlength=n_tr + 1)[1 : n_tr + 1]
                 sm = moving_average(counts.astype(float), smooth_factor)
                 top = np.nanmax(sm) if len(sm) else 0
                 if top > 0:
                     sm = sm / top
-                xsm = gc[off_s: off_s + len(sm)]
+                xsm = gc[off_s : off_s + len(sm)]
                 lick_x += [*xsm, None]
                 lick_y += [*_to_curve(sm[: len(xsm)]), None]
 
         # Reward-probability schedule (pL red, pR blue), drawn as 0..1 lines in the bottom
         # panel just like the trial-based figure; broken at session boundaries.
-        if (tr_s is not None and n_tr and len(tr_s) == n_tr
-                and {"reward_probabilityL", "reward_probabilityR"} <= set(tr_s.columns)):
+        if (
+            tr_s is not None
+            and n_tr
+            and len(tr_s) == n_tr
+            and {"reward_probabilityL", "reward_probabilityR"} <= set(tr_s.columns)
+        ):
             has_prob = True
             probL_x += [*gc, None]
             probL_y += [*tr_s["reward_probabilityL"].to_numpy(), None]
@@ -659,38 +772,91 @@ def _trial_of(times, _gc=gc, _n=n_tr):
     ht = "%%{x:.2f}s<br>trial %%{customdata[0]}<br>session %%{customdata[1]}<extra>%s</extra>"
     for key, (name, color, width, label) in ev_meta.items():
         a = ev_acc[key]
-        fig.add_trace(go.Scattergl(
-            x=a["x"], y=a["y"], customdata=a["cd"], mode="lines",
-            line=dict(color=color, width=width), name=label, hovertemplate=ht % label),
-            row=1, col=1)
+        fig.add_trace(
+            go.Scattergl(
+                x=a["x"],
+                y=a["y"],
+                customdata=a["cd"],
+                mode="lines",
+                line=dict(color=color, width=width),
+                name=label,
+                hovertemplate=ht % label,
+            ),
+            row=1,
+            col=1,
+        )
 
     for gname, gcolor in [("go cue", "green"), ("go cue (ignored)", "red")]:
         a = gocue_acc[gname]
         if a["x"]:
-            fig.add_trace(go.Scattergl(
-                x=a["x"], y=a["y"], customdata=a["cd"], mode="lines",
-                line=dict(color=gcolor, width=0.75), opacity=0.75, name=gname,
-                hovertemplate=ht % gname), row=1, col=1)
+            fig.add_trace(
+                go.Scattergl(
+                    x=a["x"],
+                    y=a["y"],
+                    customdata=a["cd"],
+                    mode="lines",
+                    line=dict(color=gcolor, width=0.75),
+                    opacity=0.75,
+                    name=gname,
+                    hovertemplate=ht % gname,
+                ),
+                row=1,
+                col=1,
+            )
 
     if frac_x:
-        fig.add_trace(go.Scattergl(x=frac_x, y=frac_y, mode="lines",
-                                   line=dict(color="gold", width=1.5), name="pR/(pL+pR)"),
-                      row=1, col=1)
+        fig.add_trace(
+            go.Scattergl(
+                x=frac_x,
+                y=frac_y,
+                mode="lines",
+                line=dict(color="gold", width=1.5),
+                name="pR/(pL+pR)",
+            ),
+            row=1,
+            col=1,
+        )
     if choice_x:
-        fig.add_trace(go.Scattergl(x=choice_x, y=choice_y, mode="lines",
-                                   line=dict(color="black", width=1.5),
-                                   name=f"choice (smooth = {smooth_factor})"), row=1, col=1)
+        fig.add_trace(
+            go.Scattergl(
+                x=choice_x,
+                y=choice_y,
+                mode="lines",
+                line=dict(color="black", width=1.5),
+                name=f"choice (smooth = {smooth_factor})",
+            ),
+            row=1,
+            col=1,
+        )
     if lick_x:
-        fig.add_trace(go.Scattergl(x=lick_x, y=lick_y, mode="lines",
-                                   line=dict(color="black", width=1.2, dash="dash"),
-                                   name=f"lick count (smooth = {smooth_factor})"), row=1, col=1)
+        fig.add_trace(
+            go.Scattergl(
+                x=lick_x,
+                y=lick_y,
+                mode="lines",
+                line=dict(color="black", width=1.2, dash="dash"),
+                name=f"lick count (smooth = {smooth_factor})",
+            ),
+            row=1,
+            col=1,
+        )
 
     # Row 2: reward-probability schedule (pR blue, pL red), 0..1.
     if has_prob:
-        fig.add_trace(go.Scattergl(x=probR_x, y=probR_y, mode="lines",
-                                   line=dict(color="blue", width=1), name="pR"), row=2, col=1)
-        fig.add_trace(go.Scattergl(x=probL_x, y=probL_y, mode="lines",
-                                   line=dict(color="red", width=1), name="pL"), row=2, col=1)
+        fig.add_trace(
+            go.Scattergl(
+                x=probR_x, y=probR_y, mode="lines", line=dict(color="blue", width=1), name="pR"
+            ),
+            row=2,
+            col=1,
+        )
+        fig.add_trace(
+            go.Scattergl(
+                x=probL_x, y=probL_y, mode="lines", line=dict(color="red", width=1), name="pL"
+            ),
+            row=2,
+            col=1,
+        )
 
     y_main_top = params["curve_top"]
 
@@ -708,8 +874,17 @@ def _trial_of(times, _gc=gc, _n=n_tr):
             d = C["data"].values - np.nanmin(C["data"].values)
             d = d / np.nanmax(d) + bottom
             color = fip_colors["_".join(channel.split("_")[:2])]
-            fig.add_trace(go.Scattergl(x=C.timestamps.values + last_off, y=d, mode="lines",
-                                       line=dict(color=color), name=channel), row=1, col=1)
+            fig.add_trace(
+                go.Scattergl(
+                    x=C.timestamps.values + last_off,
+                    y=d,
+                    mode="lines",
+                    line=dict(color=color),
+                    name=channel,
+                ),
+                row=1,
+                col=1,
+            )
             yticks.append(bottom + 0.5)
             ylabels.append(channel)
             band += 1
@@ -725,12 +900,23 @@ def _trial_of(times, _gc=gc, _n=n_tr):
     x_last = x_first + cum
     t0_view = first_gc if first_gc is not None else x_first
 
-    fig.update_yaxes(tickvals=yticks, ticktext=ylabels, fixedrange=True,
-                     range=[params["behavior_bottom"] - 0.05, y_main_top + 0.25], row=1, col=1)
+    fig.update_yaxes(
+        tickvals=yticks,
+        ticktext=ylabels,
+        fixedrange=True,
+        range=[params["behavior_bottom"] - 0.05, y_main_top + 0.25],
+        row=1,
+        col=1,
+    )
     fig.update_yaxes(title_text="p_reward", range=[0, 1], fixedrange=True, row=2, col=1)
     fig.update_xaxes(range=[t0_view, t0_view + 120], row=1, col=1)
-    fig.update_xaxes(title_text="Time (s)", range=[t0_view, t0_view + 120], row=2, col=1,
-                     rangeslider=dict(visible=True, thickness=0.06, range=[x_first, x_last]))
+    fig.update_xaxes(
+        title_text="Time (s)",
+        range=[t0_view, t0_view + 120],
+        row=2,
+        col=1,
+        rangeslider=dict(visible=True, thickness=0.06, range=[x_first, x_last]),
+    )
     if len(sess_spans) > 1:  # x tick labels restart at 0 each session
         tickvals, ticktext = [], []
         for start, dur in sess_spans:
@@ -744,12 +930,22 @@ def _trial_of(times, _gc=gc, _n=n_tr):
 
     fig.update_layout(
         # Title pinned to the very top-left so it clears the legend below it.
-        title=dict(text=title or "Session Scroller", x=0.0, xanchor="left",
-                   y=0.98, yanchor="top"),
-        showlegend=True, height=620, width=1000, template="simple_white",
+        title=dict(text=title or "Session Scroller", x=0.0, xanchor="left", y=0.98, yanchor="top"),
+        showlegend=True,
+        height=620,
+        width=1000,
+        template="simple_white",
         # Legend outside, top-left, horizontal, compact entries (narrow box).
-        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0,
-                    font=dict(size=9), entrywidthmode="pixels", entrywidth=125),
+        legend=dict(
+            orientation="h",
+            yanchor="bottom",
+            y=1.02,
+            xanchor="left",
+            x=0,
+            font=dict(size=9),
+            entrywidthmode="pixels",
+            entrywidth=125,
+        ),
         margin=dict(l=70, r=20, t=120, b=40),
     )
     return fig

From d1f9ba6d939a15db00a263947ad4134b68f1dd78 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 00:08:50 +0000
Subject: [PATCH 03/14] updated to nwb_list version

---
 .../plot/plot_foraging_session_plotly.py      | 57 ++++++++++++-------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index ddae2b6..8f7f826 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -13,6 +13,7 @@
 """
 
 import numpy as np
+import pandas as pd
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
@@ -543,7 +544,7 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
 
 
 def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
-    nwb, fip=[], adjust_time=True, title=None, smooth_factor=5
+    nwb_list, fip=[], adjust_time=True, title=None, smooth_factor=5
 ):
     """Plotly version of :func:`plot_session_scroller.plot_session_scroller` (time-based).
 
@@ -560,7 +561,7 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
 
     Parameters
     ----------
-    nwb, an nwb like object that contains attributes: df_events, session_id
+    nwb_list, a list of nwb like object that contains attributes: df_events, session_id
         and optionally contains attributes fip_df, df_licks
     fip (list), FIP channels to plot. Must be present in nwb.df_fip
     adjust_time : bool, optional
@@ -575,24 +576,40 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     -------
     plotly.graph_objects.Figure
     """
-    if not hasattr(nwb, "df_events"):
-        print("computing df_events first")
-        nwb.df_events = nu.create_df_events(nwb)
-        df_events = nwb.df_events
-    else:
-        df_events = nwb.df_events
-    if hasattr(nwb, "df_fip"):
-        fip_df = nwb.df_fip
-    if hasattr(nwb, "df_licks"):
-        df_licks = nwb.df_licks
-    else:
-        df_licks = None
-    if not hasattr(nwb, "df_trials"):
-        print("computing df_trials")
-        nwb.df_trials = nu.create_df_trials(nwb)
-        df_trials = nwb.df_trials
-    else:
-        df_trials = nwb.df_trials
+    # accumulate per-session dataframes, always add session_id
+    events_acc = []
+    trials_acc = []
+    fip_acc = []
+    for i, nwb in enumerate(nwb_list):
+        # ensure df_events
+        if not hasattr(nwb, "df_events") or nwb.df_events is None:
+            nwb.df_events = nu.create_df_events(nwb)
+        df_e = nwb.df_events.copy()
+        sid = getattr(nwb, "session_id", i)
+        df_e["session_id"] = sid
+        events_acc.append(df_e)
+
+        # optional fip dataframe
+        if hasattr(nwb, "df_fip") and nwb.df_fip is not None:
+            df_f = nwb.df_fip.copy()
+            df_f["session_id"] = sid
+            fip_acc.append(df_f)
+
+        # ensure df_trials if available / computable
+        if not hasattr(nwb, "df_trials") or nwb.df_trials is None:
+            try:
+                nwb.df_trials = nu.create_df_trials(nwb)
+            except Exception:
+                nwb.df_trials = None
+        if hasattr(nwb, "df_trials") and nwb.df_trials is not None:
+            df_t = nwb.df_trials.copy()
+            df_t["session_id"] = sid
+            trials_acc.append(df_t)
+
+    # concatenate (or None if nothing present)
+    df_events = pd.concat(events_acc, ignore_index=True) if events_acc else None
+    df_trials = pd.concat(trials_acc, ignore_index=True) if trials_acc else None
+    fip_df = pd.concat(fip_acc, ignore_index=True) if fip_acc else None
 
     # y-layout, bottom -> top:
     #   * event rows in [0, 1]: rewards at the outer edges, licks inside (right pair near the

From 95f4682375c4091b915bba7b576afb3c3c46ddb0 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 00:09:19 +0000
Subject: [PATCH 04/14] linting

---
 tests/test_plot_foraging_session_plotly.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 47b7db8..3749f41 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -116,8 +116,9 @@ def test_multi_session(self):
         e1 = self.df_events.assign(session_id="s1")
         e2 = self.df_events.assign(session_id="s2", timestamps=self.df_events["timestamps"] + 100)
         t1 = self.df_trials.assign(session_id="s1")
-        t2 = self.df_trials.assign(session_id="s2",
-                                   goCue_start_time=self.df_trials["goCue_start_time"] + 100)
+        t2 = self.df_trials.assign(
+            session_id="s2", goCue_start_time=self.df_trials["goCue_start_time"] + 100
+        )
         fig = plot_session_in_time_plotly(
             pd.concat([e1, e2], ignore_index=True),
             df_trials=pd.concat([t1, t2], ignore_index=True),

From 547ef1c3359a66c8e5d5f7614b13c746b40241c2 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 00:18:37 +0000
Subject: [PATCH 05/14] fixed tests

---
 tests/test_plot_foraging_session_plotly.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 3749f41..05cc71e 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -5,6 +5,8 @@
 
 import os
 import unittest
+from types import SimpleNamespace
+
 
 import numpy as np
 import pandas as pd
@@ -98,15 +100,15 @@ def setUp(self):
 
     def test_events_only(self):
         """Works with just an events frame (no probability band)."""
-        fig = plot_session_in_time_plotly(self.df_events)
+        nwb = SimpleNamespace(df_events=self.df_events)
+        fig = plot_session_in_time_plotly([nwb])
         self.assertIsInstance(fig, go.Figure)
         self.assertGreater(len(fig.data), 0)
 
     def test_with_trials(self):
         """Supplying df_trials adds the reward-probability band traces."""
-        fig = plot_session_in_time_plotly(
-            self.df_events, df_trials=self.df_trials, title="unit_test"
-        )
+        nwb = SimpleNamespace(df_events=self.df_events, df_trials=self.df_trials)
+        fig = plot_session_in_time_plotly([nwb], title="unit_test")
         names = [tr.name for tr in fig.data]
         self.assertIn("pR", names)
         self.assertIn("pL", names)
@@ -119,12 +121,11 @@ def test_multi_session(self):
         t2 = self.df_trials.assign(
             session_id="s2", goCue_start_time=self.df_trials["goCue_start_time"] + 100
         )
-        fig = plot_session_in_time_plotly(
-            pd.concat([e1, e2], ignore_index=True),
-            df_trials=pd.concat([t1, t2], ignore_index=True),
-        )
-        self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels
+        nwb1 = SimpleNamespace(df_events=e1, df_trials=t1)
+        nwb2 = SimpleNamespace(df_events=e2, df_trials=t2)
 
+        fig = plot_session_in_time_plotly([nwb1, nwb2])
+        self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels
 
 if __name__ == "__main__":
     unittest.main()

From b89cd2a88cd8c160c68b0cb74020ab60b5337d61 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 00:27:39 +0000
Subject: [PATCH 06/14] linting

---
 .../metrics/trial_metrics.py                              | 2 +-
 .../plot/plot_foraging_session.py                         | 3 ++-
 .../plot/plot_foraging_session_plotly.py                  | 8 +++++---
 tests/test_plot_foraging_session_plotly.py                | 1 +
 4 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py b/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py
index e4cd53b..b5ae322 100644
--- a/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py
+++ b/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py
@@ -349,7 +349,7 @@ def get_average_signal_window(
     if output_col is None:
         output_col = (
             f"{data_column}_{channel}_{offsets[0]}_"
-            f"{offsets[1]}_{alignment_event.replace('_in_session','')}"
+            f"{offsets[1]}_{alignment_event.replace('_in_session', '')}"
         )
 
     # copy df_trials, drops na values, sort trial by alignment event
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session.py
index dc2fca9..2deaac8 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session.py
@@ -66,7 +66,8 @@ def plot_foraging_session_nwb(nwb, **kwargs):
         0,
         1.05,
         f"{nwb.session_id}\n"
-        f'Total trials {len(nwb.df_trials)}, ignored {np.sum(nwb.df_trials["animal_response"]==2)},'
+        f'Total trials {len(nwb.df_trials)},' +
+        'ignored {np.sum(nwb.df_trials["animal_response"] == 2)},'
         f' left {np.sum(nwb.df_trials["animal_response"] == 0)},'
         f' right {np.sum(nwb.df_trials["animal_response"] == 1)}',
         fontsize=8,
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index 8f7f826..420b217 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -23,6 +23,8 @@
 )
 from aind_dynamic_foraging_basic_analysis.plot.plot_foraging_session import moving_average
 from aind_dynamic_foraging_basic_analysis.plot.style import PHOTOSTIM_EPOCH_MAPPING
+from aind_dynamic_foraging_data_utils import nwb_utils as nu
+
 
 # Map the matplotlib single-letter colors used by the matplotlib versions to plotly names,
 # so the two renderings line up. Anything not listed is passed through unchanged.
@@ -752,17 +754,17 @@ def _trial_of(times, _gc=gc, _n=n_tr):
                     moving_average(~np.isnan(choice), smooth_factor) + 1e-6
                 )
                 sm[sm > 100] = np.nan
-                xsm = gc[off_s : off_s + len(sm)]
+                xsm = gc[off_s: off_s + len(sm)]
                 choice_x += [*xsm, None]
                 choice_y += [*_to_curve(sm[: len(xsm)]), None]
             lt = np.concatenate([_ev("left_lick_time"), _ev("right_lick_time")])
             if len(lt):
-                counts = np.bincount(_trial_of(lt), minlength=n_tr + 1)[1 : n_tr + 1]
+                counts = np.bincount(_trial_of(lt), minlength=n_tr + 1)[1: n_tr + 1]
                 sm = moving_average(counts.astype(float), smooth_factor)
                 top = np.nanmax(sm) if len(sm) else 0
                 if top > 0:
                     sm = sm / top
-                xsm = gc[off_s : off_s + len(sm)]
+                xsm = gc[off_s: off_s + len(sm)]
                 lick_x += [*xsm, None]
                 lick_y += [*_to_curve(sm[: len(xsm)]), None]
 
diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 05cc71e..5d7cd61 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -127,5 +127,6 @@ def test_multi_session(self):
         fig = plot_session_in_time_plotly([nwb1, nwb2])
         self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels
 
+
 if __name__ == "__main__":
     unittest.main()

From 108bcad80f6ac8b211787d4d19cb03cf76104202 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 01:29:40 +0000
Subject: [PATCH 07/14] missed coverage... !

---
 .../__init__.py                               |  1 +
 .../plot/plot_foraging_session_plotly.py      |  4 +-
 tests/test_plot_foraging_session_plotly.py    | 51 +++++++++++++++++--
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/__init__.py b/src/aind_dynamic_foraging_basic_analysis/__init__.py
index c1d3ba2..9e584db 100644
--- a/src/aind_dynamic_foraging_basic_analysis/__init__.py
+++ b/src/aind_dynamic_foraging_basic_analysis/__init__.py
@@ -6,5 +6,6 @@
 from .plot.plot_foraging_session import plot_foraging_session  # noqa: F401
 from .plot.plot_foraging_session_plotly import (  # noqa: F401
     plot_foraging_session_plotly,
+    plot_foraging_session_nwb_plotly,
     plot_session_in_time_plotly,
 )
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index 420b217..9ccfd2d 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -524,7 +524,7 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
         return
 
     if "side_bias" not in nwb.df_trials:
-        fig, axes = plot_foraging_session_plotly(
+        fig = plot_foraging_session_plotly(
             [np.nan if x == 2 else x for x in nwb.df_trials["animal_response"].values],
             nwb.df_trials["earned_reward"].values,
             [nwb.df_trials["reward_probabilityL"], nwb.df_trials["reward_probabilityR"]],
@@ -533,7 +533,7 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
     else:
         if "plot_list" not in kwargs:
             kwargs["plot_list"] = ["choice", "reward_prob", "bias"]
-        fig, axes = plot_foraging_session_plotly(
+        fig = plot_foraging_session_plotly(
             [np.nan if x == 2 else x for x in nwb.df_trials["animal_response"].values],
             nwb.df_trials["earned_reward"].values,
             [nwb.df_trials["reward_probabilityL"], nwb.df_trials["reward_probabilityR"]],
diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 5d7cd61..4b64487 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -5,7 +5,6 @@
 
 import os
 import unittest
-from types import SimpleNamespace
 
 
 import numpy as np
@@ -14,10 +13,20 @@
 
 from aind_dynamic_foraging_basic_analysis import (
     plot_foraging_session_plotly,
+    plot_foraging_session_nwb_plotly,
     plot_session_in_time_plotly,
 )
 from tests.nwb_io import get_history_from_nwb
 
+class EmptyNWB:
+    """
+    Just an empty class for saving attributes to
+    """
+
+    def __init__(self, df_trials=None, df_events=None):
+        self.df_trials = df_trials
+        self.df_events = df_events
+
 
 class TestPlotForagingSessionPlotly(unittest.TestCase):
     """Test the trial-based plotly plot against a real session."""
@@ -35,6 +44,38 @@ def setUpClass(cls):
             _,
         ) = get_history_from_nwb(nwb_file)
 
+
+    def test_nwb_plot(self):
+        """ Tests plotting form nwb works"""
+
+        # Test without bias column
+        choices = np.array([0, 0, 1, 1, 2, 2])
+        rewards = np.array([True, False, True, False, False, False])
+        pL = [0.1] * 6
+        pR = [0.8] * 6
+        df = pd.DataFrame()
+        df["animal_response"] = choices
+        df["earned_reward"] = rewards
+        df["reward_probabilityL"] = pL
+        df["reward_probabilityR"] = pR
+        df["auto_waterL"] = [0] * 6
+        df["auto_waterR"] = [0] * 6
+        nwb = EmptyNWB(df_trials = df)
+        nwb.session_id = "test"
+        plot_foraging_session_nwb_plotly(nwb)
+
+        # Test with bias column
+        nwb.df_trials["side_bias"] = np.array([0, 0, 0.1, 0.1, 0.05, 0.05])
+        nwb.df_trials["side_bias_confidence_interval"] = [
+            [-1, 1],
+            [-1, 1],
+            [-1, 1],
+            [-1, 1],
+            [-1, 1],
+            [-1, 1],
+        ]
+        plot_foraging_session_nwb_plotly(nwb)
+
     def test_returns_figure(self):
         """A plotly Figure is returned with both panels populated."""
         fig = plot_foraging_session_plotly(
@@ -100,14 +141,14 @@ def setUp(self):
 
     def test_events_only(self):
         """Works with just an events frame (no probability band)."""
-        nwb = SimpleNamespace(df_events=self.df_events)
+        nwb = EmptyNWB(df_events = self.df_events)
         fig = plot_session_in_time_plotly([nwb])
         self.assertIsInstance(fig, go.Figure)
         self.assertGreater(len(fig.data), 0)
 
     def test_with_trials(self):
         """Supplying df_trials adds the reward-probability band traces."""
-        nwb = SimpleNamespace(df_events=self.df_events, df_trials=self.df_trials)
+        nwb = EmptyNWB(df_events=self.df_events, df_trials=self.df_trials)
         fig = plot_session_in_time_plotly([nwb], title="unit_test")
         names = [tr.name for tr in fig.data]
         self.assertIn("pR", names)
@@ -121,8 +162,8 @@ def test_multi_session(self):
         t2 = self.df_trials.assign(
             session_id="s2", goCue_start_time=self.df_trials["goCue_start_time"] + 100
         )
-        nwb1 = SimpleNamespace(df_events=e1, df_trials=t1)
-        nwb2 = SimpleNamespace(df_events=e2, df_trials=t2)
+        nwb1 = EmptyNWB(df_events=e1, df_trials=t1)
+        nwb2 = EmptyNWB(df_events=e2, df_trials=t2)
 
         fig = plot_session_in_time_plotly([nwb1, nwb2])
         self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels

From 1308bebea4c378b7842ec6f6775039a4fab84fd1 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 01:31:34 +0000
Subject: [PATCH 08/14] linting

---
 .../plot/plot_foraging_session_plotly.py                    | 2 ++
 tests/test_plot_foraging_session_plotly.py                  | 6 +++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index 9ccfd2d..ed22b02 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -544,6 +544,8 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
             **kwargs,
         )
 
+    return fig
+
 
 def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     nwb_list, fip=[], adjust_time=True, title=None, smooth_factor=5
diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 4b64487..84c8b97 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -18,6 +18,7 @@
 )
 from tests.nwb_io import get_history_from_nwb
 
+
 class EmptyNWB:
     """
     Just an empty class for saving attributes to
@@ -44,7 +45,6 @@ def setUpClass(cls):
             _,
         ) = get_history_from_nwb(nwb_file)
 
-
     def test_nwb_plot(self):
         """ Tests plotting form nwb works"""
 
@@ -60,7 +60,7 @@ def test_nwb_plot(self):
         df["reward_probabilityR"] = pR
         df["auto_waterL"] = [0] * 6
         df["auto_waterR"] = [0] * 6
-        nwb = EmptyNWB(df_trials = df)
+        nwb = EmptyNWB(df_trials=df)
         nwb.session_id = "test"
         plot_foraging_session_nwb_plotly(nwb)
 
@@ -141,7 +141,7 @@ def setUp(self):
 
     def test_events_only(self):
         """Works with just an events frame (no probability band)."""
-        nwb = EmptyNWB(df_events = self.df_events)
+        nwb = EmptyNWB(df_events=self.df_events)
         fig = plot_session_in_time_plotly([nwb])
         self.assertIsInstance(fig, go.Figure)
         self.assertGreater(len(fig.data), 0)

From 7c15bf5a76d2f7dac9a754995b565166d90b8dd1 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 02:26:19 +0000
Subject: [PATCH 09/14] missed doc string

---
 tests/test_plot_foraging_session_plotly.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index 84c8b97..cd5e266 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -25,6 +25,7 @@ class EmptyNWB:
     """
 
     def __init__(self, df_trials=None, df_events=None):
+        """ adds df_trials and df_events"""
         self.df_trials = df_trials
         self.df_events = df_events
 

From f700c16f0579ebee4592dfbcec9a5b9ebd8987e8 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 02:41:13 +0000
Subject: [PATCH 10/14] missed another coverage

---
 tests/test_plot_foraging_session_plotly.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index cd5e266..adfc535 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -48,6 +48,10 @@ def setUpClass(cls):
 
     def test_nwb_plot(self):
         """ Tests plotting form nwb works"""
+        # Test we have df_trials
+        nwb = EmptyNWB()
+        del nwb.df_trials
+        plot_foraging_session_nwb_plotly(nwb)
 
         # Test without bias column
         choices = np.array([0, 0, 1, 1, 2, 2])

From f30ff8dc02ee5f661ef71ae0723d11bb06d7ca20 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Tue, 30 Jun 2026 03:01:06 +0000
Subject: [PATCH 11/14] added wrapper to keep han's behavior functionality

---
 .../__init__.py                               |  2 +
 .../plot/plot_foraging_session_plotly.py      | 66 +++++++++++++++++--
 tests/test_plot_foraging_session_plotly.py    | 28 ++++++--
 3 files changed, 86 insertions(+), 10 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/__init__.py b/src/aind_dynamic_foraging_basic_analysis/__init__.py
index 9e584db..370b992 100644
--- a/src/aind_dynamic_foraging_basic_analysis/__init__.py
+++ b/src/aind_dynamic_foraging_basic_analysis/__init__.py
@@ -8,4 +8,6 @@
     plot_foraging_session_plotly,
     plot_foraging_session_nwb_plotly,
     plot_session_in_time_plotly,
+    plot_session_in_time_nwb_plotly,
+
 )
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index ed22b02..111b1c0 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -547,7 +547,7 @@ def plot_foraging_session_nwb_plotly(nwb, **kwargs):
     return fig
 
 
-def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
+def plot_session_in_time_nwb_plotly(  # noqa: C901 pragma: no cover
     nwb_list, fip=[], adjust_time=True, title=None, smooth_factor=5
 ):
     """Plotly version of :func:`plot_session_scroller.plot_session_scroller` (time-based).
@@ -566,7 +566,7 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     Parameters
     ----------
     nwb_list, a list of nwb like object that contains attributes: df_events, session_id
-        and optionally contains attributes fip_df, df_licks
+        and optionally contains attributes df_fip, df_licks
     fip (list), FIP channels to plot. Must be present in nwb.df_fip
     adjust_time : bool, optional
         If True (default), shift time so the first event is at t = 0 (always shifted when
@@ -613,7 +613,61 @@ def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
     # concatenate (or None if nothing present)
     df_events = pd.concat(events_acc, ignore_index=True) if events_acc else None
     df_trials = pd.concat(trials_acc, ignore_index=True) if trials_acc else None
-    fip_df = pd.concat(fip_acc, ignore_index=True) if fip_acc else None
+    df_fip = pd.concat(fip_acc, ignore_index=True) if fip_acc else None
+
+    return plot_session_in_time_plotly(
+        df_events, df_trials, df_fip, fip=fip,
+        adjust_time=adjust_time, title=title, smooth_factor=smooth_factor
+    )
+
+
+def plot_session_in_time_plotly(  # noqa: C901 pragma: no cover
+    df_events, df_trials=None, df_fip=None, fip=[], adjust_time=True, title=None, smooth_factor=5
+):
+    """Plotly version of :func:`plot_session_scroller.plot_session_scroller` (time-based).
+
+    Plots the session in real time (not in trial): left / right licks and rewards as ticks,
+    go cues as vertical lines (red for ignored trials), smoothed overlays above the events,
+    and -- when ``df_trials`` is supplied -- the left / right reward-probability band in the
+    rangeslider "scroller" below.
+
+    Multiple sessions: if ``df_events`` has a ``session_id`` column with more than one
+    session, the sessions are concatenated end-to-end along time in order of appearance
+    (each restarted at the running offset); a thick vertical line marks each boundary and
+    per-trial / smoothed quantities reset per session. ``df_trials`` is matched per session
+    by its own ``session_id`` column when present.
+
+    Parameters
+    ----------
+    df_events : pandas.DataFrame
+        Tidy dataframe of session events (``event`` + ``timestamps``; optional ``session_id``).
+        Recognised events: ``left_lick_time``, ``right_lick_time``, ``left_reward_delivery_time``,
+        ``right_reward_delivery_time`` and ``goCue_start_time``.
+    df_trials : pandas.DataFrame, optional
+        Per-trial dataframe for the reward-probability band / overlays / red ignored go cues
+        (and a fallback source of go-cue times). Uses ``goCue_start_time``,
+        ``reward_probabilityL/R`` and ``animal_response``; go-cue times must share the
+        ``df_events`` time base. Matched per session via ``session_id`` when present.
+    df_fip : pandas.DataFrame, optional
+        Tidy FIP measurements (single-session only); each present channel is normalised and
+        stacked above the behavior panel.
+    adjust_time : bool, optional
+        If True (default), shift time so the first event is at t = 0 (always shifted when
+        concatenating multiple sessions).
+    title : str, optional
+        Figure title.
+    smooth_factor : int, optional
+        Smoothing window for the choice / lick-count overlays, by default 5.
+
+    Returns
+    -------
+    plotly.graph_objects.Figure
+    """
+    df_events = df_events.copy()
+    if df_trials is not None:
+        df_trials = df_trials.copy()
+    if df_fip is not None:
+        df_fip = df_fip.copy()
 
     # y-layout, bottom -> top:
     #   * event rows in [0, 1]: rewards at the outer edges, licks inside (right pair near the
@@ -882,16 +936,16 @@ def _trial_of(times, _gc=gc, _n=n_tr):
     y_main_top = params["curve_top"]
 
     # FIP channels (single-session only), normalised and stacked above the behavior panel
-    if fip_df is not None and len(sessions) == 1 and len(fip) > 0:
+    if df_fip is not None and len(sessions) == 1 and len(fip) > 0:
         fip_channels = fip
         fip_colors = {"G_1": "green", "G_2": "darkgreen", "R_1": "red", "R_2": "darkred"}
-        present = set(fip_df["event"].unique())
+        present = set(df_fip["event"].unique())
         band = 0
         for channel in fip_channels:
             if channel not in present:
                 continue
             bottom = params["curve_top"] + 0.1 + band
-            C = fip_df.query("event == @channel").copy()
+            C = df_fip.query("event == @channel").copy()
             d = C["data"].values - np.nanmin(C["data"].values)
             d = d / np.nanmax(d) + bottom
             color = fip_colors["_".join(channel.split("_")[:2])]
diff --git a/tests/test_plot_foraging_session_plotly.py b/tests/test_plot_foraging_session_plotly.py
index adfc535..3918b2e 100644
--- a/tests/test_plot_foraging_session_plotly.py
+++ b/tests/test_plot_foraging_session_plotly.py
@@ -15,6 +15,8 @@
     plot_foraging_session_plotly,
     plot_foraging_session_nwb_plotly,
     plot_session_in_time_plotly,
+    plot_session_in_time_nwb_plotly,
+
 )
 from tests.nwb_io import get_history_from_nwb
 
@@ -146,15 +148,27 @@ def setUp(self):
 
     def test_events_only(self):
         """Works with just an events frame (no probability band)."""
+        fig = plot_session_in_time_plotly(self.df_events)
+        self.assertIsInstance(fig, go.Figure)
+        self.assertGreater(len(fig.data), 0)
+
+        # nwb version
         nwb = EmptyNWB(df_events=self.df_events)
-        fig = plot_session_in_time_plotly([nwb])
+        fig = plot_session_in_time_nwb_plotly([nwb])
         self.assertIsInstance(fig, go.Figure)
         self.assertGreater(len(fig.data), 0)
 
     def test_with_trials(self):
         """Supplying df_trials adds the reward-probability band traces."""
+        fig = plot_session_in_time_plotly(
+            self.df_events, df_trials=self.df_trials, title="unit_test"
+        )
+        names = [tr.name for tr in fig.data]
+        self.assertIn("pR", names)
+        self.assertIn("pL", names)
+        # nwb version
         nwb = EmptyNWB(df_events=self.df_events, df_trials=self.df_trials)
-        fig = plot_session_in_time_plotly([nwb], title="unit_test")
+        fig = plot_session_in_time_nwb_plotly([nwb], title="unit_test")
         names = [tr.name for tr in fig.data]
         self.assertIn("pR", names)
         self.assertIn("pL", names)
@@ -167,10 +181,16 @@ def test_multi_session(self):
         t2 = self.df_trials.assign(
             session_id="s2", goCue_start_time=self.df_trials["goCue_start_time"] + 100
         )
+        fig = plot_session_in_time_plotly(
+            pd.concat([e1, e2], ignore_index=True),
+            df_trials=pd.concat([t1, t2], ignore_index=True),
+        )
+        self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels
+
+        # nwb version
         nwb1 = EmptyNWB(df_events=e1, df_trials=t1)
         nwb2 = EmptyNWB(df_events=e2, df_trials=t2)
-
-        fig = plot_session_in_time_plotly([nwb1, nwb2])
+        fig = plot_session_in_time_nwb_plotly([nwb1, nwb2])
         self.assertEqual(len(fig.layout.shapes), 2)  # one boundary, drawn in both panels
 
 

From f865d460c9934f14c3ea9dd9639c68ba2d093f24 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Wed, 1 Jul 2026 00:28:30 +0000
Subject: [PATCH 12/14] fixed FIP colors

---
 .../plot/plot_foraging_session_plotly.py           |  4 ++--
 .../plot/plot_session_scroller.py                  |  2 +-
 .../plot/style.py                                  | 14 +++++++-------
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
index 111b1c0..047834c 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_foraging_session_plotly.py
@@ -22,6 +22,7 @@
     PhotostimData,
 )
 from aind_dynamic_foraging_basic_analysis.plot.plot_foraging_session import moving_average
+from aind_dynamic_foraging_basic_analysis.plot.plot_session_scroller import get_fip_color
 from aind_dynamic_foraging_basic_analysis.plot.style import PHOTOSTIM_EPOCH_MAPPING
 from aind_dynamic_foraging_data_utils import nwb_utils as nu
 
@@ -938,7 +939,6 @@ def _trial_of(times, _gc=gc, _n=n_tr):
     # FIP channels (single-session only), normalised and stacked above the behavior panel
     if df_fip is not None and len(sessions) == 1 and len(fip) > 0:
         fip_channels = fip
-        fip_colors = {"G_1": "green", "G_2": "darkgreen", "R_1": "red", "R_2": "darkred"}
         present = set(df_fip["event"].unique())
         band = 0
         for channel in fip_channels:
@@ -948,7 +948,7 @@ def _trial_of(times, _gc=gc, _n=n_tr):
             C = df_fip.query("event == @channel").copy()
             d = C["data"].values - np.nanmin(C["data"].values)
             d = d / np.nanmax(d) + bottom
-            color = fip_colors["_".join(channel.split("_")[:2])]
+            color = get_fip_color(channel)
             fig.add_trace(
                 go.Scattergl(
                     x=C.timestamps.values + last_off,
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
index 1bde1db..fbb68e9 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
@@ -548,4 +548,4 @@ def get_fip_color(channel):
     if root in FIP_COLORS:
         return FIP_COLORS.get(root)
 
-    return "k"
+    return "black"
diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/style.py b/src/aind_dynamic_foraging_basic_analysis/plot/style.py
index 712d7dd..c5a3696 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/style.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/style.py
@@ -26,14 +26,14 @@
 
 # Colorscheme for FIP channels
 FIP_COLORS = {
-    "G": "g",
-    "R": "r",
+    "G": "green",
+    "R": "red",
     "Iso": "gray",
-    "goCue_start_time": "b",
-    "left_lick_time": "m",
-    "right_lick_time": "r",
-    "left_reward_delivery_time": "b",
-    "right_reward_delivery_time": "r",
+    "goCue_start_time": "blue",
+    "left_lick_time": "magenta",
+    "right_lick_time": "red",
+    "left_reward_delivery_time": "blue",
+    "right_reward_delivery_time": "red",
 }
 
 

From 54ea13532b3697df5bc21a88aac5c05ca052d3a9 Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Wed, 1 Jul 2026 00:49:34 +0000
Subject: [PATCH 13/14] skipping coverage for plot_session_scroller.

---
 .../plot/plot_session_scroller.py                           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
index fbb68e9..30804dd 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
@@ -445,7 +445,7 @@ def on_key_press(event):
     return fig, ax
 
 
-def plot_metric(df_trials, go_cue_times, metric, ax):
+def plot_metric(df_trials, go_cue_times, metric, ax): # pragma: no cover
     """
     Plots a metric from df_trials
 
@@ -505,7 +505,7 @@ def plot_metric(df_trials, go_cue_times, metric, ax):
     ax.set_ylabel(ylabel, fontsize=12)
 
 
-def plot_fip(fip_df, channel, ax):
+def plot_fip(fip_df, channel, ax): # pragma: no cover
     """
     Plot an FIP channel
     """
@@ -534,7 +534,7 @@ def plot_fip(fip_df, channel, ax):
     ax.axhline(0, color="k", linewidth=0.5, alpha=0.25)
 
 
-def get_fip_color(channel):
+def get_fip_color(channel): # pragma: no cover
     """
     Gets the color for FIP
     if the channel is defined in style.FIP_COLORS, use that

From be61e29af6b0d0e0cbc2b4a016221d792e18e09b Mon Sep 17 00:00:00 2001
From: rachelstephlee <rachel.lee@alleninstitute.org>
Date: Wed, 1 Jul 2026 18:14:40 +0000
Subject: [PATCH 14/14] minor linting

---
 .../plot/plot_session_scroller.py                           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
index 30804dd..626ce0e 100644
--- a/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
+++ b/src/aind_dynamic_foraging_basic_analysis/plot/plot_session_scroller.py
@@ -445,7 +445,7 @@ def on_key_press(event):
     return fig, ax
 
 
-def plot_metric(df_trials, go_cue_times, metric, ax): # pragma: no cover
+def plot_metric(df_trials, go_cue_times, metric, ax):  # pragma: no cover
     """
     Plots a metric from df_trials
 
@@ -505,7 +505,7 @@ def plot_metric(df_trials, go_cue_times, metric, ax): # pragma: no cover
     ax.set_ylabel(ylabel, fontsize=12)
 
 
-def plot_fip(fip_df, channel, ax): # pragma: no cover
+def plot_fip(fip_df, channel, ax):  # pragma: no cover
     """
     Plot an FIP channel
     """
@@ -534,7 +534,7 @@ def plot_fip(fip_df, channel, ax): # pragma: no cover
     ax.axhline(0, color="k", linewidth=0.5, alpha=0.25)
 
 
-def get_fip_color(channel): # pragma: no cover
+def get_fip_color(channel):  # pragma: no cover
     """
     Gets the color for FIP
     if the channel is defined in style.FIP_COLORS, use that