Merge pull request #1376 from AllenInstitute/refactor-visbeh-sync

Refactor visbeh sync

allensdk/brain_observatory/behavior/sync/__init__.py

@@ -4,64 +4,236 @@
 @author: marinag
 """
 from itertools import chain
-from typing import Dict, Any
-from .process_sync import filter_digital, calculate_delay  # NOQA: E402
+from typing import Dict, Any, Optional, List, Union
+from allensdk.brain_observatory.behavior.sync.process_sync import (
+    filter_digital, calculate_delay)  # NOQA: E402
 from allensdk.brain_observatory.sync_dataset import Dataset as SyncDataset  # NOQA: E402
 import numpy as np
 import scipy.stats as sps
 
-def get_sync_data(sync_path):
+
+def get_raw_stimulus_frames(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> np.ndarray:
+    """ Report the raw timestamps of each stimulus frame. This corresponds to 
+    the time at which the psychopy window's flip method returned, but not 
+    necessarily to the time at which the stimulus frame was displayed.
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start).
+
+    """
+    try:
+        return dataset.get_edges("falling",'stim_vsync', "seconds")
+    except KeyError:
+        if not permissive:
+            raise
+        return
+
+
+def get_ophys_frames(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> np.ndarray:
+    """ Report the timestamps of each optical physiology video frame
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start).
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Notes
+    -----
+    use rising edge for Scientifica, falling edge for Nikon 
+    http://confluence.corp.alleninstitute.org/display/IT/Ophys+Time+Sync
+    This function uses rising edges
+
+    """
+    try:
+        return dataset.get_edges("rising", '2p_vsync', "seconds")
+    except KeyError:
+        if not permissive:
+            raise
+        return
+
+
+def get_lick_times(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> Optional[np.ndarray]:
+    """ Report the timestamps of each detected lick
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start) 
+        or None. If None, no lick timestamps were found in this sync 
+        dataset.
+
+    """
+    return dataset.get_edges(
+        "rising", ["lick_times", "lick_sensor"], "seconds", permissive)
 
-    sync_dataset = SyncDataset(sync_path)
-    meta_data = sync_dataset.meta_data
-    sample_freq = meta_data['ni_daq']['counter_output_freq']
+
+def get_stim_photodiode(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> Optional[List[float]]:
+    """ Report the timestamps of each detected sync square transition (both 
+    black -> white and white -> black) in this experiment.
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start) 
+        or None. If None, no photodiode timestamps were found in this sync 
+        dataset.
+
+    """
+    return dataset.get_edges(
+        "all", ["stim_photodiode", "photodiode"], "seconds", permissive)
+
+
+def get_trigger(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> Optional[np.ndarray]:
+    """ Returns (as a 1-element array) the time at which optical physiology 
+    acquisition was started.
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+   Returns
+    -------
+    timestamps (floating point; seconds; relative to experiment start) 
+        or None. If None, no timestamps were found in this sync dataset.
+
+    Notes
+    -----
+    Ophys frame timestamps can be recorded before acquisition start when 
+        experimenters are setting up the recording session. These do not 
+        correspond to acquired ophys frames.
+
+    """
+    return dataset.get_edges(
+        "rising", ["2p_trigger", "acq_trigger"], "seconds", permissive)
+
+
+def get_eye_tracking(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> Optional[np.ndarray]:
+    """ Report the timestamps of each frame of the eye tracking video
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start) 
+        or None. If None, no eye tracking timestamps were found in this sync 
+        dataset.
+
+    """
+    return dataset.get_edges(
+        "rising", ["cam2_exposure", "eye_tracking"], "seconds", permissive)
+
+
+def get_behavior_monitoring(
+    dataset: SyncDataset, 
+    permissive: bool = False
+) -> Optional[np.ndarray]:
+    """ Report the timestamps of each frame of the behavior 
+    monitoring video
+
+    Parameters
+    ----------
+    dataset : describes experiment timing
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
+
+    Returns
+    -------
+    array of timestamps (floating point; seconds; relative to experiment start) 
+        or None. If None, no behavior monitoring timestamps were found in this 
+        sync dataset.
+
+    """
+    return dataset.get_edges(
+        "rising", ["cam1_exposure", "behavior_monitoring"], "seconds", 
+        permissive)
+
+
+def get_sync_data(
+    sync_path: str,
+    permissive: bool = False
+) -> Dict[str, Union[List, np.ndarray, None]]:
+    """ Convenience function for extracting several timestamp arrays from a 
+    sync file.
+
+    Parameters
+    ----------
+    sync_path : The hdf5 file here ought to be a Visual Behavior sync output 
+        file. See allensdk.brain_observatory.sync_dataset for more details of 
+        this format.
+    permissive : If True, None will be returned if timestamps are not found. If 
+        False, a KeyError will be raised
     
-    # use rising edge for Scientifica, falling edge for Nikon http://confluence.corp.alleninstitute.org/display/IT/Ophys+Time+Sync
-    # 2P vsyncs
-    vs2p_r = sync_dataset.get_rising_edges('2p_vsync')
-    vs2p_f = sync_dataset.get_falling_edges('2p_vsync')  # new sync may be able to do units = 'sec', so conversion can be skipped
-    frames_2p = vs2p_r / sample_freq
-    vs2p_fsec = vs2p_f / sample_freq
-
-    stimulus_times_no_monitor_delay = sync_dataset.get_falling_edges('stim_vsync') / sample_freq
-
-    if 'lick_times' in meta_data['line_labels']:
-        lick_times = sync_dataset.get_rising_edges('lick_1') / sample_freq
-    elif 'lick_sensor' in meta_data['line_labels']:
-        lick_times = sync_dataset.get_rising_edges('lick_sensor') / sample_freq
-    else:
-        lick_times = None
-    if '2p_trigger' in meta_data['line_labels']:
-        trigger = sync_dataset.get_rising_edges('2p_trigger') / sample_freq
-    elif 'acq_trigger' in meta_data['line_labels']:
-        trigger = sync_dataset.get_rising_edges('acq_trigger') / sample_freq
-    if 'stim_photodiode' in meta_data['line_labels']:
-        a = sync_dataset.get_rising_edges('stim_photodiode') / sample_freq
-        b = sync_dataset.get_falling_edges('stim_photodiode') / sample_freq
-        stim_photodiode = sorted(list(a)+list(b))
-    elif 'photodiode' in meta_data['line_labels']:
-        a = sync_dataset.get_rising_edges('photodiode') / sample_freq
-        b = sync_dataset.get_falling_edges('photodiode') / sample_freq
-        stim_photodiode = sorted(list(a)+list(b))
-    if 'cam2_exposure' in meta_data['line_labels']:
-        eye_tracking = sync_dataset.get_rising_edges('cam2_exposure') / sample_freq
-    elif 'eye_tracking' in meta_data['line_labels']:
-        eye_tracking = sync_dataset.get_rising_edges('eye_tracking') / sample_freq
-    if 'cam1_exposure' in meta_data['line_labels']:
-        behavior_monitoring = sync_dataset.get_rising_edges('cam1_exposure') / sample_freq
-    elif 'behavior_monitoring' in meta_data['line_labels']:
-        behavior_monitoring = sync_dataset.get_rising_edges('behavior_monitoring') / sample_freq
-
-    sync_data = {'ophys_frames': frames_2p,
-                 'lick_times': lick_times,
-                 'ophys_trigger': trigger,
-                 'eye_tracking': eye_tracking,
-                 'behavior_monitoring': behavior_monitoring,
-                 'stim_photodiode': stim_photodiode,
-                 'stimulus_times_no_delay': stimulus_times_no_monitor_delay,
-                 }
-
-    return sync_data
+    Returns
+    -------
+    A dictionary with the following keys. All timestamps in seconds:
+        ophys_frames : timestamps of each optical physiology frame
+        lick_times : timestamps of each detected lick
+        ophys_trigger : The time at which ophys acquisition was started
+        eye_tracking : timestamps of each eye tracking video frame
+        behavior_monitoring : timestamps of behavior monitoring video frame
+        stim_photodiode : timestamps of each photodiode transition
+        stimulus_times_no_delay : raw stimulus frame timestamps
+    Some values may be None. This indicates that the corresponding timestamps 
+    were not located in this sync file.
+
+    """
+
+    sync_dataset = SyncDataset(sync_path)
+    return {
+        'ophys_frames': get_ophys_frames(sync_dataset, permissive),
+        'lick_times': get_lick_times(sync_dataset, permissive),
+        'ophys_trigger': get_trigger(sync_dataset, permissive),
+        'eye_tracking': get_eye_tracking(sync_dataset, permissive),
+        'behavior_monitoring': get_behavior_monitoring(sync_dataset, permissive),
+        'stim_photodiode': get_stim_photodiode(sync_dataset, permissive),
+        'stimulus_times_no_delay': get_raw_stimulus_frames(sync_dataset, permissive)
+    }
 
 
 def frame_time_offset(data: Dict[str, Any]) -> float:

allensdk/brain_observatory/sync_dataset.py

@@ -15,6 +15,7 @@
 
 """
 import collections
+from typing import Union, Sequence, Optional
 
 import h5py as h5
 import numpy as np
@@ -294,8 +295,37 @@ def get_rising_edges(self, line, units='samples'):
         changes = self.get_bit_changes(bit)
         return self.get_all_times(units)[np.where(changes == 1)]
 
-    def get_edges(self, kind, keys, units='seconds'):
+    def get_edges(
+        self, 
+        kind: str, 
+        keys: Union[str, Sequence[str]], 
+        units: str = "seconds", 
+        permissive: bool = False
+    ) -> Optional[np.ndarray]:
         """ Utility function for extracting edge times from a line
+
+        Parameters
+        ----------
+        kind : One of "rising", "falling", or "all". Should this method return 
+            timestamps for rising, falling or both edges on the appropriate 
+            line
+        keys : These will be checked in sequence. Timestamps will be returned 
+            for the first which is present in the line labels
+        units : one of "seconds", "samples", or "indices". The returned 
+            "time"stamps will be given in these units.
+        raise_missing : If True and no matching line is found, a KeyError will
+            be raised
+
+        Returns
+        -------
+        An array of edge times. If raise_missing is False and none of the keys 
+            were found, returns None.
+
+        Raises
+        ------
+        KeyError : none of the provided keys were found among this dataset's 
+            line labels
+
         """
         if kind == 'falling':
             fn = self.get_falling_edges
@@ -316,7 +346,9 @@ def get_edges(self, kind, keys, units='seconds'):
             except ValueError:
                 continue
 
-        raise KeyError(f"none of {keys} were found in this dataset's line labels")
+        if not permissive:
+            raise KeyError(
+                f"none of {keys} were found in this dataset's line labels")
 
     def get_falling_edges(self, line, units='samples'):
         """