patch_dff.py

import numpy as np

from allensdk.core.brain_observatory_nwb_data_set import BrainObservatoryNwbDataSet
from allensdk.core.brain_observatory_cache import BrainObservatoryCache
from allensdk.brain_observatory.drifting_gratings import DriftingGratings


class DffPatchingDataSet(BrainObservatoryNwbDataSet):

    def __init__(self, nwb_file, dff_traces, dff_timestamps=None):
        super(DffPatchingDataSet, self).__init__(nwb_file)
        self.dff_traces = dff_traces
        self.dff_timestamps = dff_timestamps

    def get_dff_traces(self, cell_specimen_ids=None):

        # This block reuses the existing timestamps, unless you've provided your own
        if self.dff_timestamps is not None:
            timestamps, traces = super(DffPatchingDataSet, self).get_dff_traces(cell_specimen_ids)
            del traces
        else:
            timestamps = self.dff_timestamps

        # BrainObservatoryNwbDataSet supports filtering to a subset of cells.
        # I don't think we need this for stimulus analysis, but supporting it is pretty easy
        cell_selection = slice(None) if cell_specimen_ids is None else self.get_cell_specimen_indices(cell_specimen_ids)
        
        return timestamps, self.dff_traces[cell_selection, :]
    
    
class DffPatchingCache(BrainObservatoryCache):
    
    def get_ophys_experiment_data(
        self, 
        ophys_experiment_id, 
        file_name=None, 
        dff_traces=None, 
        dff_timestamps=None
    ):
        
        base_data_set = super(DffPatchingCache, self).get_ophys_experiment_data(
            ophys_experiment_id, file_name=file_name
        )
        
        if dff_traces is None and dff_timestamps is None:
            return base_data_set
        elif dff_traces is not None:
            return DffPatchingDataSet(base_data_set.nwb_file, dff_traces, dff_timestamps)
        elif dff_timestamps is not None:
            raise ValueError("You've overriden dff timestamps without overriding the data!")
            
            
cache = DffPatchingCache()
eid = cache.get_ophys_experiments(stimuli=["drifting_gratings"])[0]["id"]

regular_data_set = cache.get_ophys_experiment_data(eid)

# make a random array of the right shape
timestamps, dff = regular_data_set.get_dff_traces()
new_dff = np.random.rand(*dff.shape)

# ... and patch it onto the dataset
patched_data_set = cache.get_ophys_experiment_data(eid, dff_traces=new_dff)

# compare speed_tuning (a running-window df/f-dependent analysis)
regular_dg_analysis = DriftingGratings(regular_data_set)
print(regular_dg_analysis.get_speed_tuning(regular_dg_analysis._binsize))

patched_dg_analysis = DriftingGratings(patched_data_set)
print(patched_dg_analysis.get_speed_tuning(patched_dg_analysis._binsize))