Source code for allensdk.brain_observatory.sync_stim_aligner

# 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}