added analysis_util enriching functions.

Author: rachelstephlee

src/rachel_analysis_utils/analysis_utils.py

@@ -1,156 +1,111 @@
-import warnings
-import glob
-import pandas as pd
 import numpy as np
+import pandas as pd 
+from analysis_wrapper.plots import summary_plots
+from aind_dynamic_foraging_basic_analysis.metrics import trial_metrics
 
 
-class dummy_nwb:
-    def __init__(self, df_trials, df_events, df_fip, ses_idx = None, df_licks = None, grouped = False) -> None:
-        if grouped is True:
-            self.df_events = df_events
-            self.df_fip = df_fip
-            self.df_trials = df_trials
-            self.session_id = ', '.join(df_trials.ses_idx.unique())
-            return
-        if ses_idx is None and grouped is False:
-
-            if len(df_trials.ses_idx.unique()) > 1 or \
-                len(df_events.ses_idx.unique()) > 1 or \
-                len(df_fip.ses_idx.unique()) > 1:
-
-                warnings.warn('multiple sessions found, only one will be attached to this nwb')
-            ses_idx = df_trials.ses_idx.unique()[0]
-             
-                
-        assert df_fip[df_fip['ses_idx'] == ses_idx].shape[0] != 0 ,(
-            "No session exists in the df_fip"
-        )
-        self.session_id = ses_idx
-        self.df_events = df_events[df_events['ses_idx'] == ses_idx]
-        self.df_fip = df_fip[df_fip['ses_idx'] == ses_idx].copy().reset_index(drop=True)
-        self.df_trials = df_trials[df_trials['ses_idx'] == ses_idx]
-        if df_licks:
-            self.df_licks = df_licks[df_licks['ses_idx'] == ses_idx]
-
-        nwb_file_name = glob.glob(f"/root/capsule/data/**{ses_idx}**/nwb/**.nwb")
-        if len(nwb_file_name):
-            self.nwb_file_loc = nwb_file_name[0]
-        else:
-            self.nwb_file_loc = None
+def add_AUC_and_rpe_slope(nwbs_by_week, parameters, data_column = 'data_z_norm', 
+                            alignment_event = 'choice_time_in_session',offsets = [0.33,1]):
+    rpe_slope_dict = {}
+    nwbs_by_week_enriched = []
+    for channel in list(parameters["channels"].keys()):
+        if parameters['preprocessing'] is not 'raw':
+            channel = channel +  '_' + parameters['preprocessing'] 
+
+        avg_signal_col = summary_plots.output_col_name(channel, data_column, alignment_event)
+        for nwb_week in nwbs_by_week:
+        
+            nwb_week_enriched = trial_metrics.get_average_signal_window_multi(
+                            nwb_week,
+                            alignment_event=alignment_event,
+                            offsets=offsets,
+                            channel=channel,
+                            data_column=data_column,
+                            output_col = avg_signal_col
+                        )
+            nwbs_by_week_enriched.append(nwb_week_enriched)
 
+        # get rpe slope per session 
+
+        df_trials_all = pd.concat([nwb.df_trials for nwb_week in nwbs_by_week_enriched for nwb in nwb_week])
+        rpe_slope = []
+        for ses_idx in sorted(df_trials_all['ses_idx'].unique()):
+            
+            data = df_trials_all[df_trials_all['ses_idx'] == ses_idx]
+            data = data.dropna(subset = [avg_signal_col, 'RPE_earned'])
+            if len(data) == 0:
+                continue
+            data_neg = data[data['RPE_earned'] < 0]
+            data_pos = data[data['RPE_earned'] >= 0]
+
+            ses_date = pd.to_datetime(ses_idx.split('_')[1])
+            (_,_, slope_pos) = summary_plots.get_RPE_by_avg_signal_fit(data_pos, avg_signal_col)
+            (_,_, slope_neg) = summary_plots.get_RPE_by_avg_signal_fit(data_neg, avg_signal_col)
+            rpe_slope.append([ses_date, slope_pos, slope_neg])
+        rpe_slope = pd.DataFrame(rpe_slope, columns=['date', 'slope (RPE >= 0)', 'slope (RPE < 0)'])
+        rpe_slope_dict[channel] = rpe_slope
+
+    subject_id = str(nwbs_by_week_enriched[0][0]).split(' ')[1].split('_')[0]
+    # Concatenate with keys, turning dict keys into an index
+    combined_rpe_slope = pd.concat(rpe_slope_dict, names=["channel"])
+    combined_rpe_slope = combined_rpe_slope.reset_index(level="channel").reset_index(drop=True)
+
+    combined_rpe_slope.to_csv(f"/results/{subject_id}_rpe_slope.csv")
+
+    return nwbs_by_week_enriched, combined_rpe_slope
+
+
+def enrich_df_trials(df_trials):
+
+##### PART I: REWARD #######
+    df_trials['reward_all'] = df_trials['earned_reward'] + df_trials['extra_reward']
+    # Compute num_reward_past and num_no_reward_past
+    df_trials['rewarded_prev'] = df_trials.groupby('ses_idx')['reward_all'].shift(1)  # Shift to look at past values
 
-    def __str__(self):
-        return f"session {self.session_id}"
+    df_trials['num_reward_past'] = df_trials.groupby(
+                            (df_trials['rewarded_prev'] != df_trials['reward_all']).cumsum()).cumcount() + 1
 
-    def __repr__(self):
-        return f"{self.session_id}"
+    # Set 'NA' for mismatched reward types
+    df_trials.loc[df_trials['reward_all'] == 0, 'num_reward_past'] = df_trials.loc[df_trials['reward_all'] == 0, 'num_reward_past']* -1 
+
+    ##### PART II: BINNING RPE #######
+    # get RPE binned columns. 
+    RPE_binned3_label_names = [str(np.round(i,2)) for i in np.arange(-1,0.99,1/3)]
+
+    bins = np.arange(-1,1.01,1/3)
+    bins[-1] = 1.001
+
+    df_trials['RPE-binned3'] = pd.cut(df_trials['RPE_earned'],# all versus earned not a huge difference
+                        bins = bins, right = True, labels=RPE_binned3_label_names)
+
+    ##### PART III: BINNING QCHOSEN #######
+    bins = [0.0, 1/3, 2/3, 1.01]
+    q_labels = ["Qch 0", "Qch 0.33", "Qch 0.66"]
+
+    q_bin = pd.cut(df_trials['Q_chosen'], bins=bins, labels=q_labels, include_lowest=True, right=True)
+    reward_label = df_trials['earned_reward'].map({True: "R+", False: "R-"})
+
+    # build combined label series (None where q_bin is NA)
+    reward_Qcat_series = pd.Series(
+        np.where(q_bin.isna(), None, reward_label.astype(str) + " (" + q_bin.astype(str) + ")"),
+        index=df_trials.index
+    )
+
+    # ordered categories you requested
+    Qch_binned3_label_names = [
+        "R- (Qch 0)", "R- (Qch 0.33)", "R- (Qch 0.66)",
+        "R+ (Qch 0)", "R+ (Qch 0.33)", "R+ (Qch 0.66)"
+    ]
+
+    # assign final ordered categorical to dataframe (no intermediate column left behind)
+    df_trials['Qch-binned3'] = pd.Categorical(reward_Qcat_series, categories=Qch_binned3_label_names, ordered=True)
 
 
-def get_dummy_nwbs(df_trials, df_events, df_fip):
-    ses_idx_list = df_trials.ses_idx.unique()
-    dummy_nwbs_list = []
-    ses_dates_order = np.argsort(pd.to_datetime([ses_idx.split('_')[1] for ses_idx in ses_idx_list]))
-
-    for ses_idx in ses_idx_list[ses_dates_order]:
-        # Check if ses_idx exists in all 3 dataframes
-        if (
-            ses_idx in df_events['ses_idx'].values and
-            ses_idx in df_fip['ses_idx'].values and
-            ses_idx in df_trials['ses_idx'].values
-        ):
-            df_trials_i = df_trials[df_trials['ses_idx'] == ses_idx]
-            df_events_i = df_events[df_events['ses_idx'] == ses_idx]
-            df_fip_i = df_fip[df_fip['ses_idx'] == ses_idx]
-
-            dummy_nwbs_list.append(dummy_nwb(df_trials_i, df_events_i, df_fip_i))
-        else:
-            warnings.warn(f"Skipping {ses_idx}: not found in all input DataFrames.", UserWarning)
-
-    return dummy_nwbs_list
-
-def get_dummy_nwbs_by_subject(df_trials, df_events, df_fip):
-    df_trials['subject_id'] =  df_trials['ses_idx'].str.split('_').str[0]
-    df_events['subject_id'] =  df_events['ses_idx'].str.split('_').str[0]
-    df_fip['subject_id'] =  df_fip['ses_idx'].str.split('_').str[0]
-    subject_id_list = df_trials.subject_id.unique()
-    dummy_nwbs_list = []
-    for subject_id in subject_id_list:
-        # Check if ses_idx exists in all 3 dataframes
-        if (
-            subject_id in df_events['subject_id'].values and
-            subject_id in df_fip['subject_id'].values and
-            subject_id in df_trials['subject_id'].values
-        ):
-            df_trials_i = df_trials[df_trials['subject_id'] == subject_id]
-            df_events_i = df_events[df_events['subject_id'] == subject_id]
-            df_fip_i = df_fip[df_fip['subject_id'] == subject_id]
-
-            dummy_nwbs_list.append(get_dummy_nwbs(df_trials_i, df_events_i, df_fip_i))
-        else:
-            warnings.warn(f"Skipping {subject_id}: not found in all input DataFrames.", UserWarning)
-
-    return dummy_nwbs_list
-
-def get_date_and_week_interval(df, start_date):
-    date_series = pd.to_datetime(df['ses_idx'].str.split('_').str[1], format='%Y-%m-%d')
-    week_interval_series = ((date_series - start_date).dt.days // 7) + 1
-    return week_interval_series
-
-def get_dummy_nwbs_by_week(df_sess,df_trials, df_events, df_fip):
-    start_date = pd.to_datetime(df_sess['session_date'].min())
-
-    df_sess['week_interval'] = get_date_and_week_interval(df_sess, start_date)
-    df_trials['week_interval'] = get_date_and_week_interval(df_trials, start_date)
-    df_events['week_interval'] = get_date_and_week_interval(df_events, start_date)
-    df_fip['week_interval'] = get_date_and_week_interval(df_fip, start_date)
-
-    week_interval_list = df_trials.week_interval.unique()
-    dummy_nwbs_list = []
-    for week_interval in week_interval_list:
-        # Check if ses_idx exists in all 3 dataframes
-        if (
-            week_interval in df_events['week_interval'].values and
-            week_interval in df_fip['week_interval'].values and
-            week_interval in df_trials['week_interval'].values
-        ):
-            df_trials_i = df_trials[df_trials['week_interval'] == week_interval]
-            df_events_i = df_events[df_events['week_interval'] == week_interval]
-            df_fip_i = df_fip[df_fip['week_interval'] == week_interval]
-
-            dummy_nwbs_list.append(get_dummy_nwbs(df_trials_i, df_events_i, df_fip_i))
-        else:
-            warnings.warn(f"Skipping {week_interval}: not found in all input DataFrames.", UserWarning)
-
-    return df_sess, dummy_nwbs_list
-
-
-
-def combine_dummy_nwbs_to_dfs(dummy_nwbs_list):
-    """
-    Given a list of dummy_nwb objects, concatenate their df_trials, df_events, and df_fip
-    into three large DataFrames.
-
-    Parameters
-    ----------
-    dummy_nwbs : list of dummy_nwb
-
-    Returns
-    -------
-    tuple of pd.DataFrame
-        (df_trials_all, df_events_all, df_fip_all)
-    """
-
-    df_trials_list = []
-    df_events_list = []
-    df_fip_list = []
-
-    for nwb in dummy_nwbs_list:
-        df_trials_list.append(nwb.df_trials)
-        df_events_list.append(nwb.df_events)
-        df_fip_list.append(nwb.df_fip)
-
-    df_trials_all = pd.concat(df_trials_list, ignore_index=True)
-    df_events_all = pd.concat(df_events_list, ignore_index=True)
-    df_fip_all = pd.concat(df_fip_list, ignore_index=True)
-
-    return df_trials_all, df_events_all, df_fip_all
+    ##### PART IV: GETTING STAY/LEAVE #######
+    _choice_shifted = df_trials.groupby('ses_idx')['choice'].shift(1)
+    df_trials['stay'] = df_trials['choice'] == _choice_shifted
+    df_trials['switch'] = df_trials['choice'] != _choice_shifted
+    df_trials['response_time'] = df_trials['choice_time_in_trial'] -  df_trials['goCue_start_time_in_trial']
+
+
+    return df_trials