ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.modules.save_synapses

    import os
    import csv
    import h5py
    import numpy as np
    from neuron import h
    from glob import glob
    from itertools import product

    from .sim_module import SimulatorMod
    from bmtk.simulator.bionet.biocell import BioCell
    from bmtk.simulator.bionet.io_tools import io
    from bmtk.utils.sonata.utils import add_hdf5_magic, add_hdf5_version
    from bmtk.simulator.bionet.pointprocesscell import PointProcessCell

    pc = h.ParallelContext()
    MPI_RANK = int(pc.id())
    N_HOSTS = int(pc.nhost())

    class SaveSynapses(SimulatorMod):
        def __init__(self, network_dir, single_file=False, **params):
            self._network_dir = network_dir
            self._virt_lookup = {}
            self._gid_lookup = {}
            self._sec_lookup = {}
            if MPI_RANK == 0:
                if not os.path.exists(network_dir):
                    os.makedirs(network_dir)
            pc.barrier()

            #if N_HOSTS > 1:
            #    io.log_exception('save_synapses module is not current supported with mpi')

            self._syn_writer = ConnectionWriter(network_dir)

        def _print_nc(self, nc, src_nid, trg_nid, cell, src_pop, trg_pop, edge_type_id):
            if isinstance(cell, BioCell):
                sec_x = nc.postloc()
                sec = h.cas()
                sec_id = self._sec_lookup[cell.gid][sec]  # cell.get_section_id(sec)
                h.pop_section()
                self._syn_writer.add_bio_conn(edge_type_id, src_nid, src_pop, trg_nid, trg_pop, nc.weight[0], sec_id, sec_x)

            else:
                self._syn_writer.add_point_conn(edge_type_id, src_nid, src_pop, trg_nid, trg_pop, nc.weight[0])

        def initialize(self, sim):
            io.log_info('Saving network connections. This may take a while.')

            # Need a way to look up virtual nodes from nc.pre()
            for pop_name, nodes_table in sim.net._virtual_nodes.items():
                for node_id, virt_node in nodes_table.items():
                    self._virt_lookup[virt_node.hobj] = (pop_name, node_id)

            # Need to figure out node_id and pop_name from nc.srcgid()
            for node_pop in sim.net.node_populations:
                pop_name = node_pop.name
                for node in node_pop[0::1]:
                    if node.model_type != 'virtual':
                        self._gid_lookup[node.gid] = (pop_name, node.node_id)

            for gid, cell in sim.net.get_local_cells().items():
                trg_pop, trg_id = self._gid_lookup[gid]
                if isinstance(cell, BioCell):
                    # sections = cell._syn_seg_ix
                    self._sec_lookup[gid] = {sec_name: sec_id for sec_id, sec_name in enumerate(cell.get_sections_id())}

                else:
                    sections = [-1]*len(cell.netcons)

                for nc, edge_type_id in zip(cell.netcons, cell._edge_type_ids):
                    src_gid = int(nc.srcgid())
                    if src_gid == -1:
                        # source is a virtual node
                        src_pop, src_id = self._virt_lookup[nc.pre()]
                    else:
                        src_pop, src_id = self._gid_lookup[src_gid]

                    self._print_nc(nc, src_id, trg_id, cell, src_pop, trg_pop, edge_type_id)

            self._syn_writer.close()
            pc.barrier()

            if N_HOSTS > 1 and MPI_RANK == 0:
                merger = H5Merger(self._network_dir, self._syn_writer._pop_groups.keys())

            pc.barrier()

            io.log_info('    Done saving network connections.')

    class H5Merger(object):
        def __init__(self, network_dir, grp_keys):
            self._network_dir = network_dir
            self._grp_keys = list(grp_keys)

            self._edge_counts = {(s, t): 0 for s, t in self._grp_keys}
            self._biophys_edge_count = {(s, t): 0 for s, t in self._grp_keys}
            self._point_edge_count = {(s, t): 0 for s, t in self._grp_keys}

            self._tmp_files = {(s, t): [] for s, t in self._grp_keys}
            for (src_pop, trg_pop), r in product(self._grp_keys, range(N_HOSTS)):
                fname = '.core{}.{}_{}_edges.h5'.format(r, src_pop, trg_pop)
                fpath = os.path.join(self._network_dir, fname)
                if not os.path.exists(fpath):
                    io.log_warning('Expected file {} is missing'.format(fpath))

                h5file = h5py.File(fpath, 'r')
                edges_grp = h5file['/edges/{}_{}'.format(src_pop, trg_pop)]
                self._tmp_files[(src_pop, trg_pop)].append(edges_grp)

                self._edge_counts[(src_pop, trg_pop)] += len(edges_grp['source_node_id'])
                self._biophys_edge_count[(src_pop, trg_pop)] += len(edges_grp['0/syn_weight'])
                self._point_edge_count[(src_pop, trg_pop)] += len(edges_grp['1/syn_weight'])

            for (src_pop, trg_pop), in_grps in self._tmp_files.items():
                out_h5 = h5py.File(os.path.join(self._network_dir, '{}_{}_edges.h5'.format(src_pop, trg_pop)), 'w')
                add_hdf5_magic(out_h5)
                add_hdf5_version(out_h5)
                pop_root = out_h5.create_group('/edges/{}_{}'.format(src_pop, trg_pop))
                n_edges_total = self._edge_counts[(src_pop, trg_pop)]
                n_edges_bio = self._biophys_edge_count[(src_pop, trg_pop)]
                n_edges_point = self._point_edge_count[(src_pop, trg_pop)]

                pop_root.create_dataset('source_node_id', (n_edges_total, ), dtype=np.uint64)
                pop_root['source_node_id'].attrs['node_population'] = src_pop

                pop_root.create_dataset('target_node_id', (n_edges_total, ), dtype=np.uint64)
                pop_root['target_node_id'].attrs['node_population'] = trg_pop

                pop_root.create_dataset('edge_group_id', (n_edges_total, ), dtype=np.uint16)
                pop_root.create_dataset('edge_group_index', (n_edges_total,), dtype=np.uint16)
                pop_root.create_dataset('edge_type_id', (n_edges_total, ), dtype=np.uint32)

                pop_root.create_dataset('0/syn_weight', (n_edges_bio, ), dtype=float)
                pop_root.create_dataset('0/sec_id', (n_edges_bio, ), dtype=np.uint64)
                pop_root.create_dataset('0/sec_x', (n_edges_bio, ), dtype=float)
                pop_root.create_dataset('1/syn_weight', (n_edges_point, ), dtype=float)

                total_offset = 0
                bio_offset = 0
                point_offset = 0
                for grp in in_grps:
                    n_ds = len(grp['source_node_id'])
                    pop_root['source_node_id'][total_offset:(total_offset + n_ds)] = grp['source_node_id'][()]
                    pop_root['target_node_id'][total_offset:(total_offset + n_ds)] = grp['target_node_id'][()]
                    pop_root['edge_group_id'][total_offset:(total_offset + n_ds)] = grp['edge_group_id'][()]
                    pop_root['edge_group_index'][total_offset:(total_offset + n_ds)] = grp['edge_group_index'][()]
                    pop_root['edge_type_id'][total_offset:(total_offset + n_ds)] = grp['edge_type_id'][()]
                    total_offset += n_ds

                    n_ds = len(grp['0/syn_weight'])
                    # print(grp['0/syn_weight'][()])
                    pop_root['0/syn_weight'][bio_offset:(bio_offset + n_ds)] = grp['0/syn_weight'][()]
                    pop_root['0/sec_id'][bio_offset:(bio_offset + n_ds)] = grp['0/sec_id'][()]
                    pop_root['0/sec_x'][bio_offset:(bio_offset + n_ds)] = grp['0/sec_x'][()]
                    bio_offset += n_ds

                    n_ds = len(grp['1/syn_weight'])
                    pop_root['1/syn_weight'][point_offset:(point_offset + n_ds)] = grp['1/syn_weight'][()]
                    point_offset += n_ds

                    fname = grp.file.filename
                    grp.file.close()
                    if os.path.exists(fname):
                        os.remove(fname)

                self._create_index(pop_root, index_type='target')
                self._create_index(pop_root, index_type='source')
                out_h5.close()

        def _create_index(self, pop_root, index_type='target'):
            if index_type == 'target':
                edge_nodes = np.array(pop_root['target_node_id'], dtype=np.int64)
                output_grp = pop_root.create_group('indices/target_to_source')
            elif index_type == 'source':
                edge_nodes = np.array(pop_root['source_node_id'], dtype=np.int64)
                output_grp = pop_root.create_group('indices/source_to_target')

            edge_nodes = np.append(edge_nodes, [-1])
            n_targets = np.max(edge_nodes)
            ranges_list = [[] for _ in range(n_targets + 1)]

            n_ranges = 0
            begin_index = 0
            cur_trg = edge_nodes[begin_index]
            for end_index, trg_gid in enumerate(edge_nodes):
                if cur_trg != trg_gid:
                    ranges_list[cur_trg].append((begin_index, end_index))
                    cur_trg = int(trg_gid)
                    begin_index = end_index
                    n_ranges += 1

            node_id_to_range = np.zeros((n_targets + 1, 2))
            range_to_edge_id = np.zeros((n_ranges, 2))
            range_index = 0
            for node_index, trg_ranges in enumerate(ranges_list):
                if len(trg_ranges) > 0:
                    node_id_to_range[node_index, 0] = range_index
                    for r in trg_ranges:
                        range_to_edge_id[range_index, :] = r
                        range_index += 1
                    node_id_to_range[node_index, 1] = range_index

            output_grp.create_dataset('range_to_edge_id', data=range_to_edge_id, dtype='uint64')
            output_grp.create_dataset('node_id_to_range', data=node_id_to_range, dtype='uint64')

    class ConnectionWriter(object):

        class H5Index(object):
            def __init__(self, file_path, src_pop, trg_pop):
                # TODO: Merge with NetworkBuilder code for building SONATA files
                self._nsyns = 0
                self._n_biosyns = 0
                self._n_pointsyns = 0
                self._block_size = 5

                self._pop_name = '{}_{}'.format(src_pop, trg_pop)
                # self._h5_file = h5py.File(os.path.join(network_dir, '{}_edges.h5'.format(self._pop_name)), 'w')
                self._h5_file = h5py.File(file_path, 'w')
                add_hdf5_magic(self._h5_file)
                add_hdf5_version(self._h5_file)
                self._pop_root = self._h5_file.create_group('/edges/{}'.format(self._pop_name))
                self._pop_root.create_dataset('edge_group_id', (self._block_size, ), dtype=np.uint16,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root.create_dataset('source_node_id', (self._block_size, ), dtype=np.uint64,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root['source_node_id'].attrs['node_population'] = src_pop
                self._pop_root.create_dataset('target_node_id', (self._block_size, ), dtype=np.uint64,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root['target_node_id'].attrs['node_population'] = trg_pop
                self._pop_root.create_dataset('edge_type_id', (self._block_size, ), dtype=np.uint32,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root.create_dataset('0/syn_weight', (self._block_size, ), dtype=float,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root.create_dataset('0/sec_id', (self._block_size, ), dtype=np.uint64,
                                              chunks=(self._block_size, ), maxshape=(None, ))
                self._pop_root.create_dataset('0/sec_x', (self._block_size, ), chunks=(self._block_size, ),
                                              maxshape=(None, ), dtype=float)
                self._pop_root.create_dataset('1/syn_weight', (self._block_size, ), dtype=float,
                                              chunks=(self._block_size, ), maxshape=(None, ))

            def _add_conn(self, edge_type_id, src_id, trg_id, grp_id):
                self._pop_root['edge_type_id'][self._nsyns] = edge_type_id
                self._pop_root['source_node_id'][self._nsyns] = src_id
                self._pop_root['target_node_id'][self._nsyns] = trg_id
                self._pop_root['edge_group_id'][self._nsyns] = grp_id

                self._nsyns += 1
                if self._nsyns % self._block_size == 0:
                    self._pop_root['edge_type_id'].resize((self._nsyns + self._block_size,))
                    self._pop_root['source_node_id'].resize((self._nsyns + self._block_size, ))
                    self._pop_root['target_node_id'].resize((self._nsyns + self._block_size, ))
                    self._pop_root['edge_group_id'].resize((self._nsyns + self._block_size, ))

            def add_bio_conn(self, edge_type_id, src_id, trg_id, syn_weight, sec_id, sec_x):
                self._add_conn(edge_type_id, src_id, trg_id, 0)
                self._pop_root['0/syn_weight'][self._n_biosyns] = syn_weight
                self._pop_root['0/sec_id'][self._n_biosyns] = sec_id
                self._pop_root['0/sec_x'][self._n_biosyns] = sec_x

                self._n_biosyns += 1
                if self._n_biosyns % self._block_size == 0:
                    self._pop_root['0/syn_weight'].resize((self._n_biosyns + self._block_size, ))
                    self._pop_root['0/sec_id'].resize((self._n_biosyns + self._block_size, ))
                    self._pop_root['0/sec_x'].resize((self._n_biosyns + self._block_size, ))

            def add_point_conn(self, edge_type_id, src_id, trg_id, syn_weight):
                self._add_conn(edge_type_id, src_id, trg_id, 1)
                self._pop_root['1/syn_weight'][self._n_pointsyns] = syn_weight

                self._n_pointsyns += 1
                if self._n_pointsyns % self._block_size == 0:
                    self._pop_root['1/syn_weight'].resize((self._n_pointsyns + self._block_size, ))

            def clean_ends(self):
                self._pop_root['source_node_id'].resize((self._nsyns,))
                self._pop_root['target_node_id'].resize((self._nsyns,))
                self._pop_root['edge_group_id'].resize((self._nsyns,))
                self._pop_root['edge_type_id'].resize((self._nsyns,))

                self._pop_root['0/syn_weight'].resize((self._n_biosyns,))
                self._pop_root['0/sec_id'].resize((self._n_biosyns,))
                self._pop_root['0/sec_x'].resize((self._n_biosyns,))

                self._pop_root['1/syn_weight'].resize((self._n_pointsyns,))

                eg_ds = self._pop_root.create_dataset('edge_group_index', (self._nsyns, ), dtype=np.uint64)
                bio_count, point_count = 0, 0
                for idx, grp_id in enumerate(self._pop_root['edge_group_id']):
                    if grp_id == 0:
                        eg_ds[idx] = bio_count
                        bio_count += 1
                    elif grp_id == 1:
                        eg_ds[idx] = point_count
                        point_count += 1

                self._create_index('target')
                self._create_index('source')

            def _create_index(self, index_type='target'):
                if index_type == 'target':
                    edge_nodes = np.array(self._pop_root['target_node_id'], dtype=np.int64)
                    output_grp = self._pop_root.create_group('indices/target_to_source')
                elif index_type == 'source':
                    edge_nodes = np.array(self._pop_root['source_node_id'], dtype=np.int64)
                    output_grp = self._pop_root.create_group('indices/source_to_target')

                edge_nodes = np.append(edge_nodes, [-1])
                n_targets = np.max(edge_nodes)
                ranges_list = [[] for _ in range(n_targets + 1)]

                n_ranges = 0
                begin_index = 0
                cur_trg = edge_nodes[begin_index]
                for end_index, trg_gid in enumerate(edge_nodes):
                    if cur_trg != trg_gid:
                        ranges_list[cur_trg].append((begin_index, end_index))
                        cur_trg = int(trg_gid)
                        begin_index = end_index
                        n_ranges += 1

                node_id_to_range = np.zeros((n_targets + 1, 2))
                range_to_edge_id = np.zeros((n_ranges, 2))
                range_index = 0
                for node_index, trg_ranges in enumerate(ranges_list):
                    if len(trg_ranges) > 0:
                        node_id_to_range[node_index, 0] = range_index
                        for r in trg_ranges:
                            range_to_edge_id[range_index, :] = r
                            range_index += 1
                        node_id_to_range[node_index, 1] = range_index

                output_grp.create_dataset('range_to_edge_id', data=range_to_edge_id, dtype='uint64')
                output_grp.create_dataset('node_id_to_range', data=node_id_to_range, dtype='uint64')

            def close_h5(self):
                self._h5_file.close()

        def __init__(self, network_dir):
            self._network_dir = network_dir
            self._pop_groups = {}

        def _group_key(self, src_pop, trg_pop):
            return (src_pop, trg_pop)

        def _get_edge_group(self, src_pop, trg_pop):
            grp_key = self._group_key(src_pop, trg_pop)
            if grp_key not in self._pop_groups:
                pop_name = '{}_{}'.format(src_pop, trg_pop)
                if N_HOSTS > 1:
                    pop_name = '.core{}.{}'.format(MPI_RANK, pop_name)

                file_path = os.path.join(self._network_dir, '{}_edges.h5'.format(pop_name))
                self._pop_groups[grp_key] = self.H5Index(file_path, src_pop, trg_pop)

            return self._pop_groups[grp_key]

        def add_bio_conn(self, edge_type_id, src_id, src_pop, trg_id, trg_pop, syn_weight, sec_id, sec_x):
            h5_grp = self._get_edge_group(src_pop, trg_pop)
            h5_grp.add_bio_conn(edge_type_id, src_id, trg_id, syn_weight, sec_id, sec_x)

        def add_point_conn(self, edge_type_id, src_id, src_pop, trg_id, trg_pop, syn_weight):
            h5_grp = self._get_edge_group(src_pop, trg_pop)
            h5_grp.add_point_conn(edge_type_id, src_id, trg_id, syn_weight)

        def close(self):
            for _, h5index in self._pop_groups.items():
                h5index.clean_ends()
                h5index.close_h5()