soltesz_extensions.py

import numpy as np
import random

from pynwb.spec import NWBDatasetSpec, NWBNamespaceBuilder, NWBGroupSpec, NWBDtypeSpec
from pynwb import get_class, load_namespaces, NWBFile, NWBHDF5IO

from datetime import datetime

ns_path = "soltesz.namespace.yaml"
ext_source = "soltesz.extensions.yaml"

gid_spec = NWBDatasetSpec(doc='global id for neuron',
                          shape=(None,1),
                          name='gid', dtype='int')
PopulationSpikeTimes = NWBGroupSpec(neurodata_type_def='PopulationSpikeTimes',
                                    doc='Population Spike Times',
                                    datasets=[gid_spec,
                                        NWBDatasetSpec(doc='spike times for all neurons', 
                                                       shape=(None,1), 
                                                       name='times', dtype='float'),
                                        NWBDatasetSpec(doc='data pointer', 
                                                       shape=(None,1),
                                                       name='data_pointer', dtype='int')],
                           neurodata_type_inc='NWBContainer')


CatCellInfo = NWBGroupSpec(neurodata_type_def='CatCellInfo',
                           doc='Categorical Cell Info',
                           datasets=[gid_spec,
                                     NWBDatasetSpec(name='indices',
                                                    doc='indices into values for each gid in order',
                                                    shape=(None,1),
                                                    dtype='int'),
                                     NWBDatasetSpec(name='values',
                                                    doc='list of unique values',
                                                    shope=(None,1), dtype='str')],
                           neurodata_type_inc='NWBContainer')


ContCellInfo = NWBGroupSpec(neurodata_type_def='ContCellInfo', 
                            datasets=[gid_spec,
                                NWBDatasetSpec(name='data',
                                               doc='continuous values',
                                               shape=(None, None), dtype='float')],
                           doc='Continuous Cell Info',
                           neurodata_type_inc='NWBContainer')


ns_builder = NWBNamespaceBuilder('soltesz extensions', "soltesz")
ns_builder.add_spec(ext_source, PopulationSpikeTimes)
ns_builder.add_spec(ext_source, CatCellInfo)
ns_builder.add_spec(ext_source, ContCellInfo)
ns_builder.export(ns_path)

############
# generate fake data
duration = 10
nspikes = 1000
nunits = 5

times = np.random.rand(nspikes) * duration
data_pointer = np.hstack([0,np.sort(random.sample(range(len(times)), nunits - 1))])
gids = np.arange(nunits).astype('int')
cell_types_vals, cell_types_inds  = np.unique(['ME']*3+['LE']*2, return_inverse=True)
x_pos = np.random.randn(nunits)

## pull types from yaml files
load_namespaces(ns_path)
PopulationSpikeTimes = get_class('PopulationSpikeTimes', 'soltesz')
CatCellInfo = get_class('CatCellInfo', 'soltesz')
ContCellInfo = get_class('ContCellInfo', 'soltesz')

pop_data = PopulationSpikeTimes(name='example_population_spikes', source='source',
                                gid=gids, times=times, data_pointer=data_pointer)
cell_types = CatCellInfo(name='cell_types',source='source',
                         values=cell_types_vals, indices=cell_types_inds, gid=gids)


cell_x_pos = ContCellInfo(name='cell_x_pos', source='source', data=x_pos, gid=gids)



## write to file
f = NWBFile(file_name='tmp.nwb',
            source='me',
            session_description='my first synthetic recording',
            identifier='EXAMPLE_ID',
            session_start_time=datetime.now(),
            experimenter='Dr. Bilbo Baggins',
            lab='Bag End Labatory',
            institution='University of Middle Earth at the Shire',
            experiment_description='empty',
            session_id='LONELYMTN')

population_module = f.create_processing_module(name='0', source='source',
                                               description='description') 
population_module.add_container(pop_data)
population_module.add_container(cell_types)
population_module.add_container(cell_x_pos)


io = NWBHDF5IO('tmp.nwb', mode='w')
io.write(f)
io.close()