ALLEN INSTITUTE

# Source code for bmtk.utils.io.spike_trains

    import os
    import sys
    import csv

    import h5py
    import pandas as pd
    import numpy as np
    from bmtk.utils.sonata.utils import add_hdf5_magic, add_hdf5_version

    class SpikeTrainWriter(object):

        class TmpFileMetadata(object):
            def __init__(self, file_name, sort_order=None):
                self.file_name = file_name
                self.sort_order = sort_order

        def __init__(self, tmp_dir, mpi_rank=0, mpi_size=1):
            # For NEST/NEURON based simulations it is prefereable not to use mpi4py, so let the parent simulator determine
            # MPI rank and size
            self._mpi_rank = mpi_rank
            self._mpi_size = mpi_size

            # used to temporary save spike files since for large simulations saving spikes into memory can crash the
            # system. Requires the user to create the directory
            self._tmp_dir = tmp_dir
            if self._tmp_dir is None or not os.path.exists(self._tmp_dir):
                raise Exception('Directory path {} does not exists'.format(self._tmp_dir))
            self._all_tmp_files = [self.TmpFileMetadata(self._get_tmp_filename(r)) for r in range(mpi_size)]
            # TODO: Determine best buffer size.
            self._tmp_file_handle = open(self._all_tmp_files[mpi_rank].file_name, 'w')

            self._tmp_spikes_handles = []  # used when sorting mulitple file
            self._spike_count = -1

            # Nest gid files uses tab seperators and a different order for tmp spike files.
            self.delimiter = ' '  # delimiter for temporary file
            self.time_col = 0
            self.gid_col = 1

        def _get_tmp_filename(self, rank):
            return os.path.join(self._tmp_dir, '_bmtk_tmp_spikes_{}.csv'.format(rank))

        def _count_spikes(self, recount=False):
            if self._mpi_rank == 0:
                if self._spike_count > -1 and not recount:
                    return self._spike_count

                self._spike_count = 0
                for tmp_file in self._all_tmp_files:
                    with open(tmp_file.file_name, 'r') as csvfile:
                        csv_reader = csv.reader(csvfile, delimiter=self.delimiter)
                        self._spike_count += sum(1 for _ in csv_reader)

        def _sort_tmp_file(self, filedata, sort_order):
            # For now load spikes into pandas, it's the fastest way but may be an issue with memory
            if sort_order is None or filedata.sort_order == sort_order:
                return

            file_name = filedata.file_name
            tmp_spikes_ds = pd.read_csv(file_name, sep=' ', names=['time', 'gid'])
            tmp_spikes_ds = tmp_spikes_ds.sort_values(by=[sort_order, 'time'])
            tmp_spikes_ds.to_csv(file_name, sep=' ', index=False, header=False)
            filedata.sort_order = sort_order

        def _next_spike(self, rank):
            try:
                val = next(self._tmp_spikes_handles[rank])
                return val[0], val[1], rank
            except StopIteration:
                return None

        def add_spike(self, time, gid):
            self._tmp_file_handle.write('{:.6f} {}\n'.format(time, gid))

        def add_spikes(self, times, gid):
            for t in times:
                self.add_spike(t, gid)

        def add_spikes_file(self, file_name, sort_order=None):
            self._all_tmp_files.append(self.TmpFileMetadata(file_name, sort_order))

        def _sort_files(self, sort_order, sort_column, file_write_fnc):
            self._tmp_spikes_handles = []
            for fdata in self._all_tmp_files:
                self._sort_tmp_file(fdata, sort_order)
                self._tmp_spikes_handles.append(csv.reader(open(fdata.file_name, 'r'), delimiter=self.delimiter))

            spikes = []
            for rank in range(len(self._tmp_spikes_handles)):  # range(self._mpi_size):
                spike = self._next_spike(rank)
                if spike is not None:
                    spikes.append(spike)

            # Iterate through all the ranks and find the first spike. Write that spike/gid to the output, then
            # replace that data point with the next spike on the selected rank
            indx = 0
            while spikes:
                # find which rank has the first spike
                selected_index = 0
                selected_val = spikes[0][sort_column]
                for i, spike in enumerate(spikes[1:]):
                    if spike[sort_column] < selected_val:
                        selected_index = i + 1
                        selected_val = spike[sort_column]

                # write the spike to the file
                row = spikes.pop(selected_index)
                file_write_fnc(float(row[self.time_col]), int(row[self.gid_col]), indx)
                indx += 1

                # get the next spike on that rank and replace in spikes table
                another_spike = self._next_spike(row[2])
                if another_spike is not None:
                    spikes.append(another_spike)

        def _merge_files(self, file_write_fnc):
            indx = 0
            for fdata in self._all_tmp_files:
                if not os.path.exists(fdata.file_name):
                    continue

                with open(fdata.file_name, 'r') as csv_file:
                    csv_reader = csv.reader(csv_file, delimiter=self.delimiter)
                    for row in csv_reader:
                        file_write_fnc(float(row[self.time_col]), int(row[self.gid_col]), indx)
                        indx += 1

        def _to_file(self, file_name, sort_order, file_write_fnc):
            if sort_order is None:
                sort_column = 0
            elif sort_order == 'time':
                sort_column = self.time_col
            elif sort_order == 'gid':
                sort_column = self.gid_col
            else:
                raise Exception('Unknown sort order {}'.format(sort_order))

            # TODO: Need to make sure an MPI_Barrier is called beforehand
            self._tmp_file_handle.close()
            if self._mpi_rank == 0:
                if sort_order is not None:
                    self._sort_files(sort_order, sort_column, file_write_fnc)
                else:
                    self._merge_files(file_write_fnc)

        def to_csv(self, csv_file, sort_order=None, gid_map=None):
            # TODO: Need to call flush and then barrier
            if self._mpi_rank == 0:
                # For the single rank case don't just copy the tmp-csv to the new name. It will fail if user calls to_hdf5
                # or to_nwb after calling to_csv.
                self._count_spikes()
                csv_handle = open(csv_file, 'w')
                csv_writer = csv.writer(csv_handle, delimiter=' ')

                def file_write_fnc_identity(time, gid, indx):
                    csv_writer.writerow([time, gid])

                def file_write_fnc_transform(time, gid, indx):
                    # For the case when NEURON/NEST ids don't match with the user's gid table
                    csv_writer.writerow([time, gid_map[gid]])

                file_write_fnc = file_write_fnc_identity if gid_map is None else file_write_fnc_transform
                self._to_file(csv_file, sort_order, file_write_fnc)
                csv_handle.close()

            # TODO: Let user pass in in barrier and use it here

        def to_nwb(self, nwb_file):
            raise NotImplementedError

        def to_hdf5(self, hdf5_file, sort_order=None, gid_map=None):
            if self._mpi_rank == 0:
                with h5py.File(hdf5_file, 'w') as h5:
                    add_hdf5_magic(h5)
                    add_hdf5_version(h5)

                    self._count_spikes(recount=True)
                    spikes_grp = h5.create_group('/spikes')
                    spikes_grp.attrs['sorting'] = 'none' if sort_order is None else sort_order
                    time_ds = spikes_grp.create_dataset('timestamps', shape=(self._spike_count,), maxshape=(None,), dtype=np.float64)
                    gid_ds = spikes_grp.create_dataset('gids', shape=(self._spike_count,), maxshape=(None,), dtype=np.uint64)

                    def resize_data():
                        # There have been (unreproducable) mpi conditons where _count_spikes() is not correct, even with
                        # proper barriers and file flushing. Add a quick fix in case when converting csv to hdf5.
                        self._count_spikes(recount=True)
                        time_ds.resize((self._spike_count, ))
                        gid_ds.resize((self._spike_count, ))

                    def file_write_fnc_identity(time, gid, indx):
                        if indx >= self._spike_count:
                            resize_data()
                        time_ds[indx] = time
                        gid_ds[indx] = gid

                    def file_write_fnc_transform(time, gid, indx):
                        if indx >= self._spike_count:
                            resize_data()
                        time_ds[indx] = time
                        gid_ds[indx] = gid_map[gid]

                    file_write_fnc = file_write_fnc_identity if gid_map is None else file_write_fnc_transform
                    self._to_file(hdf5_file, sort_order, file_write_fnc)

            # TODO: Need to make sure a barrier is used here (before close is called)

        def flush(self):
            self._tmp_file_handle.flush()

        def close_tmp_file(self):
            self._tmp_file_handle.close()

        def close(self):
            if self._mpi_rank == 0:
                for tmp_file in self._all_tmp_files:
                    if os.path.exists(tmp_file.file_name):
                        os.remove(tmp_file.file_name)

    class PoissonSpikesGenerator(object):
        def __init__(self, gids, firing_rate, tstart=0.0, tstop=1000.0):
            self._gids = gids
            self._firing_rate = firing_rate / 1000.0
            self._tstart = tstart
            self._tstop = tstop

        def to_hdf5(self, file_name, sort_order='gid'):
            if sort_order == 'gid':
                gid_list = []
                times_list = []
                if sort_order == 'gid':
                    for gid in self._gids:
                        c_time = self._tstart
                        while c_time < self._tstop:
                            interval = -np.log(1.0 - np.random.uniform()) / self._firing_rate
                            c_time += interval
                            gid_list.append(gid)
                            times_list.append(c_time)

                with h5py.File(file_name, 'w') as h5:
                    h5.create_dataset('/spikes/gids', data=gid_list, dtype=np.uint)
                    h5.create_dataset('/spikes/timestamps', data=times_list, dtype=float)
                    h5['/spikes'].attrs['sorting'] = 'by_gid'

            else:
                raise NotImplementedError

    class SpikesInput(object):

        def get_spikes(self, gid):
            raise NotImplementedError()

        @staticmethod
        def load(name, module, input_type, params):
            module_lc = module.lower()
            if module_lc == 'nwb':
                return SpikesInputNWBv1(name, module, input_type, params)
            elif module_lc == 'h5' or module_lc == 'hdf5':
                return SpikesInputH5(name, module, input_type, params)
            elif module_lc == 'csv':
                return SpikesInputCSV(name, module, input_type, params)
            else:
                raise Exception('Unable to load spikes for module type {}'.format(module))

    class SpikesInputNWBv1(SpikesInput):
        def __init__(self, name, module, input_type, params):
            self.input_file = params['input_file']
            self._h5_handle = h5py.File(self.input_file, 'r')
            self.trial = params['trial']
            self._trial_grp = self._h5_handle['processing'][self.trial]['spike_train']

        def get_spikes(self, gid):
            return self._trial_grp[str(gid)]['data']

    class SONATAIndexer(object):

        def get_spikes(self, gid):
            raise NotImplementedError()

    class DictIndexedGIDs(SONATAIndexer):
        def __init__(self, spikes_input_h5):
            self._gid_indicies = {}
            self._parent = spikes_input_h5

            indx_beg = 0
            c_gid = self._parent.gids[0]
            for indx, gid in enumerate(self._parent.gids):
                # go through the gids dataset, determine slices for each gid
                if gid != c_gid:
                    self._gid_indicies[c_gid] = slice(indx_beg, indx)
                    c_gid = gid
                    indx_beg = indx
            self._gid_indicies[c_gid] = slice(indx_beg, indx+1)  # saves the last entry

        def get_spikes(self, gid):
            if gid in self._gid_indicies:
                return self._parent.timestamps[self._gid_indicies[gid]]
            else:
                return []

    class DFIndexedGIDs(SONATAIndexer):
        def __init__(self, spikes_input_h5):
            self._parent = spikes_input_h5

            index_df = pd.DataFrame(data={'gids': self._parent.gids}, )
            index_df.drop_duplicates(inplace=True)
            index_df = index_df.stack().reset_index()
            index_df.columns = ['indx_beg', 'tmp', 'gid']
            index_df.set_index('gid', inplace=True)
            index_df = index_df.drop('tmp', axis=1)
            index_df['indx_end'] = index_df['indx_beg'].shift(-1).fillna(len(self._parent.gids)).astype(np.int64)
            self._gid_indicies = index_df  # index_df.to_dict(orient='index')

        def get_spikes(self, gid):
            if gid in self._gid_indicies.index:
                indx_beg = self._gid_indicies.loc[gid]['indx_beg']
                indx_end = self._gid_indicies.loc[gid]['indx_end']
                return self._parent.timestamps[indx_beg:indx_end]
            else:
                return []

    class HDF5IndexedGIDs(SONATAIndexer):
        def __init__(self, spikes_input_h5):
            raise NotImplementedError()

    class UnindexedGIDs(SONATAIndexer):
        def __init__(self, spikes_input_h5):
            raise NotImplementedError()

    class SpikesInputH5(SpikesInput):
        def __init__(self, name, module, input_type, params):
            self._input_file = params['input_file']
            self._h5_handle = h5py.File(self._input_file, 'r')
            self._gid_ds = self._h5_handle['/spikes/gids']
            self._timestamps_ds = self._h5_handle['/spikes/timestamps']
            self._sort_order = self._h5_handle['/spikes'].attrs.get('sorting', None)
            if sys.version_info[0] >= 3 and isinstance(self._sort_order, bytes) and self._sort_order is not None:
                # h5py attributes return str in py 2, bytes in py 3
                self._sort_order = self._sort_order.decode()

            # Create an index for fetching spike-trains based on gid
            if 'index' in self._h5_handle['/spikes'] and isinstance(self._h5_handle['/spikes/index'], h5py.Dataset):
                # In the case the index is built into the hdf5 file
                self._gid_index = HDF5IndexedGIDs(self)

            elif self._sort_order in ['gid', 'gids', 'by_gid']:
                # In the case when the sort_order == 'gid'
                # self._gid_index = DictIndexedGIDs(self)
                self._gid_index = DFIndexedGIDs(self)

            else:
                # In the case where the spike-trains are sorted by time or unsorted
                self._gid_index = UnindexedGIDs(self)

        @property
        def gids(self):
            return self._gid_ds

        @property
        def timestamps(self):
            return self._timestamps_ds

        def get_spikes(self, gid):
            return self._gid_index.get_spikes(gid)

    class SpikesInputCSV(SpikesInput):
        def __init__(self, name, module, input_type, params):
            self._spikes_df = pd.read_csv(params['input_file'], index_col='gid', sep=' ')

        def get_spikes(self, gid):
            spike_times_str = self._spikes_df.iloc[gid]['spike-times']
            return np.array(spike_times_str.split(','), dtype=float)