From 4945819f1bcea4718c29c5c721e5d3c72a92af19 Mon Sep 17 00:00:00 2001 From: nkeesey <140021626+nkeesey@users.noreply.github.com> Date: Mon, 2 Dec 2024 12:18:49 -0800 Subject: [PATCH] refactoring session_metrics --- .../metrics/session_metrics.py | 430 ++++++++++-------- .../metrics/trial_metrics.py | 181 +++++--- 2 files changed, 350 insertions(+), 261 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 fdb5ff8..90dd665 100644 --- a/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py +++ b/src/aind_dynamic_foraging_basic_analysis/metrics/session_metrics.py @@ -1,11 +1,12 @@ """ Consolidated session metric tool - dfs = nwb_utils.create_trials_df(your_nwb) + df_session_meta = compute_session_metadata(nwb) + df_session = compute_session_metrics(nwb) """ -import re +import logging import numpy as np import pandas as pd @@ -16,12 +17,11 @@ # NOTE: finished_rate_with_autowater is the same as the calculated total response rate LEFT, RIGHT, IGNORE = 0, 1, 2 +logger = logging.getLogger(__name__) -def session_metrics(nwb): - """ - Compute all session metadata and performance metrics - -- Metadata -- +def compute_session_metadata(nwb): + """ block structure metrics, block, contrast, and effective probability metrics duration metrics, gocue, delay period, and iti @@ -31,173 +31,224 @@ def session_metrics(nwb): 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 df_trials: 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() + + df_trials = nwb.df_trials.copy() # 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]))) + 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 = nwb.subject.subject_id - + # -- Block and probability analysis -- - p_L = df.reward_probabilityL.values - p_R = df.reward_probabilityR.values + p_L = df_trials.reward_probabilityL.values + p_R = df_trials.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)): + 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, - - + "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)), - + "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(), - + "duration_gocue_stop_mean": df_trials.loc[:, "duration_gocue_stop"].mean(), + "duration_gocue_stop_std": df_trials.loc[:, "duration_gocue_stop"].std(), + "duration_gocue_stop_median": df_trials.loc[:, "duration_gocue_stop"].median(), + "duration_gocue_stop_min": df_trials.loc[:, "duration_gocue_stop"].min(), + "duration_gocue_stop_max": df_trials.loc[:, "duration_gocue_stop"].max(), + "duration_delay_period_mean": df_trials.loc[:, "duration_delay_period"].mean(), + "duration_delay_period_std": df_trials.loc[:, "duration_delay_period"].std(), + "duration_delay_period_median": df_trials.loc[:, "duration_delay_period"].median(), + "duration_delay_period_min": df_trials.loc[:, "duration_delay_period"].min(), + "duration_delay_period_max": df_trials.loc[:, "duration_delay_period"].max(), + "duration_iti_mean": df_trials.loc[:, "duration_iti"].mean(), + "duration_iti_std": df_trials.loc[:, "duration_iti"].std(), + "duration_iti_median": df_trials.loc[:, "duration_iti"].median(), + "duration_iti_min": df_trials.loc[:, "duration_iti"].min(), + "duration_iti_max": df_trials.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(), - + "reward_volume_left_mean": df_trials.loc[df_trials.reward, "reward_size_left"].mean(), + "reward_volume_right_mean": df_trials.loc[df_trials.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'}, + **{ + f"lickspout_movement_range_{axis}": np.ptp(df_trials[f"lickspout_position_{axis}"]) + for axis in "xyz" + }, + **{ + f"lickspout_initial_pos_{axis}": df_trials[f"lickspout_position_{axis}"][0] + for axis in "xyz" + }, + **{ + f"lickspout_median_pos_{axis}": np.median(df_trials[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 + if "bpod" in nwb.session_description: + dict_meta["old_bpod_session"] = True # Create metadata DataFrame - df_meta = pd.DataFrame(dict_meta, index=[0]) - + df_session_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] - + if "auto_train_engaged" in df_trials.columns: + df_session_meta["auto_train", "curriculum_name"] = ( + np.nan + if df_trials.auto_train_curriculum_name.mode()[0].lower() == "none" + else df_trials.auto_train_curriculum_name.mode()[0] + ) + df_session_meta["auto_train", "curriculum_version"] = ( + np.nan + if df_trials.auto_train_curriculum_version.mode()[0].lower() == "none" + else df_trials.auto_train_curriculum_version.mode()[0] + ) + df_session_meta["auto_train", "curriculum_schema_version"] = ( + np.nan + if df_trials.auto_train_curriculum_schema_version.mode()[0].lower() == "none" + else df_trials.auto_train_curriculum_schema_version.mode()[0] + ) + df_session_meta["auto_train", "current_stage_actual"] = ( + np.nan + if df_trials.auto_train_stage.mode()[0].lower() == "none" + else df_trials.auto_train_stage.mode()[0] + ) + df_session_meta["auto_train", "if_overriden_by_trainer"] = ( + np.nan + if all(df_trials.auto_train_stage_overridden.isna()) + else df_trials.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 + df_session_meta["auto_train", "if_consistent_within_session"] = ( + len(df_trials.groupby([col for col in df_trials.columns if "auto_train" in col])) == 1 + ) else: - for field in ['curriculum_name', 'curriculum_version', 'curriculum_schema_version', 'current_stage_actual', 'if_overriden_by_trainer']: - df_meta['auto_train', field] = None - + for field in [ + "curriculum_name", + "curriculum_version", + "curriculum_schema_version", + "current_stage_actual", + "if_overriden_by_trainer", + ]: + df_session_meta["auto_train", field] = None + + return df_session_meta + + +def compute_session_metrics(nwb): + """ + Compute all session metadata and performance metrics + + 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 df_trials: nwb_utils.create_trials_df(nwb)") + return + + df_trials = nwb.df_trials.copy() + + # Add session Metdata + df_session_meta = compute_session_metadata(nwb) + # -- Performance Metrics -- - n_total_trials = len(df) - n_finished_trials = (df.animal_response != IGNORE).sum() - + n_total_trials = len(df_trials) + n_finished_trials = (df_trials.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 + n_total_trials_non_autowater = df_trials.non_autowater_trial.sum() + n_finished_trials_non_autowater = df_trials.non_autowater_finished_trial.sum() + + n_reward_trials_non_autowater = df_trials.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 + ) - # 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) + baited="without bait" not in nwb.protocol.lower(), + choice_history=df_trials.animal_response.map({0: 0, 1: 1, 2: np.nan}).values, + reward_history=df_trials.rewarded_historyL | df_trials.rewarded_historyR, + p_reward=[ + df_trials.reward_probabilityL.values, + df_trials.reward_probabilityR.values, + ], + random_number=[ + df_trials.reward_random_number_left.values, + df_trials.reward_random_number_right.values, + ], + autowater_offered=(df_trials.auto_waterL == 1) | (df_trials.auto_waterR == 1), + ) + + all_lick_number = len(nwb.acquisition["left_lick_time"].timestamps) + len( + nwb.acquisition["right_lick_time"].timestamps ) - 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() + n_left = ((df_trials.animal_response == LEFT) & (df_trials.non_autowater_trial)).sum() + n_right = ((df_trials.animal_response == RIGHT) & (df_trials.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 + 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 = [] probability_not_chosen = [] - for _, row in df.iterrows(): + for _, row in df_trials.iterrows(): if row.animal_response == 2: probability_chosen.append(np.nan) probability_not_chosen.append(np.nan) @@ -210,18 +261,18 @@ def _get_block_starts(p_L, p_R): probability_chosen.append(row.reward_probabilityR) probability_not_chosen.append(row.reward_probabilityL) - df["probability_chosen"] = probability_chosen - df["probability_not_chosen"] = probability_not_chosen + df_trials["probability_chosen"] = probability_chosen + df_trials["probability_not_chosen"] = probability_not_chosen # Calculate chosen probability - average = df["probability_chosen"] - df["probability_not_chosen"] + average = df_trials["probability_chosen"] - df_trials["probability_not_chosen"] p_larger_global = max( - df["probability_chosen"].max(), df["probability_not_chosen"].max() + df_trials["probability_chosen"].max(), df_trials["probability_not_chosen"].max() ) p_smaller_global = min( - df["probability_chosen"].min(), df["probability_not_chosen"].min() + df_trials["probability_chosen"].min(), df_trials["probability_not_chosen"].min() ) mean_difference = average.mean() @@ -230,64 +281,69 @@ def _get_block_starts(p_L, p_R): # 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, - + "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_trials.non_autowater_trial & (df_trials.animal_response != IGNORE) + ).sum(), + "autowater_ignored": ( + ~df_trials.non_autowater_trial & (df_trials.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, - + "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, - + "reaction_time_median": df_trials.loc[:, "reaction_time"].median(), + "reaction_time_mean": df_trials.loc[:, "reaction_time"].mean(), + "early_lick_rate": (df_trials.loc[:, "n_lick_all_delay_period"] > 0).sum() / n_total_trials, + "invalid_lick_ratio": (all_lick_number - df_trials.loc[:, "n_lick_all_gocue_stop"].sum()) + / all_lick_number, # Lick consistency 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(), - '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]) + "double_dipping_rate_finished_trials": ( + df_trials.loc[(df_trials.animal_response != IGNORE), "n_lick_switches_gocue_stop"] > 0 + ).sum() + / (df_trials.animal_response != IGNORE).sum(), + "double_dipping_rate_finished_reward_trials": ( + df_trials.loc[df_trials.reward, "n_lick_switches_gocue_stop"] > 0 + ).sum() + / df_trials.reward.sum(), + "double_dipping_rate_finished_noreward_trials": ( + df_trials.loc[ + (df_trials.animal_response != IGNORE) & (~df_trials.reward), + "n_lick_switches_gocue_stop", + ] + > 0 + ).sum() + / ((df_trials.animal_response != IGNORE) & (~df_trials.reward)).sum(), + "lick_consistency_mean_finished_trials": df_trials.loc[ + (df_trials.animal_response != IGNORE), "n_lick_consistency_gocue_stop" + ].mean(), + "lick_consistency_mean_finished_reward_trials": df_trials.loc[ + df_trials.reward, "n_lick_consistency_gocue_stop" + ].mean(), + "lick_consistency_mean_finished_noreward_trials": df_trials.loc[ + (df_trials.animal_response != IGNORE) & (~df_trials.reward), + "n_lick_consistency_gocue_stop", + ].mean(), + } + + # Create performance Dataframe + df_session_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 + df_session = pd.concat([df_session_meta, df_session_performance], axis=1).reset_index(drop=True) + + return df_session 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 1761ace..2478640 100644 --- a/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py +++ b/src/aind_dynamic_foraging_basic_analysis/metrics/trial_metrics.py @@ -12,7 +12,8 @@ LEFT, RIGHT, IGNORE = 0, 1, 2 -def compute_all_trial_metrics(nwb): + +def compute_trial_metrics(nwb): """ Computes all trial by trial metrics @@ -35,112 +36,144 @@ def compute_all_trial_metrics(nwb): print("You need to compute df_trials: nwb_utils.create_trials_df(nwb)") return - df = nwb.df_trials.copy() + df_trials = 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[:] - + df_trials["reward"] = False + df_trials.loc[ + ( + df_trials.rewarded_historyL + | df_trials.rewarded_historyR + | df_trials.auto_waterR + | df_trials.auto_waterL + ) + & (df_trials.animal_response != IGNORE), + "reward", + ] = True + + df_trials["reward_non_autowater"] = False + df_trials.loc[ + (df_trials.rewarded_historyL | df_trials.rewarded_historyR), "reward_non_autowater" + ] = True + + df_trials["non_autowater_trial"] = False + df_trials.loc[ + (df_trials.auto_waterL == 0) & (df_trials.auto_waterR == 0), "non_autowater_trial" + ] = True + + df_trials["non_autowater_finished_trial"] = df_trials["non_autowater_trial"] & ( + df_trials["animal_response"] != IGNORE + ) + df_trials["ignored_non_autowater"] = df_trials["non_autowater_trial"] & ( + df_trials["animal_response"] == IGNORE + ) + df_trials["ignored_autowater"] = ~df_trials["non_autowater_trial"] & ( + df_trials["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 + # Using _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'], + "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 i in range(len(df_trials)): 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]] - + start_time, stop_time = df_trials.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)] + 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) - + df_trials.loc[i, f"duration_{epoch_name}"] = stop_time - start_time + df_trials.loc[i, f"n_lick_left_{epoch_name}"] = len(left_licks) + df_trials.loc[i, f"n_lick_right_{epoch_name}"] = len(right_licks) + df_trials.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_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_trials.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}'] = ( + choice = df_trials.loc[i, "animal_response"] + df_trials.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 + 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 - + df_trials.loc[i, f"n_lick_switches_{epoch_name}"] = 0 + df_trials.loc[i, f"n_lick_consistency_{epoch_name}"] = np.nan + # Special treatment for gocue to stop epoch - if epoch_name == 'gocue_stop': + 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 + df_trials.loc[i, "reaction_time"] = ( + first_lick - df_trials.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] + if df_trials.loc[i, "animal_response"] == IGNORE: + df_trials.loc[i, "reaction_time"] = np.nan + df_trials.loc[i, "n_valid_licks_left"] = 0 + df_trials.loc[i, "n_valid_licks_right"] = 0 + df_trials.loc[i, "n_valid_licks_all"] = 0 + + df_trials["RESPONDED"] = [x in [0, 1] for x in df_trials["animal_response"].values] # Rolling fraction of goCues with a response - df["response_rate"] = ( - df["RESPONDED"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() + df_trials["response_rate"] = ( + df_trials["RESPONDED"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() ) # Rolling fraction of goCues with a response - df["gocue_reward_rate"] = ( - df["earned_reward"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() + df_trials["gocue_reward_rate"] = ( + df_trials["earned_reward"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() ) # Rolling fraction of responses with a response - df["RESPONSE_REWARD"] = [ - x[0] if x[1] else np.nan for x in zip(df["earned_reward"], df["RESPONDED"]) + df_trials["RESPONSE_REWARD"] = [ + x[0] if x[1] else np.nan for x in zip(df_trials["earned_reward"], df_trials["RESPONDED"]) ] - df["response_reward_rate"] = ( - df["RESPONSE_REWARD"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() + df_trials["response_reward_rate"] = ( + df_trials["RESPONSE_REWARD"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() ) # Rolling fraction of choosing right - df["WENT_RIGHT"] = [x if x in [0, 1] else np.nan for x in df["animal_response"]] - df["choose_right_rate"] = ( - df["WENT_RIGHT"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() + df_trials["WENT_RIGHT"] = [x if x in [0, 1] else np.nan for x in df_trials["animal_response"]] + df_trials["choose_right_rate"] = ( + df_trials["WENT_RIGHT"].rolling(WIN_DUR, min_periods=MIN_EVENTS, center=True).mean() ) # Clean up temp columns - drop_cols = ["RESPONDED", "RESPONSE_REWARD", "WENT_RIGHT"] - df = df.drop(columns=drop_cols) + drop_cols = [ + "RESPONDED", + "RESPONSE_REWARD", + "WENT_RIGHT", + ] + df_trials = df_trials.drop(columns=drop_cols) - return df + return df_trials