Source code for allensdk.brain_observatory.behavior.data_objects.rewards

from typing import Optional

import pandas as pd
import numpy as np
from pynwb import NWBFile, TimeSeries, ProcessingModule

from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.core import DataObject
from allensdk.brain_observatory.behavior.data_objects import StimulusTimestamps
from allensdk.core import \
    NwbReadableInterface
from allensdk.brain_observatory.behavior.data_files.stimulus_file import \
    StimulusFileReadableInterface
from allensdk.core import \
    NwbWritableInterface




class Rewards(DataObject, StimulusFileReadableInterface, NwbReadableInterface,
              NwbWritableInterface):
    def __init__(self, rewards: pd.DataFrame):
        super().__init__(name='rewards', value=rewards)



    @classmethod
    def from_stimulus_file(
            cls, stimulus_file: BehaviorStimulusFile,
            stimulus_timestamps: StimulusTimestamps) -> "Rewards":
        """Get reward data from pkl file, based on timestamps
        (not sync file).
        """

        if not np.isclose(stimulus_timestamps.monitor_delay, 0.0):
            msg = ("Instantiating rewards with monitor_delay = "
                   f"{stimulus_timestamps.monitor_delay: .2e}; "
                   "monitor_delay should be zero for Rewards "
                   "data object")
            raise RuntimeError(msg)

        data = stimulus_file.data

        trial_df = pd.DataFrame(data["items"]["behavior"]["trial_log"])
        rewards_dict = {"volume": [], "timestamps": [], "auto_rewarded": []}
        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_timestamps.value[rewards[0][2]])
                auto_rwrd = trial["trial_params"]["auto_reward"]
                rewards_dict["auto_rewarded"].append(auto_rwrd)

        df = pd.DataFrame(rewards_dict)
        return cls(rewards=df)




    @classmethod
    def from_nwb(cls, nwbfile: NWBFile) -> Optional["Rewards"]:
        if 'rewards' in nwbfile.processing:
            rewards = nwbfile.processing['rewards']
            time = rewards.get_data_interface('autorewarded').timestamps[:]
            autorewarded = rewards.get_data_interface('autorewarded').data[:]
            volume = rewards.get_data_interface('volume').data[:]
        else:
            volume = []
            time = []
            autorewarded = []

        df = pd.DataFrame({
                'volume': volume,
                'timestamps': time,
                'auto_rewarded': autorewarded})
        return cls(rewards=df)




    def to_nwb(self, nwbfile: NWBFile) -> NWBFile:

        # If there is no rewards data, do not
        # write anything to the NWB file (this
        # is expected for passive sessions)
        if len(self.value['timestamps']) == 0:
            return nwbfile

        reward_volume_ts = TimeSeries(
            name='volume',
            data=self.value['volume'].values,
            timestamps=self.value['timestamps'].values,
            unit='mL'
        )

        autorewarded_ts = TimeSeries(
            name='autorewarded',
            data=self.value['auto_rewarded'].values,
            timestamps=reward_volume_ts.timestamps,
            unit='mL'
        )

        rewards_mod = ProcessingModule('rewards',
                                       'Licking behavior processing module')
        rewards_mod.add_data_interface(reward_volume_ts)
        rewards_mod.add_data_interface(autorewarded_ts)
        nwbfile.add_processing_module(rewards_mod)

        return nwbfile