Source code for allensdk.brain_observatory.behavior.rewards_processing

import pandas as pd
from collections import defaultdict


def get_rewards(data, stimulus_rebase_function):
    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(stimulus_rebase_function(rewards[0][1]))
            rewards_dict["autorewarded"].append('auto_rewarded' in trial['trial_params'])

    df = pd.DataFrame(rewards_dict).set_index('timestamps', drop=True)

    return df