From 78476a4db4f02c3e60f07286fec97c2820a7e78c Mon Sep 17 00:00:00 2001
From: nkeesey <140021626+nkeesey@users.noreply.github.com>
Date: Mon, 2 Dec 2024 10:07:17 -0800
Subject: [PATCH] adding lick stats to trial_metrics, adding rest of
 process_nwb stats to session_metrics

---
 .../metrics/session_metrics.py                | 354 +++++++++++-------
 .../metrics/trial_metrics.py                  |  85 +++++
 2 files changed, 300 insertions(+), 139 deletions(-)

diff --git a/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py b/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py
index 5c4b73d..fdb5ff8 100644
--- a/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py
+++ b/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py
@@ -19,104 +19,179 @@
 
 def session_metrics(nwb):
     """
-    Compute all session level metrics
-
-    Includes session level metadata as a temporary organizer
-    Includes majority of metrics from process_nwbs.py
+    Compute all session metadata and performance metrics 
+
+    -- Metadata --
+    block structure metrics,        block, contrast, and effective
+                                    probability metrics
+    duration metrics,               gocue, delay period, and iti
+                                    metrics
+    reward size metrics,            left and right reward volumes
+    lick spout movement,            movement of lickspout during session:
+                                    range, initial position, median position
+    autotrain categories,           curriculum version, name, schema,
+                                    current_stage_actual, and if overriden
+
+    -- Performance --
+    basic performance metrics,      both autowater and non-autowater specific
+                                    rates (total, finished, ignored, finished rate,
+                                    ignored rate, reward rate)
+    calculated metrics,             foraging efficiency, foraging performance
+                                    (both normal and random seed), bias naive,
+                                    chosen probability
+    lick metrics,                   reaction mean and median, early lick rate, invalid
+                                    lick ratio, double dipping finished rates (reward and total), 
+                                    lick consistency means (total, reward, and non-rewarded)
 
     New addition: chosen_probability - average difference between the chosen probability
     and non-chosen probability / the difference between the largest and smallest probability
     in the session
     """
 
-    if not hasattr(nwb, "df_trials"):
-        print("You need to compute dfs: nwb_utils.create_trials_df(nwb)")
+    if not hasattr(nwb, 'df_trials'):
+        print('You need to compute df_trials: nwb_utils.create_trials_df(nwb)')
         return
-
+    
     df = nwb.df_trials.copy()
 
-    # METADATA PLACEHOLDER
+    # Block information
+    def _get_block_starts(p_L, p_R):
+        """Find the indices of block starts"""
+        block_start_ind_left = np.where(np.hstack([True, np.diff(p_L) != 0]))[0]
+        block_start_ind_right = np.where(np.hstack([True, np.diff(p_R) != 0]))[0]
+        block_start_ind_effective = np.sort(np.unique(np.hstack([block_start_ind_left, block_start_ind_right])))
+        return block_start_ind_left, block_start_ind_right, block_start_ind_effective
+
+    # -- Key meta data --
     session_start_time = nwb.session_start_time
     session_date = session_start_time.strftime("%Y-%m-%d")
-    subject_id_from_meta = nwb.subject.subject_id
-
-    # Parse the file name for suffix
-    old_re = re.match(
-        r"(?P<subject_id>\d+)_(?P<date>\d{4}-\d{2}-\d{2})(?:_(?P<n>\d+))?\.json", nwb.session_id
-    )
-
-    if old_re is not None:
-        subject_id, session_date, nwb_suffix = old_re.groups()
-        nwb_suffix = int(nwb_suffix) if nwb_suffix is not None else 0
+    subject_id = nwb.subject.subject_id
+    
+    # -- Block and probability analysis --
+    p_L = df.reward_probabilityL.values
+    p_R = df.reward_probabilityR.values
+    p_contrast = np.max([p_L, p_R], axis=0) / (np.min([p_L, p_R], axis=0) + 1e-6)
+    p_contrast[p_contrast > 100] = 100  # Cap the contrast at 100
+    
+    # Parse effective block
+    block_start_left, block_start_right, block_start_effective = _get_block_starts(p_L, p_R)
+    if 'uncoupled' not in nwb.protocol.lower():
+        if not (len(block_start_left) == len(block_start_right) 
+                and all(block_start_left == block_start_right)):
+            logger.warning("Blocks are not fully aligned in a Coupled task!")
+
+    # -- Metadata dictionary --
+    dict_meta = {
+        'subject_id': subject_id,
+        'session_date': session_date,        
+        'user_name': nwb.experimenter[0],
+        'task': nwb.protocol,
+        'session_start_time': session_start_time,
+
+        
+        # Block structure metrics
+        'p_reward_sum_mean': np.mean(p_L + p_R),
+        'p_reward_sum_std': np.std(p_L + p_R),
+        'p_reward_sum_median': np.median(p_L + p_R),
+        
+        'p_reward_contrast_mean': np.mean(p_contrast),
+        'p_reware_contrast_median': np.median(p_contrast),
+        
+        'effective_block_length_mean': np.mean(np.diff(block_start_effective)),
+        'effective_block_length_std': np.std(np.diff(block_start_effective)),
+        'effective_block_length_median': np.median(np.diff(block_start_effective)),
+        'effective_block_length_min': np.min(np.diff(block_start_effective)),
+        'effective_block_length_max': np.max(np.diff(block_start_effective)),
+        
+        # Duration metrics
+        'duration_gocue_stop_mean': df.loc[:, 'duration_gocue_stop'].mean(),
+        'duration_gocue_stop_std': df.loc[:, 'duration_gocue_stop'].std(),
+        'duration_gocue_stop_median': df.loc[:, 'duration_gocue_stop'].median(),
+        'duration_gocue_stop_min': df.loc[:, 'duration_gocue_stop'].min(),
+        'duration_gocue_stop_max': df.loc[:, 'duration_gocue_stop'].max(),
+
+        'duration_delay_period_mean': df.loc[:, 'duration_delay_period'].mean(),
+        'duration_delay_period_std': df.loc[:, 'duration_delay_period'].std(),
+        'duration_delay_period_median': df.loc[:, 'duration_delay_period'].median(),
+        'duration_delay_period_min': df.loc[:, 'duration_delay_period'].min(),
+        'duration_delay_period_max': df.loc[:, 'duration_delay_period'].max(),
+
+        'duration_iti_mean': df.loc[:, 'duration_iti'].mean(),
+        'duration_iti_std': df.loc[:, 'duration_iti'].std(),
+        'duration_iti_median': df.loc[:, 'duration_iti'].median(),
+        'duration_iti_min': df.loc[:, 'duration_iti'].min(),
+        'duration_iti_max': df.loc[:, 'duration_iti'].max(),
+        
+        # Reward size metrics
+        'reward_volume_left_mean': df.loc[df.reward, 'reward_size_left'].mean(),
+        'reward_volume_right_mean': df.loc[df.reward, 'reward_size_right'].mean(),
+        
+        # Lickspouts movement range (in um)
+        **{f'lickspout_movement_range_{axis}': 
+            np.ptp(df[f'lickspout_position_{axis}']) for axis in 'xyz'},
+        **{f'lickspout_initial_pos_{axis}': 
+            df[f'lickspout_position_{axis}'][0] for axis in 'xyz'},
+        **{f'lickspout_median_pos_{axis}': 
+            np.median(df[f'lickspout_position_{axis}']) for axis in 'xyz'},
+    }
+    
+    # Add flag for old bpod session
+    if 'bpod' in nwb.session_description:
+        dict_meta['old_bpod_session'] = True
+
+    # Create metadata DataFrame
+    df_meta = pd.DataFrame(dict_meta, index=[0])
+    
+    # Add automatic training info
+    if 'auto_train_engaged' in df.columns:       
+        df_meta['auto_train', 'curriculum_name'] = np.nan if df.auto_train_curriculum_name.mode()[0].lower() == 'none' else df.auto_train_curriculum_name.mode()[0]
+        df_meta['auto_train', 'curriculum_version'] = np.nan if df.auto_train_curriculum_version.mode()[0].lower() == 'none' else df.auto_train_curriculum_version.mode()[0]
+        df_meta['auto_train', 'curriculum_schema_version'] = np.nan if df.auto_train_curriculum_schema_version.mode()[0].lower() == 'none' else df.auto_train_curriculum_schema_version.mode()[0]
+        df_meta['auto_train', 'current_stage_actual'] = np.nan if df.auto_train_stage.mode()[0].lower() == 'none' else df.auto_train_stage.mode()[0]
+        df_meta['auto_train', 'if_overriden_by_trainer'] = np.nan if all(df.auto_train_stage_overridden.isna()) else df.auto_train_stage_overridden.mode()[0]
+        
+        # Check consistency of auto train settings
+        df_meta['auto_train', 'if_consistent_within_session'] = len(df.groupby(
+            [col for col in df.columns if 'auto_train' in col]
+        )) == 1
     else:
-        subject_id, session_date, session_json_time = re.match(
-            r"(?P<subject_id>\d+)_(?P<date>\d{4}-\d{2}-\d{2})(?:_(?P<time>.*))\.json",
-            nwb.session_id,
-        ).groups()
-        nwb_suffix = int(session_json_time.replace("-", ""))
-
-    # Verify metadata matches
-    assert subject_id == subject_id_from_meta, (
-        f"Subject name from the metadata ({subject_id_from_meta}) does not match "
-        f"that from json name ({subject_id})!!"
-    )
-    assert session_date == session_date, (
-        f"Session date from the metadata ({session_date}) does not match "
-        f"that from json name ({session_date})!!"
-    )
-
-    # Create session index
-    session_index = pd.MultiIndex.from_tuples(
-        [(subject_id, session_date, nwb_suffix)], names=["subject_id", "session_date", "nwb_suffix"]
-    )
-
-    # Get metadata from nwb.scratch
-    meta_dict = nwb.scratch["metadata"].to_dataframe().iloc[0].to_dict()
-
-    # Calculate performance metrics
+        for field in ['curriculum_name', 'curriculum_version', 'curriculum_schema_version', 'current_stage_actual', 'if_overriden_by_trainer']:
+            df_meta['auto_train', field] = None
+    
+    # -- Performance Metrics --
     n_total_trials = len(df)
     n_finished_trials = (df.animal_response != IGNORE).sum()
-
+    
     # Actual foraging trials (autowater excluded)
     n_total_trials_non_autowater = df.non_autowater_trial.sum()
     n_finished_trials_non_autowater = df.non_autowater_finished_trial.sum()
-
+        
     n_reward_trials_non_autowater = df.reward_non_autowater.sum()
-    reward_rate_non_autowater_finished = (
-        n_reward_trials_non_autowater / n_finished_trials_non_autowater
-        if n_finished_trials_non_autowater > 0
-        else np.nan
-    )
+    reward_rate_non_autowater_finished = n_reward_trials_non_autowater / n_finished_trials_non_autowater if n_finished_trials_non_autowater > 0 else np.nan
 
-    # Calculate foraging efficiency
+    # Foraging efficiency 
     foraging_eff, foraging_eff_random_seed = compute_foraging_efficiency(
-        baited="without bait" not in nwb.protocol.lower(),
-        choice_history=df.animal_response.map({0: 0, 1: 1, 2: np.nan}).values,
-        reward_history=df.rewarded_historyL | df.rewarded_historyR,
-        p_reward=[
-            df.reward_probabilityL.values,
-            df.reward_probabilityR.values,
-        ],
-        random_number=[
-            df.reward_random_number_left.values,
-            df.reward_random_number_right.values,
-        ],
-        autowater_offered=(df.auto_waterL == 1) | (df.auto_waterR == 1),
-    )
-
-    # Override foraging_eff_random_seed for old bpod sessions
-    if "bpod" in nwb.session_description:
-        foraging_eff_random_seed = nwb.get_scratch("metadata")[
-            "foraging_efficiency_with_actual_random_seed"
-        ].values[0]
-
-    finished_rate = (
-        n_finished_trials_non_autowater / n_total_trials_non_autowater
-        if n_total_trials_non_autowater > 0
-        else np.nan
+    baited='without bait' not in nwb.protocol.lower(),
+    choice_history=df.animal_response.map({0: 0, 1: 1, 2: np.nan}).values,
+    reward_history=df.rewarded_historyL | df.rewarded_historyR,
+    p_reward=[
+        df.reward_probabilityL.values,
+        df.reward_probabilityR.values,
+    ],
+    random_number=[
+        df.reward_random_number_left.values,
+        df.reward_random_number_right.values,
+    ],
+    autowater_offered=(df.auto_waterL == 1) | (df.auto_waterR == 1)
     )
 
-    # New Metrics
+    all_lick_number = len(nwb.acquisition['left_lick_time'].timestamps) + len(nwb.acquisition['right_lick_time'].timestamps)
+    
+    # Naive bias calculation
+    n_left = ((df.animal_response == LEFT) & (df.non_autowater_trial)).sum()
+    n_right = ((df.animal_response == RIGHT) & (df.non_autowater_trial)).sum()
+    bias_naive = 2 * (n_right / (n_left + n_right) - 0.5) if n_left + n_right > 0 else np.nan
+    finished_rate = n_finished_trials_non_autowater / n_total_trials_non_autowater if n_total_trials_non_autowater > 0 else np.nan
 
     # Probability chosen calculation
     probability_chosen = []
@@ -138,7 +213,7 @@ def session_metrics(nwb):
     df["probability_chosen"] = probability_chosen
     df["probability_not_chosen"] = probability_not_chosen
 
-    # Calculate the chosen probability
+    # Calculate chosen probability
     average = df["probability_chosen"] - df["probability_not_chosen"]
 
     p_larger_global = max(
@@ -152,66 +227,67 @@ def session_metrics(nwb):
     mean_difference = average.mean()
     chosen_probability = mean_difference / (p_larger_global - p_smaller_global)
 
-    # Pack all data
-    dict_meta = {
-        # Basic metadata
-        "rig": meta_dict["box"],
-        "user_name": nwb.experimenter[0],
-        "task": nwb.protocol,
-        # New metric
-        "chosen_probability": chosen_probability,
-        # Trial counts and rates
-        "total_trials": n_total_trials_non_autowater,
-        "finished_trials": n_finished_trials_non_autowater,
-        "finished_rate": finished_rate,
-        "total_trials_with_autowater": n_total_trials,
-        "finished_trials_with_autowater": n_finished_trials,
-        "finished_rate_with_autowater": n_finished_trials / n_total_trials,
-        # Reward and foraging metrics
-        "reward_trials": n_reward_trials_non_autowater,
-        "reward_rate": reward_rate_non_autowater_finished,
-        "foraging_eff": foraging_eff,
-        "foraging_eff_random_seed": foraging_eff_random_seed,
-        "foraging_performance": foraging_eff * finished_rate,
-        "foraging_performance_random_seed": foraging_eff_random_seed * finished_rate,
-        # Timing metrics
-        "reaction_time_median": df.loc[:, "reaction_time"].median(),
-        "reaction_time_mean": df.loc[:, "reaction_time"].mean(),
-        "early_lick_rate": (df.loc[:, "n_lick_all_delay_period"] > 0).sum() / n_total_trials,
-        # Double dipping metrics
-        "double_dipping_rate_finished_trials": (
-            df.loc[(df.animal_response != IGNORE), "n_lick_switches_gocue_stop"] > 0
-        ).sum()
-        / (df.animal_response != IGNORE).sum(),
-        "double_dipping_rate_finished_reward_trials": (
-            df.loc[df.reward, "n_lick_switches_gocue_stop"] > 0
-        ).sum()
-        / df.reward.sum(),
-        "double_dipping_rate_finished_noreward_trials": (
-            df.loc[
-                (df.animal_response != IGNORE) & (~df.reward),
-                "n_lick_switches_gocue_stop",
-            ]
-            > 0
-        ).sum()
-        / ((df.animal_response != IGNORE) & (~df.reward)).sum(),
+    # Performance dictionary
+    dict_performance = {
+        # Basic performance metrics
+        'total_trials_with_autowater': n_total_trials,
+        'finished_trials_with_autowater': n_finished_trials,
+        'finished_rate_with_autowater': n_finished_trials / n_total_trials,
+        'ignore_rate_with_autowater': 1 - n_finished_trials / n_total_trials,
+        'autowater_collected': (~df.non_autowater_trial & (df.animal_response != IGNORE)).sum(),
+        'autowater_ignored': (~df.non_autowater_trial & (df.animal_response == IGNORE)).sum(),
+        
+        'total_trials': n_total_trials_non_autowater,
+        'finished_trials': n_finished_trials_non_autowater,
+        'ignored_trials': n_total_trials_non_autowater - n_finished_trials_non_autowater,
+        'finished_rate': finished_rate,
+        'ignore_rate': 1 - finished_rate,
+        
+        'reward_trials': n_reward_trials_non_autowater,
+        'reward_rate': reward_rate_non_autowater_finished,
+        'foraging_eff': foraging_eff,
+        'foraging_eff_random_seed': foraging_eff_random_seed,
+        
+        'foraging_performance': foraging_eff * finished_rate,
+        'foraging_performance_random_seed': foraging_eff_random_seed * finished_rate,
+        
+        'bias_naive': bias_naive,
+
+        # New Metrics
+        'chosen_probability': chosen_probability,
+
+        # Lick timing metrics
+        'reaction_time_median': df.loc[:, 'reaction_time'].median(),
+        'reaction_time_mean': df.loc[:, 'reaction_time'].mean(),
+        
+        'early_lick_rate': 
+            (df.loc[:, 'n_lick_all_delay_period'] > 0).sum() / n_total_trials,
+        
+        'invalid_lick_ratio':   
+            (all_lick_number - df.loc[:, 'n_lick_all_gocue_stop'].sum()) / all_lick_number,
+            
         # Lick consistency metrics
-        "lick_consistency_mean_finished_trials": df.loc[
-            (df.animal_response != IGNORE), "n_lick_consistency_gocue_stop"
-        ].mean(),
-        "lick_consistency_mean_finished_reward_trials": df.loc[
-            df.reward, "n_lick_consistency_gocue_stop"
-        ].mean(),
-        "lick_consistency_mean_finished_noreward_trials": df.loc[
-            (df.animal_response != IGNORE) & (~df.reward),
-            "n_lick_consistency_gocue_stop",
-        ].mean(),
-    }
-
-    # Create DataFrame with hierarchical columns
-    df_meta = pd.DataFrame(dict_meta, index=session_index)
-    df_meta.columns = pd.MultiIndex.from_product(
-        [["metadata"], dict_meta.keys()], names=["type", "variable"]
-    )
-
-    return df_meta
+        'double_dipping_rate_finished_trials':  
+            (df.loc[(df.animal_response != IGNORE), 'n_lick_switches_gocue_stop'] > 0).sum() 
+            / (df.animal_response != IGNORE).sum(),
+        'double_dipping_rate_finished_reward_trials':  
+            (df.loc[df.reward, 'n_lick_switches_gocue_stop'] > 0).sum()  
+            / df.reward.sum(),
+        'double_dipping_rate_finished_noreward_trials':   
+            (df.loc[(df.animal_response != IGNORE) & (~df.reward), 'n_lick_switches_gocue_stop'] > 0).sum() 
+            / ((df.animal_response != IGNORE) & (~df.reward)).sum(),
+        'lick_consistency_mean_finished_trials': 
+            df.loc[(df.animal_response != IGNORE), 'n_lick_consistency_gocue_stop'].mean(),
+        'lick_consistency_mean_finished_reward_trials': 
+            df.loc[df.reward, 'n_lick_consistency_gocue_stop'].mean(),
+        'lick_consistency_mean_finished_noreward_trials': 
+            df.loc[(df.animal_response != IGNORE) & (~df.reward), 'n_lick_consistency_gocue_stop'].mean(),
+        }
+        
+    # Generate performance DataFrame
+    df_performance = pd.DataFrame(dict_performance, index=[0])
+
+    # Create session DataFrame
+    session_df = pd.DataFrame({**dict_meta, **dict_performance}, index=[0])
+    
+    return session_df
\ No newline at end of file
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 a19eb3d..1761ace 100644
--- a/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py
+++ b/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py
@@ -10,6 +10,7 @@
 WIN_DUR = 15
 MIN_EVENTS = 2
 
+LEFT, RIGHT, IGNORE = 0, 1, 2
 
 def compute_all_trial_metrics(nwb):
     """
@@ -22,6 +23,13 @@ def compute_all_trial_metrics(nwb):
     choose_right_rate,      fraction of trials where chose right,
                             computed only on trials with a response
 
+    reward columns,         boolean: trials where reward reward is
+                            given, autowater and non_autowater trials
+                            are included
+    lick columns,           duration, delay, and iti periods in session,
+                            right, left, and total lick counts, intertrial
+                            choices: switches, consistency, and reaction time
+
     """
     if not hasattr(nwb, "df_trials"):
         print("You need to compute df_trials: nwb_utils.create_trials_df(nwb)")
@@ -29,6 +37,83 @@ def compute_all_trial_metrics(nwb):
 
     df = nwb.df_trials.copy()
 
+    # --- Add reward-related columns ---
+    df['reward'] = False
+    df.loc[(df.rewarded_historyL | df.rewarded_historyR
+            | df.auto_waterR | df.auto_waterL 
+            & (df.animal_response != IGNORE)) > 0, 'reward'] = True
+    
+    df['reward_non_autowater'] = False
+    df.loc[(df.rewarded_historyL | df.rewarded_historyR), 'reward_non_autowater'] = True
+    
+    df['non_autowater_trial'] = False
+    df.loc[(df.auto_waterL==0) & (df.auto_waterR==0), 'non_autowater_trial'] = True
+    
+    df['non_autowater_finished_trial'] = df['non_autowater_trial'] & (df['animal_response'] != IGNORE)
+    df['ignored_non_autowater'] = df['non_autowater_trial'] & (df['animal_response'] == IGNORE)
+    df['ignored_autowater'] = ~df['non_autowater_trial'] & (df['animal_response'] == IGNORE)
+    
+    # --- Lick-related stats ---    
+    all_left_licks = nwb.acquisition['left_lick_time'].timestamps[:]
+    all_right_licks = nwb.acquisition['right_lick_time'].timestamps[:]
+    
+    # Define the start and stop time for each epoch
+    # Use _in_session columns
+    lick_stats_epochs = {
+        'gocue_stop': ['goCue_start_time_in_session', 'stop_time_in_session'],
+        'delay_period': ['delay_start_time_in_session', 'goCue_start_time_in_session'],
+        'iti': ['start_time_in_session', 'delay_start_time_in_session'],        
+    }
+    
+    # Trial-by-trial counts
+    for i in range(len(df)):
+        for epoch_name, (start_time_name, stop_time_name) in lick_stats_epochs.items():
+            start_time, stop_time = df.loc[i, [start_time_name, stop_time_name]]
+            
+            # Lick analysis for the specific epoch
+            left_licks = all_left_licks[(all_left_licks > start_time) & (all_left_licks < stop_time)]
+            right_licks = all_right_licks[(all_right_licks > start_time) & (all_right_licks < stop_time)]
+            all_licks = np.hstack([left_licks, right_licks])
+            
+            # Lick counts
+            df.loc[i, f'duration_{epoch_name}'] = stop_time - start_time
+            df.loc[i, f'n_lick_left_{epoch_name}'] = len(left_licks)
+            df.loc[i, f'n_lick_right_{epoch_name}'] = len(right_licks)
+            df.loc[i, f'n_lick_all_{epoch_name}'] = len(all_licks)
+            
+            # Lick switches
+            if len(all_licks) > 1:
+                _lick_identity = np.hstack([np.ones(len(left_licks)) * LEFT, np.ones(len(right_licks)) * RIGHT])
+                _lick_identity_sorted = [x for x, _ in sorted(zip(_lick_identity, all_licks), key=lambda pairs: pairs[1])]
+                df.loc[i, f'n_lick_switches_{epoch_name}'] = np.sum(np.diff(_lick_identity_sorted) != 0)
+                
+                # Lick consistency
+                choice = df.loc[i, 'animal_response']
+                df.loc[i, f'n_lick_consistency_{epoch_name}'] = (
+                    np.sum(_lick_identity_sorted == choice) / len(_lick_identity_sorted)
+                    if len(_lick_identity_sorted) > 0 else np.nan
+                )
+            else:
+                df.loc[i, f'n_lick_switches_{epoch_name}'] = 0
+                df.loc[i, f'n_lick_consistency_{epoch_name}'] = np.nan
+            
+            # Special treatment for gocue to stop epoch
+            if epoch_name == 'gocue_stop':
+                # Reaction time
+                first_lick = all_licks.min() if len(all_licks) > 0 else np.nan
+                df.loc[i, 'reaction_time'] = (
+                    first_lick - df.loc[i, 'goCue_start_time_in_session']
+                    if not np.isnan(first_lick) else np.nan
+                )
+                
+                # Handle ignored trials
+                if df.loc[i, 'animal_response'] == IGNORE:
+                    df.loc[i, 'reaction_time'] = np.nan
+                    df.loc[i, 'n_valid_licks_left'] = 0
+                    df.loc[i, 'n_valid_licks_right'] = 0
+                    df.loc[i, 'n_valid_licks_all'] = 0
+
+    # Response and reward rate calculation 
     df["RESPONDED"] = [x in [0, 1] for x in df["animal_response"].values]
     # Rolling fraction of goCues with a response
     df["response_rate"] = (