from typing import List, Optional, Tuple
import numpy as np
import pandas as pd
from allensdk.brain_observatory import dict_to_indexed_array
from allensdk.brain_observatory.behavior.data_files import (
BehaviorStimulusFile,
SyncFile,
)
from allensdk.brain_observatory.behavior.data_files.stimulus_file import (
StimulusFileReadableInterface,
)
from allensdk.brain_observatory.behavior.data_objects import StimulusTimestamps
from allensdk.brain_observatory.behavior.data_objects.licks import Licks
from allensdk.brain_observatory.behavior.data_objects.rewards import Rewards
from allensdk.brain_observatory.behavior.data_objects.task_parameters import (
TaskParameters,
)
from allensdk.brain_observatory.behavior.data_objects.trials.trial import Trial
from allensdk.brain_observatory.behavior.dprime import (
get_false_alarm_rate,
get_hit_rate,
get_rolling_dprime,
get_trial_count_corrected_false_alarm_rate,
get_trial_count_corrected_hit_rate,
)
from allensdk.core import (
DataObject,
NwbReadableInterface,
NwbWritableInterface,
)
from allensdk.core.dataframe_utils import enforce_df_int_typing
from pynwb import NWBFile
[docs]
class Trials(
DataObject,
StimulusFileReadableInterface,
NwbReadableInterface,
NwbWritableInterface,
):
[docs]
@classmethod
def trial_class(cls):
"""
Return the class to be used to represent a single Trial
"""
return Trial
def __init__(self, trials: pd.DataFrame, response_window_start: float):
"""
Parameters
----------
trials
response_window_start
[seconds] relative to the non-display-lag-compensated presentation
of the change-image
"""
trials = trials.rename(columns={"stimulus_change": "is_change"})
super().__init__(name="trials", value=None, is_value_self=True)
trials = enforce_df_int_typing(trials, ["change_frame"])
self._trials = trials
self._response_window_start = response_window_start
@property
def data(self) -> pd.DataFrame:
return self._trials
@property
def trial_count(self) -> int:
"""Number of trials
(including all 'go', 'catch', and 'aborted' trials)"""
return self._trials.shape[0]
@property
def go_trial_count(self) -> int:
"""Number of 'go' trials"""
return self._trials["go"].sum()
@property
def catch_trial_count(self) -> int:
"""Number of 'catch' trials"""
return self._trials["catch"].sum()
@property
def hit_trial_count(self) -> int:
"""Number of trials with a hit behavior response"""
return self._trials["hit"].sum()
@property
def miss_trial_count(self) -> int:
"""Number of trials with a hit behavior response"""
return self._trials["miss"].sum()
@property
def false_alarm_trial_count(self) -> int:
"""Number of trials where the mouse had a false alarm
behavior response"""
return self._trials["false_alarm"].sum()
@property
def correct_reject_trial_count(self) -> int:
"""Number of trials with a correct reject behavior
response"""
return self._trials["correct_reject"].sum()
[docs]
def to_nwb(self, nwbfile: NWBFile) -> NWBFile:
trials = self.data
order = list(trials.index)
for _, row in trials[["start_time", "stop_time"]].iterrows():
row_dict = row.to_dict()
nwbfile.add_trial(**row_dict)
for c in trials.columns:
if c in ["start_time", "stop_time"]:
continue
index, data = dict_to_indexed_array(trials[c].to_dict(), order)
if data.dtype == "<U1": # data type is composed of unicode
# characters
data = trials[c].tolist()
if not len(data) == len(order):
if len(data) == 0:
data = [""]
nwbfile.add_trial_column(
name=c,
description="NOT IMPLEMENTED: %s" % c,
data=data,
index=index,
)
else:
nwbfile.add_trial_column(
name=c, description="NOT IMPLEMENTED: %s" % c, data=data
)
return nwbfile
[docs]
@classmethod
def from_nwb(cls, nwbfile: NWBFile) -> "Trials":
trials = nwbfile.trials.to_dataframe()
if "lick_events" in trials.columns:
trials.drop("lick_events", inplace=True, axis=1)
trials.index = trials.index.rename("trials_id")
return cls(
trials=trials,
response_window_start=TaskParameters.from_nwb(
nwbfile=nwbfile
).response_window_sec[0],
)
[docs]
@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",
"go",
"catch",
"lick_times",
"response_time",
"response_latency",
"reward_time",
"reward_volume",
"hit",
"false_alarm",
"miss",
"correct_reject",
"aborted",
"auto_rewarded",
"change_frame",
"start_time",
"stop_time",
"trial_length",
]
[docs]
@classmethod
def from_stimulus_file(
cls,
stimulus_file: BehaviorStimulusFile,
stimulus_timestamps: StimulusTimestamps,
licks: Licks,
rewards: Rewards,
sync_file: Optional[SyncFile] = None,
) -> "Trials":
bsf = stimulus_file.data
stimuli = bsf["items"]["behavior"]["stimuli"]
trial_log = bsf["items"]["behavior"]["trial_log"]
trial_bounds = cls._get_trial_bounds(trial_log=trial_log)
all_trial_data = [None] * len(trial_log)
for idx, trial in enumerate(trial_log):
trial_start, trial_end = trial_bounds[idx]
t = cls.trial_class()(
trial=trial,
start=trial_start,
end=trial_end,
behavior_stimulus_file=stimulus_file,
index=idx,
stimulus_timestamps=stimulus_timestamps,
licks=licks,
rewards=rewards,
stimuli=stimuli,
sync_file=sync_file,
)
all_trial_data[idx] = t.data
trials = pd.DataFrame(all_trial_data).set_index("trial")
trials.index = trials.index.rename("trials_id")
# Order/Filter columns
trials = trials[cls.columns_to_output()]
return cls(
trials=trials,
response_window_start=TaskParameters.from_stimulus_file(
stimulus_file=stimulus_file
).response_window_sec[0],
)
@staticmethod
def _get_trial_bounds(trial_log: List) -> List[Tuple[int, int]]:
"""
Adjust trial boundaries from a trial_log so that there is no dead time
between trials.
Parameters
----------
trial_log: list
The trial_log read in from the well known behavior stimulus
pickle file
Returns
-------
list
Each element in the list is a tuple of the form
(start_frame, end_frame) so that the ith element
of the list gives the start and end frames of
the ith trial. The endframe of the last trial will
be -1, indicating that it should map to the last
timestamp in the session
"""
start_frames = []
for trial in trial_log:
start_f = None
for event in trial["events"]:
if event[0] == "trial_start":
start_f = event[-1]
break
if start_f is None:
msg = "Could not find a 'trial_start' event "
msg += "for all trials in the trial log\n"
msg += f"{trial}"
raise ValueError(msg)
if len(start_frames) > 0 and start_f < start_frames[-1]:
msg = "'trial_start' frames in trial log "
msg += "are not in ascending order"
msg += f"\ntrial_log: {trial_log}"
raise ValueError(msg)
start_frames.append(start_f)
end_frames = [idx for idx in start_frames[1:] + [-1]]
return list([(s, e) for s, e in zip(start_frames, end_frames)])
@property
def index(self) -> pd.Index:
return self.data.index
@property
def change_time(self) -> pd.Series:
if "change_time" in self.data:
return self.data["change_time"]
elif "change_time_no_display_delay" in self.data:
return self.data["change_time_no_display_delay"]
@property
def lick_times(self) -> pd.Series:
return self.data["lick_times"]
@property
def start_time(self) -> pd.Series:
return self.data["start_time"]
@property
def aborted(self) -> pd.Series:
return self.data["aborted"]
@property
def hit(self) -> pd.Series:
return self.data["hit"]
@property
def miss(self) -> pd.Series:
return self.data["miss"]
@property
def false_alarm(self) -> pd.Series:
return self.data["false_alarm"]
@property
def correct_reject(self) -> pd.Series:
return self.data["correct_reject"]
@property
def rolling_performance(self) -> pd.DataFrame:
"""Return a DataFrame containing trial by trial behavior response
performance metrics.
Returns
-------
pd.DataFrame
A pandas DataFrame containing:
trials_id [index]:
Index of the trial. All trials, including aborted trials,
are assigned an index starting at 0 for the first trial.
reward_rate:
Rewards earned in the previous 25 trials, normalized by
the elapsed time of the same 25 trials. Units are
rewards/minute.
hit_rate_raw:
Fraction of go trials where the mouse licked in the
response window, calculated over the previous 100
non-aborted trials. Without trial count correction applied.
hit_rate:
Fraction of go trials where the mouse licked in the
response window, calculated over the previous 100
non-aborted trials. With trial count correction applied.
false_alarm_rate_raw:
Fraction of catch trials where the mouse licked in the
response window, calculated over the previous 100
non-aborted trials. Without trial count correction applied.
false_alarm_rate:
Fraction of catch trials where the mouse licked in
the response window, calculated over the previous 100
non-aborted trials. Without trial count correction applied.
rolling_dprime:
d prime calculated using the rolling hit_rate and
rolling false_alarm _rate.
"""
reward_rate = self.calculate_reward_rate()
# Indices to build trial metrics dataframe:
trials_index = self.data.index
not_aborted_index = self.data[np.logical_not(self.aborted)].index
# Initialize dataframe:
performance_metrics_df = pd.DataFrame(index=trials_index)
# Reward rate:
performance_metrics_df["reward_rate"] = pd.Series(
reward_rate, index=self.data.index
)
# Hit rate raw:
hit_rate_raw = get_hit_rate(
hit=self.hit, miss=self.miss, aborted=self.aborted
)
performance_metrics_df["hit_rate_raw"] = pd.Series(
hit_rate_raw, index=not_aborted_index
)
# Hit rate with trial count correction:
hit_rate = get_trial_count_corrected_hit_rate(
hit=self.hit, miss=self.miss, aborted=self.aborted
)
performance_metrics_df["hit_rate"] = pd.Series(
hit_rate, index=not_aborted_index
)
# False-alarm rate raw:
false_alarm_rate_raw = get_false_alarm_rate(
false_alarm=self.false_alarm,
correct_reject=self.correct_reject,
aborted=self.aborted,
)
performance_metrics_df["false_alarm_rate_raw"] = pd.Series(
false_alarm_rate_raw, index=not_aborted_index
)
# False-alarm rate with trial count correction:
false_alarm_rate = get_trial_count_corrected_false_alarm_rate(
false_alarm=self.false_alarm,
correct_reject=self.correct_reject,
aborted=self.aborted,
)
performance_metrics_df["false_alarm_rate"] = pd.Series(
false_alarm_rate, index=not_aborted_index
)
# Rolling-dprime:
is_passive_session = (self.data["reward_volume"] == 0).all() and (
self.data["lick_times"].apply(lambda x: len(x)) == 0
).all()
if is_passive_session:
# It does not make sense to calculate d' for a passive session
# So just set it to zeros
rolling_dprime = np.zeros(len(hit_rate))
else:
rolling_dprime = get_rolling_dprime(hit_rate, false_alarm_rate)
performance_metrics_df["rolling_dprime"] = pd.Series(
rolling_dprime, index=not_aborted_index
)
return performance_metrics_df
def _calculate_response_latency_list(self) -> List:
"""per trial, determines a response latency
Returns
-------
response_latency_list: List
len() = trials.shape[0]
value is 'inf' if there are no valid licks in the trial
Note
-----
response_window_start is listed as
"relative to the non-display-lag-compensated..." because it
comes directly from the stimulus file, which knows nothing
about the display lag. However, response_window_start is
only ever compared to the difference between
trial.lick_times and trial.change_time, both of which are
corrected for monitor delay, so it does not matter
(the two instance of monitor delay cancel out in the
difference).
"""
df = pd.DataFrame(
{"lick_times": self.lick_times, "change_time": self.change_time}
)
df["valid_response_licks"] = df.apply(
lambda trial: [
lt
for lt in trial["lick_times"]
if lt - trial["change_time"] > self._response_window_start
],
axis=1,
)
response_latency = df.apply(
lambda trial: trial["valid_response_licks"][0]
- trial["change_time"]
if len(trial["valid_response_licks"]) > 0
else float("inf"),
axis=1,
)
return response_latency.tolist()
[docs]
def calculate_reward_rate(
self, window=0.75, trial_window=25, initial_trials=10
):
response_latency = self._calculate_response_latency_list()
starttime = self.start_time.values
assert len(response_latency) == len(starttime)
df = pd.DataFrame(
{"response_latency": response_latency, "starttime": starttime}
)
# adds a column called reward_rate to the input dataframe
# the reward_rate column contains a rolling average of rewards/min
# window sets the window in which a response is considered correct,
# so a window of 1.0 means licks before 1.0 second are considered
# correct
# Reorganized into this unit-testable form by Nick Cain April 25 2019
reward_rate = np.zeros(len(df))
# make the initial reward rate infinite,
# so that you include the first trials automatically.
reward_rate[:initial_trials] = np.inf
for trial_number in range(initial_trials, len(df)):
min_index = np.max((0, trial_number - trial_window))
max_index = np.min((trial_number + trial_window, len(df)))
df_roll = df.iloc[min_index:max_index]
# get a rolling number of correct trials
correct = len(df_roll[df_roll.response_latency < window])
# get the time elapsed over the trials
time_elapsed = (
df_roll.starttime.iloc[-1] - df_roll.starttime.iloc[0]
)
# calculate the reward rate, rewards/min
reward_rate_on_this_lap = correct / time_elapsed * 60
reward_rate[trial_number] = reward_rate_on_this_lap
reward_rate[np.isinf(reward_rate)] = float("nan")
return reward_rate
def _get_engaged_trials(
self, engaged_trial_reward_rate_threshold: float = 2.0
) -> pd.Series:
"""
Gets `Series` where each trial that is considered "engaged" is set to
`True`
Parameters
----------
engaged_trial_reward_rate_threshold:
The number of rewards per minute that needs to be attained
before a subject is considered 'engaged', by default 2.0
Returns
-------
`pd.Series`
"""
rolling_performance = self.rolling_performance
engaged_trial_mask = (
rolling_performance["reward_rate"]
> engaged_trial_reward_rate_threshold
)
return engaged_trial_mask
[docs]
def get_engaged_trial_count(
self, engaged_trial_reward_rate_threshold: float = 2.0
) -> int:
"""Gets count of trials considered "engaged"
Parameters
----------
engaged_trial_reward_rate_threshold:
The number of rewards per minute that needs to be attained
before a subject is considered 'engaged', by default 2.0
Returns
-------
count of trials considered "engaged"
"""
engaged_trials = self._get_engaged_trials(
engaged_trial_reward_rate_threshold=(
engaged_trial_reward_rate_threshold
)
)
return engaged_trials.sum()