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