# Here we will define a class for aligning the timesteps in a sync
# file with the frames listed in a stimulus pickle file.
from typing import Tuple, Union, List, Dict, Any
import numpy as np
import logging
import pathlib
from allensdk.brain_observatory import sync_dataset
from allensdk.internal.core.lims_utilities import safe_system_path
from allensdk.brain_observatory.behavior.data_files.stimulus_file import (
_StimulusFile)
def _choose_line(
data: sync_dataset.Dataset,
sync_lines: Union[str, Tuple[str]]) -> str:
"""
Scan through sync_lines in order. Select the first one
that is present in the sync file. Raise an exception if
none are present.
Parameters
----------
data: sync_dataset.Dataset
sync_lines: Union[str, Tuple[str]]
Returns
-------
chosen_line: str
The first line in sync_lines that is present in
the sync file.
"""
if isinstance(sync_lines, str):
sync_lines = (sync_lines, )
chosen_line = None
for this_line in sync_lines:
if this_line in data.line_labels:
chosen_line = this_line
break
if chosen_line is None:
msg = ("Could not find one of "
f"{sync_lines} in sync dataset. "
f"available lines:\n{data.line_labels}")
raise RuntimeError(msg)
return chosen_line
def _get_rising_times(
data: sync_dataset.Dataset,
sync_lines: Union[str, Tuple[str]]):
"""
Get the timestamps, in seconds, associated with the rising
edges in a specific line in a sync file
Parameters
----------
data: sync_dataset.Dataset
sync_lines: Union[str, Tuple[str]]
The line to look for in the sync file.
If a str, return the rising edges in that line.
If a Tuple, work through the tuple **in order** until
a line is found that is present in the sync file. That
is the line for which timestamps will be returned.
Returns
-------
timestamps: np.ndarray
The times, in seconds, associated with the rising edges
of the chosen line.
"""
chosen_line = _choose_line(
data=data,
sync_lines=sync_lines)
timestamps = data.get_rising_edges(
line=chosen_line,
units='seconds')
return timestamps
def _get_falling_times(
data: sync_dataset.Dataset,
sync_lines: Union[str, Tuple[str]]):
"""
Get the timestamps, in seconds, associated with the falling
edges in a specific line in a sync file.
Note: only falling edges that occur after rising edges are
returned. This is a quality control measure.
Parameters
----------
data: sync_dataset.Dataset
sync_lines: Union[str, Tuple[str]]
The line to look for in the sync file.
If a str, return the falling edges in that line.
If a Tuple, work through the tuple **in order** until
a line is found that is present in the sync file. That
is the line for which timestamps will be returned.
Returns
-------
timestamps: np.ndarray
The times, in seconds, associated with the rising edges
of the chosen line.
"""
chosen_line = _choose_line(
data=data,
sync_lines=sync_lines)
rising_edges = data.get_rising_edges(
line=chosen_line,
units='seconds')
falling_edges = data.get_falling_edges(
line=chosen_line,
units='seconds')
valid = (falling_edges > rising_edges[0])
return falling_edges[valid]
def _get_line_starts_and_ends(
data: sync_dataset.Dataset,
sync_lines: Union[str, Tuple[str]]) -> Tuple[np.ndarray, np.ndarray]:
"""
Parameters
----------
data: sync_dataset.Dataset
sync_lines: Union[str, Tuple[str]]
The line to look for in the sync file.
If a str, return the falling edges in that line.
If a Tuple, work through the tuple **in order** until
a line is found that is present in the sync file. That
is the line for which timestamps will be returned.
Returns
-------
(start_times, end_times): Tuple[np.ndarray, np.ndarray]
np.ndarrays of times (in seconds) that the given
line turns on (rises) and turns off (falls).
"""
start_times = _get_rising_times(
data=data,
sync_lines=sync_lines)
end_times = _get_falling_times(
data=data,
sync_lines=sync_lines)
return (start_times, end_times)
def _get_start_frames(
data: sync_dataset.Dataset,
raw_frame_times: np.ndarray,
stimulus_frame_counts: List[int],
tolerance: float) -> List[int]:
"""
Find the start frames for a series of stimuli that need to be
registered to a single sync file.
Parameters
----------
data: sync_dataset.Dataset
A representation of the sync file being registered
raw_frame_times: np.ndarray
The timestamps (in seconds) of the events being registered
stimulus_frame_counts: List[int]
The number of events occuring during each stimulus
as read from the pickle file, **in the order** that
the stimuli were run.
tolerance: float
The tolerance within which the length of an epoch
and an element of stimulus_frame_counts will be
considered the same.
Returns
-------
start_frames: List[int]
The global index of the starting frames associated with the
stimuli represented by stimulus_frame_counts
Notes
-----
Ideally, the vsync_stim line in the sync file represents when
frames from a stimulus are presented to the mouse. Unfortunately,
they do not carry any information about which stimulus block
the frames correspond to (behavior, mapping, or replay in a fiducial
VBN session). The stim_running line, however, can be used to find
the dividing times between the three stimulus presentations.
stim_running is high when a stimulus block is being presented
and low when it is not. What this method does is:
1) Take a list of raw frame times as input (probably the falling
edges of the vsync_stim line, but possibly the rising edges,
depending on our use case)
2) Find the timestamps associated with the rising and falling
edges of stim_running. These are taken to be the breaks between
stimulus blocks.
3) Divide the raw frame times into blocks that fall between the
stimulus block start/end times from (2).
4) Try to match the blocks of frame times against the expected
number of frames in each block as specified in stimulus_frame_counts.
See discussion here
http://confluence.corp.alleninstitute.org/pages/viewpage.action?spaceKey=IT&title=Addressing+variable+monitor+lag+in+sync+data+for+Visual+Behavior+Neuropixels
This is a modified copy of a method originally implemented in
ecephys_etl_pipelines/.../vbn_create_stimulus_table/create_stim_table.py
The purpose of the modified copy is to allow us to use rising *or* falling
edges as raw_stim_times, depending on the use case.
"""
frame_count_arr = np.array(stimulus_frame_counts)
stim_starts, stim_ends = _get_line_starts_and_ends(
data=data,
sync_lines=('stim_running', 'sweep'))
# break raw_frame_times into epochs based on stim_starts and stim_ends
epoch_frame_counts = []
epoch_start_frames = []
for start, end in zip(stim_starts, stim_ends):
# Inner expression returns a bool array where conditions are True
# np.where evaluates bool array to return indices where bool array True
epoch_frames = np.where((raw_frame_times >= start)
& (raw_frame_times < end))[0]
epoch_frame_counts.append(len(epoch_frames))
epoch_start_frames.append(epoch_frames[0])
if len(epoch_frame_counts) == len(frame_count_arr):
# There is a 1:1 mapping between the epochs found by sub-dividing
# the stim_running line and the stimulus blocks expected based on
# the stimulus pickle files.
if not np.allclose(frame_count_arr, epoch_frame_counts):
logging.warning(
f"Number of frames derived from sync file "
f"({epoch_frame_counts})for each epoch not matching up with "
f"frame counts derived from pkl files ({frame_count_arr})!"
)
start_frames = epoch_start_frames
elif len(epoch_frame_counts) > len(frame_count_arr):
# There were, for some reason, more epochs found by sub-dividing the
# stim_running line than there were expected based on the stimulus
# pickle files.
logging.warning(
f"Number of stim presentations obtained from sync "
f"({len(epoch_frame_counts)}) higher than number expected "
f"({len(frame_count_arr)}). Inferring start frames."
)
start_frames = []
for stim_idx, fc in enumerate(frame_count_arr):
logging.info(f"Finding stim start for stim with index: {stim_idx}")
# Get index of stimulus whose frame counts most closely match
# the expected number of frames
best_match = int(
np.argmin([np.abs(efc - fc) for efc in epoch_frame_counts])
)
lower_tol = fc * (1.0 - tolerance)
upper_tol = fc * (1.0 + tolerance)
if lower_tol <= epoch_frame_counts[best_match] <= upper_tol:
_ = epoch_frame_counts.pop(best_match)
start_frame = epoch_start_frames.pop(best_match)
start_frames.append(start_frame)
logging.info(
f"Found stim start for stim with index ({stim_idx})"
f"at vsync ({start_frame})"
)
else:
raise RuntimeError(
"Could not find matching sync frames "
f"for stim: {stim_idx}\n"
f"expected n_frames {fc}; "
f"best_match {epoch_frame_counts[best_match]}; "
f"tolerance {tolerance}"
)
else:
raise RuntimeError(
f"Do not know how to handle more pkl frame count entries "
f"({frame_count_arr}) than sync derived epoch frame count "
f"entries ({epoch_frame_counts})!"
)
return start_frames
[docs]
def get_stim_timestamps_from_stimulus_blocks(
stimulus_files: Union[_StimulusFile, List[_StimulusFile]],
sync_file: Union[str, pathlib.Path],
raw_frame_time_lines: Union[str, List[str]],
raw_frame_time_direction: str,
frame_count_tolerance: float) -> Dict[str, Any]:
"""
Find the timestamps associated a set of stimulus blocks
that have to be aligned with a single sync file
Parameters
----------
stimulus_files: Union[_StimulusFile, List[_StimulusFile]]
The _StimulusFile objects being registered to the sync file
sync_file: Union[str, pathlib.Path]
The path to the sync file
raw_frame_time_lines: Union[str, List[str]]
The line to be used to find raw frame times (usually 'vsync_stim').
If a list, the code will scan the list in order until a line
that is present in the sync file is found. That line will be used.
raw_frame_time_direction: str
Either 'rising' or 'falling' indicating which edge to use in finding
the raw frame times
frame_count_tolerance: float
The tolerance to within two blocks of frame counts are considered
equal
Returns
-------
A dict in which
"timestamps" -> List[np.ndarray]
The list of timestamp arrays corresponding to the provided
_StimulusFiles.
"start_frames" -> List[int]
The list of starting frames for the provided _StimulusFiles
**The order of stimulus_files will dictate the order of these
lists.**
Notes
-----
This method operates by finding the start frames associated
with each stimulus block according to _get_start_frames and
then assigning the timestamps associated with
stimulus_block.num_frames to each stimulus block.
"""
if raw_frame_time_direction == 'rising':
frame_time_fn = _get_rising_times
elif raw_frame_time_direction == 'falling':
frame_time_fn = _get_falling_times
else:
msg = ("Cannot parse raw_frame_time_direction = "
f"'{raw_frame_time_direction}'\n"
"must be either 'rising' or 'falling'")
raise ValueError(msg)
if not isinstance(stimulus_files, list):
stimulus_files = [stimulus_files, ]
if isinstance(sync_file, pathlib.Path):
str_path = str(sync_file.resolve().absolute())
else:
str_path = sync_file
safe_sync_path = safe_system_path(file_name=str_path)
list_of_timestamps = []
with sync_dataset.Dataset(safe_sync_path) as sync_data:
raw_frame_times = frame_time_fn(
data=sync_data,
sync_lines=raw_frame_time_lines)
frame_count_list = [s.num_frames for s in stimulus_files]
start_frames = _get_start_frames(
data=sync_data,
raw_frame_times=raw_frame_times,
stimulus_frame_counts=frame_count_list,
tolerance=frame_count_tolerance)
for f0, nf in zip(start_frames, frame_count_list):
this_array = raw_frame_times[f0:f0+nf]
list_of_timestamps.append(this_array)
return {"timestamps": list_of_timestamps,
"start_frames": start_frames}