timeline example

timeline.py

import numpy as np
import xarray

from allensdk.core.brain_observatory_cache import BrainObservatoryCache
from allensdk.brain_observatory import stimulus_info as si

TIME = 'time'
SEC = 'second'
MS = 'millisecond'
IDX = 'index'
CELL_INDEX = 'cell_index'

temporal_conversion_dict = {MS:.001, SEC:1., IDX:None}

class Timeline(object):
    
    def __init__(self, timestamps=None, unit=None, label=None):
        
        assert not timestamps is None
        assert not unit is None

        if label is None:
            label = str(uuid.uuid4())[:6]


        timestamps = np.array(timestamps).astype(np.float)
        self._len = len(timestamps)
        self._coords = {(label, unit):timestamps}

    def __len__(self):
        return self._len

    def get_coord_list(self):
        return ['time:%s(%s)' % (unit, label) for label, unit in self._coords.keys()]

    def add_relative_coordinate(self, reference_label, new_label, offset, offset_unit=None, new_unit=None):
        
        if offset_unit is None:
            raise NotImplementedError

        # Break out early and do nothing if requested coord already is in:
        if new_unit is None:
            raise NotImplementedError
        else:

            # print (new_label, new_unit), self.coords.keys()
            if ((new_label, new_unit) in self._coords.keys()):
                return

        reference_unit_list = sorted([x[1]  for x in self._coords.keys() if x[0] == reference_label])
        if reference_unit_list == 0:
            raise Exception('Label not found')
        reference_unit = reference_unit_list[0]
        reference_timestamps = self._coords[(reference_label, reference_unit)]

        # Shift current timeline back to master:
        curr_unit_convert = float(temporal_conversion_dict[reference_unit])
        offset_unit_convert = float(temporal_conversion_dict[offset_unit])
        new_timestamps = reference_timestamps - offset*(offset_unit_convert/curr_unit_convert)

        # Scale current timeline consistent with master:
        master_unit_convert = float(temporal_conversion_dict[new_unit])
        new_timestamps /= (master_unit_convert/curr_unit_convert)

        self._coords[(new_label, new_unit)] = new_timestamps

    def get_xarray(self, dim='time'):

        coords ={}
        for coord_tuple, timestamps in self._coords.items():
            key = 'time:%s(%s)' % coord_tuple
            coords[key] = (dim, timestamps)

        return xarray.DataArray(np.arange(len(self), 
                                dtype=np.int),
                                dims=(dim,),
                                coords=coords)

    @property
    def coords(self):
        return self.get_xarray().coords


    @property
    def dims(self):
        return self.get_xarray().dims

# Settings:
oeid = 501474098

# Initializations:
boc = BrainObservatoryCache()
nwb_dataset = boc.get_ophys_experiment_data(oeid)

# Get Data:
stimulus_epoch_table = nwb_dataset.get_stimulus_epoch_table()
metadata = nwb_dataset.get_metadata()
timestamps, dff = nwb_dataset.get_dff_traces()

mstl = Timeline(timestamps=timestamps,
              unit=SEC, 
              label='example')
mstl.add_relative_coordinate('example', 'master', offset=.1, offset_unit=SEC, new_unit=MS)

dff_dataset = xarray.DataArray(dff, dims=(CELL_INDEX, TIME), coords=mstl.coords)
print dff_dataset