Source code for allensdk.brain_observatory.behavior.rewards_processing

from typing import Dict
import numpy as np
import pandas as pd

[docs]def get_rewards(data: Dict, timestamps: np.ndarray) -> pd.DataFrame: """ Construct and return a pandas DataFrame containing reward data for this session Parameters --------- data: Dict The dict that results from reading the stimulus pickle file associated with the session timestamps: np.ndarray[1d] A numpy array of timestamps associated with the stimulus frames in this session. timestamps[ii] is the clock time of the iith frame. Returns ------- pd.DataFrame containing the data associated with rewards given in this session """ trial_df = pd.DataFrame(data["items"]["behavior"]["trial_log"]) rewards_dict = {"volume": [], "timestamps": [], "autorewarded": []} for idx, trial in trial_df.iterrows(): rewards = trial["rewards"] # as i write this there can only ever be one reward per trial if rewards: rewards_dict["volume"].append(rewards[0][0]) rewards_dict["timestamps"].append(timestamps[rewards[0][2]]) auto_rwrd = trial["trial_params"]["auto_reward"] rewards_dict["autorewarded"].append(auto_rwrd) df = pd.DataFrame(rewards_dict) return df