import pickle
import warnings
from typing import Dict, List, Optional, Set, Tuple, Union
import numpy as np
import pandas as pd
from allensdk.brain_observatory.behavior.data_files import BehaviorStimulusFile
from allensdk.brain_observatory.behavior.data_objects.stimuli.stimulus_templates import ( # noqa: E501
StimulusTemplate,
StimulusTemplateFactory,
)
from allensdk.brain_observatory.behavior.data_objects.stimuli.util import (
convert_filepath_caseinsensitive,
get_image_set_name,
)
from allensdk.brain_observatory.ophys.project_constants import (
PROJECT_CODES,
VBO_ACTIVE_MAP,
VBO_PASSIVE_MAP,
)
from allensdk.core.dataframe_utils import INT_NULL
[docs]
def load_pickle(pstream):
return pickle.load(pstream, encoding="bytes")
[docs]
def get_stimulus_presentations(data, stimulus_timestamps) -> pd.DataFrame:
"""
This function retrieves the stimulus presentation dataframe and
renames the columns, adds a stop_time column, and set's index to
stimulus_presentation_id before sorting and returning the dataframe.
:param data: stimulus file associated with experiment id
:param stimulus_timestamps: timestamps indicating when stimuli switched
during experiment
:return: stimulus_table: dataframe containing the stimuli metadata as well
as what stimuli was presented
"""
stimulus_table = get_visual_stimuli_df(data, stimulus_timestamps)
# workaround to rename columns to harmonize with visual
# coding and rebase timestamps to sync time
stimulus_table.insert(
loc=0, column="flash_number", value=np.arange(0, len(stimulus_table))
)
stimulus_table = stimulus_table.rename(
columns={
"frame": "start_frame",
"time": "start_time",
"flash_number": "stimulus_presentations_id",
}
)
stimulus_table.start_time = [
stimulus_timestamps[int(start_frame)]
for start_frame in stimulus_table.start_frame.values
]
end_time = []
for end_frame in stimulus_table.end_frame.values:
if not np.isnan(end_frame):
end_time.append(stimulus_timestamps[int(end_frame)])
else:
end_time.append(float("nan"))
stimulus_table.insert(loc=4, column="stop_time", value=end_time)
stimulus_table.set_index("stimulus_presentations_id", inplace=True)
stimulus_table = stimulus_table[sorted(stimulus_table.columns)]
return stimulus_table
[docs]
def get_images_dict(pkl) -> Dict:
"""
Gets the dictionary of images that were presented during an experiment
along with image set metadata and the image specific metadata. This
function uses the path to the image pkl file to read the images and their
metadata from the pkl file and return this dictionary.
Parameters
----------
pkl: The pkl file containing the data for the stimuli presented during
experiment
Returns
-------
Dict:
A dictionary containing keys images, metadata, and image_attributes.
These correspond to paths to image arrays presented, metadata
on the whole set of images, and metadata on specific images,
respectively.
"""
# Sometimes the source is a zipped pickle:
pkl_stimuli = pkl["items"]["behavior"]["stimuli"]
metadata = {"image_set": pkl_stimuli["images"]["image_path"]}
# Get image file name;
# These are encoded case-insensitive in the pickle file :/
filename = convert_filepath_caseinsensitive(metadata["image_set"])
image_set = load_pickle(open(filename, "rb"))
images = []
images_meta = []
ii = 0
for cat, cat_images in image_set.items():
for img_name, img in cat_images.items():
meta = dict(
image_category=cat.decode("utf-8"),
image_name=img_name.decode("utf-8"),
orientation=np.nan,
phase=np.nan,
spatial_frequency=np.nan,
image_index=ii,
)
images.append(img)
images_meta.append(meta)
ii += 1
images_dict = dict(
metadata=metadata,
images=images,
image_attributes=images_meta,
)
return images_dict
[docs]
def get_stimulus_templates(
pkl: dict,
grating_images_dict: Optional[dict] = None,
limit_to_images: Optional[List] = None,
) -> Optional[StimulusTemplate]:
"""
Gets images presented during experiments from the behavior stimulus file
(*.pkl)
Parameters
----------
pkl : dict
Loaded pkl dict containing data for the presented stimuli.
grating_images_dict : Optional[dict]
Because behavior pkl files do not contain image versions of grating
stimuli, they must be obtained from an external source. The
grating_images_dict is a nested dictionary where top level keys
correspond to grating image names (e.g. 'gratings_0.0',
'gratings_270.0') as they would appear in table returned by
get_gratings_metadata(). Sub-nested dicts are expected to have 'warped'
and 'unwarped' keys where values are numpy image arrays
of aforementioned warped or unwarped grating stimuli.
limit_to_images: Optional[list]
Only return images given by these image names
Returns
-------
StimulusTemplate:
StimulusTemplate object containing images that were presented during
the experiment
"""
pkl_stimuli = pkl["items"]["behavior"]["stimuli"]
if "images" in pkl_stimuli:
images = get_images_dict(pkl)
image_set_filepath = images["metadata"]["image_set"]
image_set_name = get_image_set_name(image_set_path=image_set_filepath)
image_set_name = convert_filepath_caseinsensitive(image_set_name)
attrs = images["image_attributes"]
image_values = images["images"]
if limit_to_images is not None:
keep_idxs = [
i
for i in range(len(images))
if attrs[i]["image_name"] in limit_to_images
]
attrs = [attrs[i] for i in keep_idxs]
image_values = [image_values[i] for i in keep_idxs]
return StimulusTemplateFactory.from_unprocessed(
image_set_name=image_set_name,
image_attributes=attrs,
images=image_values,
)
elif "grating" in pkl_stimuli:
if (grating_images_dict is None) or (not grating_images_dict):
raise RuntimeError(
"The 'grating_images_dict' param MUST "
"be provided to get stimulus templates "
"because this pkl data contains "
"gratings presentations."
)
gratings_metadata = get_gratings_metadata(pkl_stimuli).to_dict(
orient="records"
)
unwarped_images = []
warped_images = []
for image_attrs in gratings_metadata:
image_name = image_attrs["image_name"]
grating_imgs_sub_dict = grating_images_dict[image_name]
unwarped_images.append(grating_imgs_sub_dict["unwarped"])
warped_images.append(grating_imgs_sub_dict["warped"])
return StimulusTemplateFactory.from_processed(
image_set_name="grating",
image_attributes=gratings_metadata,
unwarped=unwarped_images,
warped=warped_images,
)
else:
warnings.warn(
"Could not determine stimulus template images from pkl file. "
f"The pkl stimuli nested dict "
"(pkl['items']['behavior']['stimuli']) contained neither "
"'images' nor 'grating' but instead: "
f"'{pkl_stimuli.keys()}'"
)
return None
def _resolve_image_category(change_log, frame):
for change in (unpack_change_log(c) for c in change_log):
if frame < change["frame"]:
return change["from_category"]
return change["to_category"]
def _get_stimulus_epoch(
set_log: List[Tuple[str, Union[str, int], int, int]],
current_set_index: int,
start_frame: int,
n_frames: int,
) -> Tuple[int, int]:
"""
Gets the frame range for which a stimuli was presented and the transition
to the next stimuli was ongoing. Returns this in the form of a tuple.
Parameters
----------
set_log: List[Tuple[str, Union[str, int], int, int
The List of Tuples in the form of
(stimuli_type ('Image' or 'Grating'),
stimuli_descriptor (image_name or orientation of grating in degrees),
nonsynced_time_of_display (not sure, it's never used),
display_frame (frame that stimuli was displayed))
current_set_index: int
Index of stimuli set to calculate window
start_frame: int
frame where stimuli was set, set_log[current_set_index][3]
n_frames: int
number of frames for which stimuli were displayed
Returns
-------
Tuple[int, int]:
A tuple where index 0 is start frame of stimulus window and index 1 is
end frame of stimulus window
"""
try:
next_set_event = set_log[current_set_index + 1]
except IndexError: # assume this is the last set event
next_set_event = (
None,
None,
None,
n_frames,
)
return start_frame, next_set_event[3] # end frame isn't inclusive
def _get_draw_epochs(
draw_log: List[int], start_frame: int, stop_frame: int
) -> List[Tuple[int, int]]:
"""
Gets the frame numbers of the active frames within a stimulus window.
Stimulus epochs come in the form [0, 0, 1, 1, 0, 0] where the stimulus is
active for some amount of time in the window indicated by int 1 at that
frame. This function returns the ranges for which the set_log is 1 within
the draw_log window.
Parameters
----------
draw_log: List[int]
A list of ints indicating for what frames stimuli were active
start_frame: int
The start frame to search within the draw_log for active values
stop_frame: int
The end frame to search within the draw_log for active values
Returns
-------
List[Tuple[int, int]]
A list of tuples indicating the start and end frames of every
contiguous set of active values within the specified window
of the draw log.
"""
draw_epochs = []
current_frame = start_frame
while current_frame <= stop_frame:
epoch_length = 0
while current_frame < stop_frame and draw_log[current_frame] == 1:
epoch_length += 1
current_frame += 1
else:
current_frame += 1
if epoch_length:
draw_epochs.append(
(
current_frame - epoch_length - 1,
current_frame - 1,
)
)
return draw_epochs
[docs]
def unpack_change_log(change):
(
(from_category, from_name),
(
to_category,
to_name,
),
time,
frame,
) = change
return dict(
frame=frame,
time=time,
from_category=from_category,
to_category=to_category,
from_name=from_name,
to_name=to_name,
)
[docs]
def get_visual_stimuli_df(data, time) -> pd.DataFrame:
"""
This function loads the stimuli and the omitted stimuli into a dataframe.
These stimuli are loaded from the input data, where the set_log and
draw_log contained within are used to calculate the epochs. These epochs
are used as start_frame and end_frame and converted to times by input
stimulus timestamps. The omitted stimuli do not have a end_frame by design
though there duration is always 250ms.
:param data: the behavior data file
:param time: the stimulus timestamps indicating when each stimuli is
displayed
:return: df: a pandas dataframe containing the stimuli and omitted stimuli
that were displayed with their frame, end_frame, start_time,
and duration
"""
stimuli = data["items"]["behavior"]["stimuli"]
n_frames = len(time)
visual_stimuli_data = []
for stim_dict in stimuli.values():
for idx, (attr_name, attr_value, _, frame) in enumerate(
stim_dict["set_log"]
):
orientation = attr_value if attr_name.lower() == "ori" else np.nan
image_name = attr_value if attr_name.lower() == "image" else np.nan
stimulus_epoch = _get_stimulus_epoch(
stim_dict["set_log"],
idx,
frame,
n_frames,
)
draw_epochs = _get_draw_epochs(
stim_dict["draw_log"], *stimulus_epoch
)
for epoch_start, epoch_end in draw_epochs:
visual_stimuli_data.append(
{
"orientation": orientation,
"image_name": image_name,
"frame": epoch_start,
"end_frame": epoch_end,
"time": time[epoch_start],
"duration": time[epoch_end] - time[epoch_start],
# this will always work because an epoch
# will never occur near the end of time
"omitted": False,
}
)
visual_stimuli_df = pd.DataFrame(data=visual_stimuli_data)
# Add omitted flash info:
try:
omitted_flash_frame_log = data["items"]["behavior"][
"omitted_flash_frame_log"
]
except KeyError:
# For sessions for which there were no omitted flashes
omitted_flash_frame_log = dict()
omitted_flash_list = []
for _, omitted_flash_frames in omitted_flash_frame_log.items():
stim_frames = visual_stimuli_df["frame"].values
omitted_flash_frames = np.array(omitted_flash_frames)
# Test offsets of omitted flash frames
# to see if they are in the stim log
offsets = np.arange(-3, 4)
offset_arr = np.add(
np.repeat(
omitted_flash_frames[:, np.newaxis], offsets.shape[0], axis=1
),
offsets,
)
matched_any_offset = np.any(np.isin(offset_arr, stim_frames), axis=1)
# Remove omitted flashes that also exist in the stimulus log
was_true_omitted = np.logical_not(matched_any_offset) # bool
omitted_flash_frames_to_keep = omitted_flash_frames[was_true_omitted]
# Have to remove frames that are double-counted in omitted log
omitted_flash_list += list(np.unique(omitted_flash_frames_to_keep))
omitted = np.ones_like(omitted_flash_list).astype(bool)
time = [time[fi] for fi in omitted_flash_list]
omitted_df = pd.DataFrame(
{
"omitted": omitted,
"frame": omitted_flash_list,
"time": time,
"image_name": "omitted",
}
)
df = (
pd.concat((visual_stimuli_df, omitted_df), sort=False)
.sort_values("frame")
.reset_index()
)
return df
[docs]
def get_image_names(behavior_stimulus_file: BehaviorStimulusFile) -> Set[str]:
"""Gets set of image names shown during behavior session"""
stimuli = behavior_stimulus_file.stimuli
image_names = set()
for stim_dict in stimuli.values():
for attr_name, attr_value, _, _ in stim_dict["set_log"]:
if attr_name.lower() == "image":
image_names.add(attr_value)
return image_names
[docs]
def is_change_event(stimulus_presentations: pd.DataFrame) -> pd.Series:
"""
Returns whether a stimulus is a change stimulus
A change stimulus is defined as the first presentation of a new image_name
Omitted stimuli are ignored
The first stimulus in the session is ignored
:param stimulus_presentations
The stimulus presentations table
:return: is_change: pd.Series indicating whether a given stimulus is a
change stimulus
"""
stimuli = stimulus_presentations["image_name"]
# exclude omitted stimuli
stimuli = stimuli[~stimulus_presentations["omitted"]]
prev_stimuli = stimuli.shift()
# exclude first stimulus
stimuli = stimuli.iloc[1:]
prev_stimuli = prev_stimuli.iloc[1:]
is_change = stimuli != prev_stimuli
# reset back to original index
is_change = is_change.reindex(stimulus_presentations.index).rename(
"is_change"
)
# Excluded stimuli are not change events
is_change = is_change.fillna(False)
return is_change
[docs]
def get_flashes_since_change(
stimulus_presentations: pd.DataFrame,
) -> pd.Series:
"""Calculate the number of times an images is flashed between changes.
Parameters
----------
stimulus_presentations : pandas.DataFrame
Table of presented stimuli with ``is_change`` column already
calculated.
Returns
-------
flashes_since_change : pandas.Series
Number of times the same image is flashed between image changes.
"""
flashes_since_change = pd.Series(
data=np.zeros(len(stimulus_presentations), dtype=float),
index=stimulus_presentations.index,
name="flashes_since_change",
dtype="int",
)
for idx, (pd_index, row) in enumerate(stimulus_presentations.iterrows()):
omitted = row["omitted"]
if pd.isna(row["omitted"]):
omitted = False
if row["image_name"] == "omitted" or omitted:
flashes_since_change.iloc[idx] = flashes_since_change.iloc[idx - 1]
else:
if row["is_change"] or idx == 0:
flashes_since_change.iloc[idx] = 0
else:
flashes_since_change.iloc[idx] = (
flashes_since_change.iloc[idx - 1] + 1
)
return flashes_since_change
[docs]
def add_active_flag(
stim_pres_table: pd.DataFrame, trials: pd.DataFrame
) -> pd.DataFrame:
"""Mark the active stimuli by lining up the stimulus times with the
trials times.
Parameters
----------
stim_pres_table : pandas.DataFrame
Stimulus table to add active column to.
trials : pandas.DataFrame
Trials table to align with the stimulus table.
Returns
-------
stimulus_table : pandas.DataFrame
Copy of ``stim_pres_table`` with added acive column.
"""
if "active" in stim_pres_table.columns:
return stim_pres_table
else:
active = pd.Series(
data=np.zeros(len(stim_pres_table), dtype=bool),
index=stim_pres_table.index,
name="active",
)
stim_mask = (
(stim_pres_table.start_time > trials.start_time.min())
& (stim_pres_table.start_time < trials.stop_time.max())
& (~stim_pres_table.image_name.isna())
)
active[stim_mask] = True
# Clean up potential stimuli that fall outside in time of the trials
# but are part of the "active" stimulus block.
if "stimulus_block" in stim_pres_table.columns:
for stim_block in stim_pres_table["stimulus_block"].unique():
block_mask = stim_pres_table["stimulus_block"] == stim_block
if np.any(active[block_mask]):
active[block_mask] = True
stim_pres_table["active"] = active
return stim_pres_table
[docs]
def compute_trials_id_for_stimulus(
stim_pres_table: pd.DataFrame, trials_table: pd.DataFrame
) -> pd.Series:
"""Add an id to allow for merging of the stimulus presentations
table with the trials table.
If stimulus_block is not available as a column in the input table, return
an empty set of trials_ids.
Parameters
----------
stim_pres_table : pandas.DataFrame
Pandas stimulus table to create trials_id from.
trials_table : pandas.DataFrame
Trials table to create id from using trial start times.
Returns
-------
trials_ids : pd.Series
Unique id to allow merging of the stim table with the trials table.
Null values are represented by -1.
Note
----
``trials_id`` values are copied from active stimulus blocks into
passive stimulus/replay blocks that contain the same image ordering and
length.
"""
# Create a placeholder for the trials_id.
trials_ids = pd.Series(
data=np.full(len(stim_pres_table), INT_NULL, dtype=int),
index=stim_pres_table.index,
name="trials_id",
).astype("int")
# Find stimulus blocks that start within a trial. Copy the trial_id
# into our new trials_ids series. For some sessions there are gaps in
# between one trial's end and the next's stop time so we account for this
# by only using the max time for all trials as the limit.
max_trials_stop = trials_table.stop_time.max()
for idx, trial in trials_table.iterrows():
stim_mask = (
(stim_pres_table.start_time > trial.start_time)
& (stim_pres_table.start_time < max_trials_stop)
& (~stim_pres_table.image_name.isna())
)
trials_ids[stim_mask] = idx
# Return input frame if the stimulus_block or active is not available.
if (
"stimulus_block" not in stim_pres_table.columns
or "active" not in stim_pres_table.columns
):
return trials_ids
active_sorted = stim_pres_table.active
# The code below finds all stimulus blocks that contain images/trials
# and attempts to detect blocks that are identical to copy the associated
# trials_ids into those blocks. In the parlance of the data this is
# copying the active stimulus block data into the passive stimulus block.
# Get the block ids for the behavior trial presentations
stim_blocks = stim_pres_table.stimulus_block
stim_image_names = stim_pres_table.image_name
active_stim_blocks = stim_blocks[active_sorted].unique()
# Find passive blocks that show images for potential copying of the active
# into a passive stimulus block.
passive_stim_blocks = stim_blocks[
np.logical_and(~active_sorted, ~stim_image_names.isna())
].unique()
# Copy the trials_id into the passive block if it exists.
if len(passive_stim_blocks) > 0:
for active_stim_block in active_stim_blocks:
active_block_mask = stim_blocks == active_stim_block
active_images = stim_image_names[active_block_mask].values
for passive_stim_block in passive_stim_blocks:
passive_block_mask = stim_blocks == passive_stim_block
if np.array_equal(
active_images, stim_image_names[passive_block_mask].values
):
trials_ids.loc[passive_block_mask] = trials_ids[
active_block_mask
].values
return trials_ids.sort_index()
[docs]
def fix_omitted_end_frame(stim_pres_table: pd.DataFrame) -> pd.DataFrame:
"""Fill NaN ``end_frame`` values for omitted frames.
Additionally, change type of ``end_frame`` to int.
Parameters
----------
stim_pres_table : `pandas.DataFrame`
Input stimulus table to fix/fill omitted ``end_frame`` values.
Returns
-------
output : `pandas.DataFrame`
Copy of input DataFrame with filled omitted, ``end_frame`` values and
fixed typing.
"""
median_stim_frame_duration = np.nanmedian(
stim_pres_table["end_frame"] - stim_pres_table["start_frame"]
)
omitted_end_frames = (
stim_pres_table[stim_pres_table["omitted"]]["start_frame"]
+ median_stim_frame_duration
)
stim_pres_table.loc[
stim_pres_table["omitted"], "end_frame"
] = omitted_end_frames
stim_dtypes = stim_pres_table.dtypes.to_dict()
stim_dtypes["start_frame"] = int
stim_dtypes["end_frame"] = int
return stim_pres_table.astype(stim_dtypes)
[docs]
def produce_stimulus_block_names(
stim_df: pd.DataFrame, session_type: str, project_code: str
) -> pd.DataFrame:
"""Add a column stimulus_block_name to explicitly reference the kind
of stimulus block in addition to the numbered blocks.
Only implemented currently for the VBO dataset. Will not add the column
if it is not in the defined set of project codes.
Parameters
----------
stim_df : pandas.DataFrame
Input stimulus presentations DataFrame with stimulus_block column
session_type : str
Full type name of session.
project_code : str
Full name of the project this session belongs to. As this function
is currently only written for VBO, if a non-VBO project name is
presented, the function will result in a noop.
Returns
-------
modified_df : pandas.DataFrame
Stimulus presentations DataFrame with added stimulus_block_name
column if the session is from a project that makes up the VBO release.
The data frame is return the same as the input if not.
"""
if project_code not in PROJECT_CODES:
return stim_df
vbo_map = VBO_PASSIVE_MAP if "passive" in session_type else VBO_ACTIVE_MAP
for stim_block in stim_df.stimulus_block.unique():
# If we have a single block then this is a training session and we
# add +1 to the block number to reuse the general VBO map and get the
# correct task.
block_id = stim_block
if len(stim_df.stimulus_block.unique()) == 1:
block_id += 1
stim_df.loc[
stim_df["stimulus_block"] == stim_block, "stimulus_block_name"
] = vbo_map[block_id]
return stim_df
[docs]
def compute_is_sham_change(
stim_df: pd.DataFrame, trials: pd.DataFrame
) -> pd.DataFrame:
"""Add is_sham_change to stimulus presentation table.
Parameters
----------
stim_df : pandas.DataFrame
Stimulus presentations table to add is_sham_change to.
trials : pandas.DataFrame
Trials data frame to pull info from to create
Returns
-------
stimulus_presentations : pandas.DataFrame
Input ``stim_df`` DataFrame with the is_sham_change column added.
"""
if (
"trials_id" not in stim_df.columns
or "active" not in stim_df.columns
or "stimulus_block" not in stim_df.columns
):
return stim_df
stim_trials = stim_df.merge(
trials, left_on="trials_id", right_index=True, how="left"
)
catch_frames = stim_trials[stim_trials["catch"].fillna(False)][
"change_frame"
].unique()
stim_df["is_sham_change"] = False
catch_flashes = stim_df[
stim_df["start_frame"].isin(catch_frames)
].index.values
stim_df.loc[catch_flashes, "is_sham_change"] = True
stim_blocks = stim_df.stimulus_block
stim_image_names = stim_df.image_name
active_stim_blocks = stim_blocks[stim_df.active].unique()
# Find passive blocks that show images for potential copying of the active
# into a passive stimulus block.
passive_stim_blocks = stim_blocks[
np.logical_and(~stim_df.active, ~stim_image_names.isna())
].unique()
# Copy the trials_id into the passive block if it exists.
if len(passive_stim_blocks) > 0:
for active_stim_block in active_stim_blocks:
active_block_mask = stim_blocks == active_stim_block
active_images = stim_image_names[active_block_mask].values
for passive_stim_block in passive_stim_blocks:
passive_block_mask = stim_blocks == passive_stim_block
if np.array_equal(
active_images, stim_image_names[passive_block_mask].values
):
stim_df.loc[
passive_block_mask, "is_sham_change"
] = stim_df[active_block_mask]["is_sham_change"].values
return stim_df.sort_index()