added visualization function to be used for manual and passive center…

aopy/visualization/overnight_script_visualizations.py

@@ -0,0 +1,132 @@
+import numpy as np
+from matplotlib import pyplot as plt
+import aopy
+from aopy.data import base, db
+from aopy.data.bmi3d import get_target_locations
+
+def visualize_co_behavior_data(preproc_dir, te_ids):
+    """
+    Plot eye and cursor/hand trajectories for center-out behavioral sessions, split by trial outcome.
+
+    For each task entry (TE) ID, this loads preprocessed experiment data and tabulated
+    behavioral data, then splits trials into three mutually exclusive outcome categories:
+
+        - 'completed_trials': trials where the delay was completed, the reach was
+          completed, and the trial was rewarded (segmented from go_cue_time).
+        - 'failed_delay': trials where the delay was NOT completed
+          (segmented from center_target_on_time).
+        - 'failed_reach': trials where the delay was completed but the reach was NOT
+          completed (segmented from go_cue_time).
+
+    Trials are assigned to categories in the order listed above, and once a trial has
+    been claimed by an earlier category it is excluded from later ones (via the
+    `no_repeats_bool` mask), so each trial appears in at most one category/column.
+
+    For each category, trial segments run from the category's start event to either
+    `penalty_start_time` (if present) or `trial_end_time` (if `penalty_start_time` is
+    NaN). Eye and cursor kinematic data are tabulated over these segments and plotted
+    as trajectories overlaid on the session's unique target locations.
+
+    Produces one figure per TE ID with a 2x3 grid of subplots:
+        - Row 0: eye trajectories for each category (completed / failed_delay / failed_reach)
+        - Row 1: cursor/hand trajectories for the same categories
+
+    Parameters
+    ----------
+    preproc_dir : str
+        Path to the directory containing preprocessed experiment data.
+    te_ids : Sequence[int]
+        Iterable of task entry IDs to look up, process, and plot. One figure is
+        generated per ID.
+
+    Returns
+    -------
+    None
+        Displays/creates one matplotlib figure per TE ID; does not return a value.
+
+    Notes
+    -----
+    - Requires `db`, `aopy`, `plt`, `np`, and `get_target_locations` to be available
+      in the calling scope/module.
+    - If a category has no matching trials, a message is printed and that category
+      is skipped (its subplot axes remain unpopulated), but it still consumes a
+      column in the 2x3 grid.
+    - Hand/optitrack segment tabulation is currently disabled (`if False:`), so the
+      'user_world' datatype is not tabulated even when optitrack features are present.
+    - Assumes exactly one matching session entry is returned by `db.lookup_sessions(id=id)`.
+    """
+    bounds = [-10, 10, -10, 10]
+
+    categories = {'completed_trials': ('reward', True, 'go_cue_time', 'reward_start_time'), 
+                  'Pre-movement Failure': ('delay_completed', False, 'hold_start_time', 'penalty_start_time'),
+                  'Post-movement Failure': ('reach_completed', False, 'go_cue_time', 'penalty_start_time')}
+
+    for id in te_ids:
+        entry = db.lookup_sessions(id=id)[0]
+        df = db.summarize_entries([entry])
+
+        fig, ax = plt.subplots(2,3, figsize=(15, 10))
+        fig.suptitle(f'TE ID: {id} - {entry.task_desc}')
+
+        b = aopy.data.base.load_preproc_exp_data(preproc_dir, df['subject'][0], df['te_id'][0], df['date'][0])
+
+
+        bh_df = aopy.data.tabulate_behavior_data_center_out(preproc_dir, df['subject'], df['te_id'], df['date'])
+        bh_df = bh_df[bh_df['trial_initiated']].reset_index(drop=True)
+
+        no_repeats_bool = np.ones((bh_df.shape[0],), dtype=bool)
+
+        for i,key in enumerate(categories.keys()):
+            print(f'Generating plots for {key}...')
+            inclusion_boolean, inclusion_value, start_event, end_event = categories[key]
+
+            tmp_bool = bh_df[inclusion_boolean] == inclusion_value
+
+            filtered_df = bh_df[tmp_bool & no_repeats_bool]
+
+            start_times = filtered_df[start_event].values
+            end_times = filtered_df[end_event].values
+            #end_times[np.isnan(end_times)] = filtered_df['reward_start_time'].values[np.isnan(end_times)]
+            #end_times[np.isnan(end_times)] = filtered_df['trial_end_time'].values[np.isnan(end_times)]
+
+            if filtered_df.shape[0] == 0:
+                print(f'No trials found for {key}. Skipping...')
+                continue
+            no_repeats_bool = no_repeats_bool & ~tmp_bool
+
+            cursor_segments_reward = aopy.data.tabulate_kinematic_data(
+                preproc_dir, filtered_df['subject'], filtered_df['te_id'], filtered_df['date'],
+                start_times, end_times,
+                datatype='cursor'
+            )
+
+            eye_segments_reward = aopy.data.tabulate_kinematic_data(
+                preproc_dir, filtered_df['subject'], filtered_df['te_id'], filtered_df['date'],
+                start_times, end_times,
+                datatype='eye'
+            )
+            """
+            TODO: Add 3d hand movements 
+            if False:#'optitrack' in b[1]['features'].astype('str'):
+                hand_segments_reward = aopy.data.tabulate_kinematic_data(
+                                preproc_dir, filtered_df['subject'], filtered_df['te_id'], filtered_df['date'],
+                                start_times, end_times,
+                                datatype='user_world'
+                            )
+            """
+
+            subject, te_id, date = aopy.data.db.list_entry_details([entry])
+            targs = range(0,9)
+
+            eye_ax = ax[0, i]
+            eye_ax.set_title(f'{key} - Eye')
+            hand_ax = ax[1, i]
+            hand_ax.set_title(f'{key} - Hand')
+
+            unique_targets = get_target_locations(preproc_dir, subject[0], te_id[0], date[0], targs)
+            aopy.visualization.plot_targets(unique_targets, 2, bounds, ax=eye_ax)
+            aopy.visualization.plot_trajectories(eye_segments_reward, bounds, ax=eye_ax)
+
+
+            aopy.visualization.plot_targets(unique_targets, 2, bounds, ax=hand_ax)
+            aopy.visualization.plot_trajectories(cursor_segments_reward, bounds, ax=hand_ax)