import abc
import copy
import datetime
import uuid
from pathlib import Path
from typing import Dict, Optional, Tuple, Union
from allensdk.core import DataObject
from allensdk.core.pickle_utils import load_and_sanitize_pickle
from allensdk.internal.api import PostgresQueryMixin
from allensdk.internal.core import DataFile
from allensdk.internal.core.lims_utilities import safe_system_path
from cachetools import LRUCache, cached
from cachetools.keys import hashkey
# Query returns path to StimulusPickle file for given behavior session
BEHAVIOR_STIMULUS_FILE_QUERY_TEMPLATE = """
SELECT
wkf.storage_directory || wkf.filename AS stim_file
FROM
well_known_files wkf
WHERE
wkf.attachable_id = {behavior_session_id}
AND wkf.attachable_type = 'BehaviorSession'
AND wkf.well_known_file_type_id IN (
SELECT id
FROM well_known_file_types
WHERE name = 'StimulusPickle');
"""
[docs]
def from_json_cache_key(cls, stimulus_file_path: str):
return hashkey(stimulus_file_path, cls.file_path_key())
[docs]
def from_lims_cache_key(cls, db, behavior_session_id: int):
return hashkey(behavior_session_id, cls.file_path_key())
class _StimulusFile(DataFile):
"""A DataFile which contains methods for accessing and loading visual
behavior stimulus *.pkl files.
This file type contains a number of parameters collected during a behavior
session including information about stimulus presentations, rewards,
trials, and timing for all of the above.
"""
@classmethod
def file_path_key(cls) -> str:
"""
The key in the dict_repr that maps to the path
to this StimulusFile's pickle file on disk.
"""
raise NotImplementedError()
def __init__(self, filepath: Union[str, Path]):
super().__init__(filepath=filepath)
@classmethod
def from_json(cls, dict_repr: dict) -> "_StimulusFile":
filepath = dict_repr[cls.file_path_key()]
return cls._from_json(stimulus_file_path=filepath)
@classmethod
@cached(cache=LRUCache(maxsize=10), key=from_json_cache_key)
def _from_json(cls, stimulus_file_path: str) -> "_StimulusFile":
return cls(filepath=stimulus_file_path)
@classmethod
@cached(cache=LRUCache(maxsize=10), key=from_lims_cache_key)
def from_lims(
cls, db: PostgresQueryMixin, behavior_session_id: Union[int, str]
) -> "_StimulusFile":
raise NotImplementedError()
@staticmethod
def load_data(filepath: Union[str, Path]) -> dict:
filepath = safe_system_path(file_name=filepath)
return load_and_sanitize_pickle(pickle_path=filepath)
@property
def num_frames(self) -> int:
"""
Return the number of frames associated with this StimulusFile
"""
return len(self.data["intervalsms"]) + 1
[docs]
class BehaviorStimulusFile(_StimulusFile):
def __init__(self, filepath: Union[str, Path]):
super().__init__(filepath=filepath)
[docs]
@classmethod
def file_path_key(cls) -> str:
return "behavior_stimulus_file"
[docs]
@classmethod
@cached(cache=LRUCache(maxsize=10), key=from_lims_cache_key)
def from_lims(
cls, db: PostgresQueryMixin, behavior_session_id: Union[int, str]
) -> "BehaviorStimulusFile":
query = BEHAVIOR_STIMULUS_FILE_QUERY_TEMPLATE.format(
behavior_session_id=behavior_session_id
)
filepath = db.fetchone(query, strict=True)
return cls(filepath=filepath)
def _validate_frame_data(self):
"""
Make sure that self.data['intervalsms'] does not exist and that
self.data['items']['behavior']['intervalsms'] does exist.
"""
msg = ""
if "intervalsms" in self.data:
msg += "self.data['intervalsms'] present; did not expect that\n"
if "items" not in self.data:
msg += "self.data['items'] not present\n"
else:
if "behavior" not in self.data["items"]:
msg += "self.data['items']['behavior'] not present\n"
else:
if "intervalsms" not in self.data["items"]["behavior"]:
msg += (
"self.data['items']['behavior']['intervalsms'] "
"not present\n"
)
if len(msg) > 0:
full_msg = f"When getting num_frames from {type(self)}\n"
full_msg += msg
full_msg += f"\nfilepath: {self.filepath}"
raise RuntimeError(full_msg)
return None
@property
def behavior_session_uuid(self) -> Optional[uuid.UUID]:
"""Return the behavior session UUID either from the uuid field
or foraging_id field.
"""
bs_uuid = self.data.get("session_uuid")
if bs_uuid is None:
try:
bs_uuid = self._retrieve_from_params("foraging_id")["value"]
except (KeyError, RuntimeError):
bs_uuid = None
if bs_uuid:
try:
bs_uuid = uuid.UUID(bs_uuid)
except ValueError:
bs_uuid = None
return bs_uuid
@property
def date_of_acquisition(self) -> datetime.datetime:
"""
Return the date_of_acquisition as a datetime.datetime.
This will be read from self.data['start_time']
"""
assert isinstance(self.data, dict)
if "start_time" not in self.data:
raise KeyError(
"No 'start_time' listed in pickle file " f"{self.filepath}"
)
return copy.deepcopy(self.data["start_time"])
@property
def mouse_id(self) -> str:
"""Retrieve the mouse_id value from the stimulus pickle file.
This can be read either from:
data['items']['behavior']['params']['stage']
or
data['items']['behavior']['cl_params']['stage']
if both are present and they disagree, raise an exception.
"""
return self._retrieve_from_params("mouse_id")
@property
def num_frames(self) -> int:
"""
Return the number of frames associated with this StimulusFile
"""
self._validate_frame_data()
return len(self.data["items"]["behavior"]["intervalsms"]) + 1
@property
def session_type(self) -> str:
"""
Return the session type as read from the pickle file. This can
be read either from
data['items']['behavior']['params']['stage']
or
data['items']['behavior']['cl_params']['stage']
if both are present and they disagree, raise an exception
"""
return self._retrieve_from_params("stage")
@property
def stimulus_name(self) -> str:
"""
Get the image stimulus name by parsing the file path of the image set.
If no image set, check for gratings and return "behavior" if not found.
Parameters
----------
stimulus_file : BehaviorStimulusFile
Stimulus pickle file to parse.
Returns
-------
stimulus_name : str
Name of the image stimulus from the image file path set shown to
the mouse.
"""
try:
stimulus_name = Path(
self.stimuli["images"]["image_set"]
).stem.split(".")[0]
except KeyError:
# if we can't find the images key in the stimuli, check for the
# name ``grating`` as the stimulus. If not add generic
# ``behavior``.
if "grating" in self.stimuli.keys():
stimulus_name = "grating"
else:
stimulus_name = "behavior"
return stimulus_name
def _retrieve_from_params(self, key_name: str):
"""Retrieve data from either data['items']['behavior']['params'] or
data['items']['behavior']['cl_params'].
Test for conflicts or missing data and raise if issues found.
Parameters
----------
key_name : str
Name of data to attempt to retrieve from the behavior stimulus
file data.
Returns
-------
value : various
Value from the stimulus file.
"""
param_value = None
if "params" in self.data["items"]["behavior"]:
if key_name in self.data["items"]["behavior"]["params"]:
param_value = self.data["items"]["behavior"]["params"][
key_name
]
cl_value = None
if "cl_params" in self.data["items"]["behavior"]:
if key_name in self.data["items"]["behavior"]["cl_params"]:
cl_value = self.data["items"]["behavior"]["cl_params"][
key_name
]
if cl_value is None and param_value is None:
raise RuntimeError(
f"Could not find {key_name} in pickle file " f"{self.filepath}"
)
if param_value is None:
return cl_value
if cl_value is None:
return param_value
if cl_value != param_value:
raise RuntimeError(
f"Conflicting {key_name} values in pickle file "
f"{self.filepath}\n"
f"cl_params: {cl_value}\n"
f"params: {param_value}\n"
)
return param_value
@property
def session_duration(self) -> float:
"""
Gets session duration in seconds
Returns
-------
session duration in seconds
"""
start_time = self.data["start_time"]
stop_time = self.data["stop_time"]
if not isinstance(start_time, datetime.datetime):
start_time = datetime.datetime.fromtimestamp(start_time)
if not isinstance(stop_time, datetime.datetime):
stop_time = datetime.datetime.fromtimestamp(stop_time)
delta = stop_time - start_time
return delta.total_seconds()
@property
def stimuli(self) -> Dict[str, Tuple[str, Union[str, int], int, int]]:
"""Stimuli shown during session
Returns
-------
stimuli:
(stimulus type ('Image' or 'Grating'),
stimulus descriptor (image_name or orientation of grating in
degrees),
nonsynced time of display,
display frame (frame that stimuli was displayed))
"""
# TODO implement return value as class (i.e. Image, Grating)
return self.data["items"]["behavior"]["stimuli"]
[docs]
def validate(self) -> "BehaviorStimulusFile":
if "items" not in self.data or "behavior" not in self.data["items"]:
raise MalformedStimulusFileError(
f'Expected to find key "behavior" in "items" dict. '
f'Found {self.data["items"].keys()}'
)
return self
[docs]
class ReplayStimulusFile(_StimulusFile):
[docs]
@classmethod
def file_path_key(cls) -> str:
return "replay_stimulus_file"
[docs]
class MappingStimulusFile(_StimulusFile):
[docs]
@classmethod
def file_path_key(cls) -> str:
return "mapping_stimulus_file"
[docs]
class StimulusFileReadableInterface(abc.ABC):
"""Marks a data object as readable from stimulus file"""
[docs]
@classmethod
@abc.abstractmethod
def from_stimulus_file(
cls, stimulus_file: BehaviorStimulusFile
) -> "DataObject":
"""Populate a DataObject from the stimulus file
Returns
-------
DataObject:
An instantiated DataObject which has `name` and `value` properties
"""
raise NotImplementedError()
[docs]
class StimulusFileLookup(object):
"""
A container class to carry around the StimulusFile(s) associated
with a BehaviorSession
"""
def __init__(self):
self._values = dict()
@property
def behavior_stimulus_file(self) -> BehaviorStimulusFile:
if "behavior" not in self._values:
raise ValueError(
"This StimulusFileLookup has no " "BehaviorStimulusFile"
)
return self._values["behavior"]
@behavior_stimulus_file.setter
def behavior_stimulus_file(self, value: BehaviorStimulusFile):
if not isinstance(value, BehaviorStimulusFile):
raise ValueError(
"Trying to set behavior_stimulus_file to "
f"value of type {type(value)}; type should "
"be BehaviorStimulusFile"
)
self._values["behavior"] = value
@property
def replay_stimulus_file(self) -> ReplayStimulusFile:
if "replay" not in self._values:
raise ValueError(
"This StimulusFileLookup has no " "ReplayStimulusFile"
)
return self._values["replay"]
@replay_stimulus_file.setter
def replay_stimulus_file(self, value: ReplayStimulusFile):
if not isinstance(value, ReplayStimulusFile):
raise ValueError(
"Trying to set replay_stimulus_file to "
f"value of type {type(value)}; type should "
"be ReplayStimulusFile"
)
self._values["replay"] = value
@property
def mapping_stimulus_file(self) -> MappingStimulusFile:
if "mapping" not in self._values:
raise ValueError(
"This StimulusFileLookup has no " "MappingStimulusFile"
)
return self._values["mapping"]
@mapping_stimulus_file.setter
def mapping_stimulus_file(self, value: MappingStimulusFile):
if not isinstance(value, MappingStimulusFile):
raise ValueError(
"Trying to set mapping_stimulus_file to "
f"value of type {type(value)}; type should "
"be MappingStimulusFile"
)
self._values["mapping"] = value
[docs]
def stimulus_lookup_from_json(dict_repr: dict) -> StimulusFileLookup:
"""
Load a lookup table of the stimulus files associated with a
BehaviorSession from the dict representation of that session's
session_data
"""
lookup_table = StimulusFileLookup()
if BehaviorStimulusFile.file_path_key() in dict_repr:
stim = BehaviorStimulusFile.from_json(dict_repr=dict_repr)
lookup_table.behavior_stimulus_file = stim
if MappingStimulusFile.file_path_key() in dict_repr:
stim = MappingStimulusFile.from_json(dict_repr=dict_repr)
lookup_table.mapping_stimulus_file = stim
if ReplayStimulusFile.file_path_key() in dict_repr:
stim = ReplayStimulusFile.from_json(dict_repr=dict_repr)
lookup_table.replay_stimulus_file = stim
return lookup_table