Source code for allensdk.brain_observatory.behavior.data_objects.task_parameters
from enum import Enum
import numpy as np
from typing import List, Optional
from pynwb import NWBFile
from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.core import DataObject
from allensdk.core import \
NwbReadableInterface
from allensdk.brain_observatory.behavior.data_files.stimulus_file import \
StimulusFileReadableInterface
from allensdk.core import \
NwbWritableInterface
from allensdk.brain_observatory.behavior.schemas import \
BehaviorTaskParametersSchema
from allensdk.brain_observatory.nwb import load_pynwb_extension
from allensdk.brain_observatory.behavior.utils.metadata_parsers import ( # noqa E501
parse_stimulus_set
)
[docs]
class BehaviorStimulusType(Enum):
IMAGES = 'images'
GRATING = 'grating'
[docs]
class StimulusDistribution(Enum):
EXPONENTIAL = 'exponential'
GEOMETRIC = 'geometric'
[docs]
class TaskType(Enum):
CHANGE_DETECTION = 'change detection'
[docs]
class TaskParameters(DataObject, StimulusFileReadableInterface,
NwbReadableInterface, NwbWritableInterface):
def __init__(self,
blank_duration_sec: List[float],
stimulus_duration_sec: float,
omitted_flash_fraction: float,
response_window_sec: List[float],
reward_volume: float,
auto_reward_volume: float,
session_type: str,
stimulus: str,
stimulus_distribution: StimulusDistribution,
task_type: TaskType,
n_stimulus_frames: int,
stimulus_name: Optional[str] = None):
super().__init__(name='task_parameters', value=None,
is_value_self=True)
self._blank_duration_sec = blank_duration_sec
self._stimulus_duration_sec = stimulus_duration_sec
self._omitted_flash_fraction = omitted_flash_fraction
self._response_window_sec = response_window_sec
self._reward_volume = reward_volume
self._auto_reward_volume = auto_reward_volume
self._session_type = session_type
self._stimulus = BehaviorStimulusType(stimulus)
self._stimulus_distribution = StimulusDistribution(
stimulus_distribution)
self._task = TaskType(task_type)
self._n_stimulus_frames = n_stimulus_frames
self._stimulus_name = stimulus_name
self._image_set = parse_stimulus_set(session_type)
@property
def blank_duration_sec(self) -> List[float]:
return self._blank_duration_sec
@property
def stimulus_duration_sec(self) -> float:
return self._stimulus_duration_sec
@property
def omitted_flash_fraction(self) -> float:
return self._omitted_flash_fraction
@property
def response_window_sec(self) -> List[float]:
return self._response_window_sec
@property
def reward_volume(self) -> float:
return self._reward_volume
@property
def auto_reward_volume(self) -> float:
return self._auto_reward_volume
@property
def session_type(self) -> str:
return self._session_type
@property
def stimulus(self) -> str:
return self._stimulus
@property
def stimulus_distribution(self) -> float:
return self._stimulus_distribution
@property
def task(self) -> TaskType:
return self._task
@property
def n_stimulus_frames(self) -> int:
return self._n_stimulus_frames
@property
def stimulus_name(self) -> Optional[str]:
return self._stimulus_name
@property
def image_set(self) -> str:
return self._image_set
[docs]
def to_nwb(self, nwbfile: NWBFile) -> NWBFile:
nwb_extension = load_pynwb_extension(
BehaviorTaskParametersSchema, 'ndx-aibs-behavior-ophys'
)
task_parameters = self.to_dict()['task_parameters']
task_parameters_clean = BehaviorTaskParametersSchema().dump(
task_parameters
)
new_task_parameters_dict = {}
for key, val in task_parameters_clean.items():
if isinstance(val, list):
new_task_parameters_dict[key] = np.array(val)
else:
new_task_parameters_dict[key] = val
nwb_task_parameters = nwb_extension(
name='task_parameters', **new_task_parameters_dict)
nwbfile.add_lab_meta_data(nwb_task_parameters)
return nwbfile
[docs]
@classmethod
def from_nwb(cls, nwbfile: NWBFile) -> "TaskParameters":
metadata_nwb_obj = nwbfile.lab_meta_data['task_parameters']
data = BehaviorTaskParametersSchema().dump(metadata_nwb_obj)
data['task_type'] = data['task']
del data['task']
return TaskParameters(**data)
[docs]
@classmethod
def from_stimulus_file(
cls,
stimulus_file: BehaviorStimulusFile) -> "TaskParameters":
data = stimulus_file.data
behavior = data["items"]["behavior"]
config = behavior["config"]
doc = config["DoC"]
blank_duration_sec = [float(x) for x in doc['blank_duration_range']]
stim_duration = cls._calculate_stimulus_duration(
stimulus_file=stimulus_file)
omitted_flash_fraction = \
behavior['params'].get('flash_omit_probability', float('nan'))
response_window_sec = [float(x) for x in doc["response_window"]]
reward_volume = config["reward"]["reward_volume"]
auto_reward_volume = doc['auto_reward_volume']
session_type = behavior["params"]["stage"]
stimulus = next(iter(behavior["stimuli"]))
stimulus_name = stimulus_file.stimulus_name
stimulus_distribution = doc["change_time_dist"]
task = cls._parse_task(stimulus_file=stimulus_file)
n_stimulus_frames = cls._calculuate_n_stimulus_frames(
stimulus_file=stimulus_file)
return TaskParameters(
blank_duration_sec=blank_duration_sec,
stimulus_duration_sec=stim_duration,
omitted_flash_fraction=omitted_flash_fraction,
response_window_sec=response_window_sec,
reward_volume=reward_volume,
auto_reward_volume=auto_reward_volume,
session_type=session_type,
stimulus=stimulus,
stimulus_distribution=stimulus_distribution,
task_type=task,
n_stimulus_frames=n_stimulus_frames,
stimulus_name=stimulus_name
)
@staticmethod
def _calculate_stimulus_duration(
stimulus_file: BehaviorStimulusFile) -> float:
data = stimulus_file.data
behavior = data["items"]["behavior"]
stimuli = behavior['stimuli']
def _parse_stimulus_key():
if 'images' in stimuli:
stim_key = 'images'
elif 'grating' in stimuli:
stim_key = 'grating'
else:
msg = "Cannot get stimulus_duration_sec\n"
msg += "'images' and/or 'grating' not a valid "
msg += "key in pickle file under "
msg += "['items']['behavior']['stimuli']\n"
msg += f"keys: {list(stimuli.keys())}"
raise RuntimeError(msg)
return stim_key
stim_key = _parse_stimulus_key()
stim_duration = stimuli[stim_key]['flash_interval_sec']
# from discussion in
# https://github.com/AllenInstitute/AllenSDK/issues/1572
#
# 'flash_interval' contains (stimulus_duration, gray_screen_duration)
# (as @matchings said above). That second value is redundant with
# 'blank_duration_range'. I'm not sure what would happen if they were
# set to be conflicting values in the params. But it looks like
# they're always consistent. It should always be (0.25, 0.5),
# except for TRAINING_0 and TRAINING_1, which have statically
# displayed stimuli (no flashes).
if stim_duration is None:
stim_duration = np.nan
else:
stim_duration = stim_duration[0]
return stim_duration
@staticmethod
def _parse_task(
stimulus_file: BehaviorStimulusFile) -> TaskType:
data = stimulus_file.data
config = data["items"]["behavior"]["config"]
task_id = config['behavior']['task_id']
if 'DoC' in task_id:
task = TaskType.CHANGE_DETECTION
else:
msg = "metadata.get_task_parameters does not "
msg += f"know how to parse 'task_id' = {task_id}"
raise RuntimeError(msg)
return task
@staticmethod
def _calculuate_n_stimulus_frames(
stimulus_file: BehaviorStimulusFile) -> int:
data = stimulus_file.data
behavior = data["items"]["behavior"]
n_stimulus_frames = 0
for stim_type, stim_table in behavior["stimuli"].items():
n_stimulus_frames += sum(stim_table.get("draw_log", []))
return n_stimulus_frames