Source code for allensdk.brain_observatory.behavior.rewards_processing

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


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