AllenSDK/allensdk/brain_observatory/ecephys/data_objects/trials.py at 05328c2f4f4079af4df2325bea9ebf143922b875 · galenlynch/AllenSDK · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
from typing import Union, Tuple, List

import numpy as np

from allensdk.brain_observatory.behavior.data_objects.trials.trial import (
    Trial)
from allensdk.brain_observatory.behavior.data_objects.\
    trials.trials import Trials


class VBNTrial(Trial):

    def calculate_change_frame(
            self,
            event_dict: dict,
            go: bool,
            catch: bool,
            auto_rewarded: bool) -> Union[int, float]:

        """
        Calculate the frame index of a stimulus change
        associated with a specific event.

        Parameters
        ----------
        event_dict: dict
            Dictionary of trial events in the well-known `pkl` file
        go: bool
            True if "go" trial, False otherwise. Mutually exclusive with
            `catch`.
        catch: bool
            True if "catch" trial, False otherwise. Mutually exclusive
            with `go.`
        auto_rewarded: bool
            True if "auto_rewarded" trial, False otherwise.

        Returns
        -------
        change_frame: Union[int, float]
            Index of the change frame; NaN if there is no change

        Notes
        -----
        This is its own method so that child classes of Trial
        can implement different logic as needed.
        """

        if go or auto_rewarded:
            change_frame = event_dict.get(('stimulus_changed', ''))['frame']
        elif catch:
            change_frame = event_dict.get(('sham_change', ''))['frame']
        else:
            change_frame = float("nan")

        return change_frame

    def add_change_time(self, trial_dict: dict) -> Tuple[dict, float]:
        """
        Add change_time_no_display_delay to a dict representing
        a single trial.

        This implementation will just take change_frame and
        select the value of self._stimulus_timestamps corresponding
        to that frame.

        Parameters
        ----------
        trial_dict:
            dict containing all trial parameters except
            change_time

        Returns
        -------
        trial_dict:
            Same as input, except change_time_no_display_delay
            field has been added

        change_time: float
            The change time value that was added
            (this is returned separately so that child classes have the
            option of naming the column something different than
            'change_time')

        Note
        ----
        Modified trial_dict in-place, in addition to returning it
        """
        change_frame = trial_dict['change_frame']
        if np.isnan(change_frame):
            change_time = np.nan
        else:
            no_delay = self._stimulus_timestamps.subtract_monitor_delay()
            change_frame = int(change_frame)
            change_time = no_delay.value[change_frame]

        trial_dict['change_time_no_display_delay'] = change_time
        return trial_dict, change_time


class VBNTrials(Trials):

    @classmethod
    def trial_class(cls):
        """
        Return the class to be used to represent a single Trial
        """
        return VBNTrial

    @classmethod
    def columns_to_output(cls) -> List[str]:
        """
        Return the list of columns to be output in this table
        """
        return ['initial_image_name', 'change_image_name',
                'stimulus_change', 'change_time_no_display_delay',
                'go', 'catch', 'lick_times', 'response_time',
                'reward_time', 'reward_volume',
                'hit', 'false_alarm', 'miss', 'correct_reject',
                'aborted', 'auto_rewarded', 'change_frame',
                'start_time', 'stop_time', 'trial_length']

    @property
    def change_time(self):
        return self.data['change_time_no_display_delay']