ALLEN INSTITUTE

# Source code for bmtk.simulator.core.sonata_reader.edge_adaptor

    import os
    import ast
    import json
    import types

    import numpy as np

    class SonataBaseEdge(object):
        def __init__(self, sonata_edge, edge_adaptor):
            self._edge = sonata_edge
            self._prop_adaptor = edge_adaptor

        @property
        def source_node_id(self):
            return self._edge.source_node_id

        @property
        def source_population(self):
            return self._edge.source_population

        @property
        def target_node_id(self):
            return self._edge.target_node_id

        @property
        def target_population(self):
            return self._edge.target_population

        @property
        def dynamics_params(self):
            return self._prop_adaptor.dynamics_params(self._edge)

        @property
        def delay(self):
            return self._edge['delay']

        @property
        def weight_function(self):
            return self._prop_adaptor.weight_function(self._edge)

        @property
        def preselected_targets(self):
            return self._prop_adaptor.preselected_targets

        @property
        def target_sections(self):
            return self._edge['target_sections']

        @property
        def target_distance(self):
            return self._edge['distance_range']

        @property
        def is_gap_junction(self):
            try:
                return self._edge['is_gap_junction']
            except:
                return False

        @property
        def edge_type_id(self):
            return self._edge.edge_type_id

        @property
        def nsyns(self):
            return self._prop_adaptor.nsyns(self._edge)

        @property
        def model_template(self):
            return self._edge['model_template']

        def syn_weight(self, src_node, trg_node):
            return self._prop_adaptor.syn_weight(self, src_node=src_node, trg_node=trg_node)

        def __getitem__(self, item):
            return self._edge[item]

    class EdgeAdaptor(object):
        def __init__(self, network):
            self._network = network
            self._func_caches = self._network.py_function_caches

        @property
        def batch_process(self):
            return False

        @batch_process.setter
        def batch_process(self, flag):
            pass

        def get_edge(self, sonata_node):
            return SonataBaseEdge(sonata_node, self)

        @staticmethod
        def preprocess_edge_types(network, edge_population):
            edge_types_table = edge_population.types_table
            edge_type_ids = np.unique(edge_population.type_ids)
            cached_dynamics = {}

            for et_id in edge_type_ids:
                edge_type = edge_types_table[et_id]
                if 'dynamics_params' in edge_types_table.columns and edge_type['dynamics_params'] != None:
                    dynamics_params = edge_type['dynamics_params']
                    params_dir = network.get_component('synaptic_models_dir')

                    params_path = os.path.join(params_dir, dynamics_params)

                    # see if we can load the dynamics_params as a dictionary. Otherwise just save the file path and let the
                    # cell_model loader function handle the extension.
                    # Cache the loaded dynamics_params to minimize file access.
                    if params_path in cached_dynamics:
                        edge_type['dynamics_params'] = cached_dynamics[params_path]
                    else:
                        try:
                            params_val = json.load(open(params_path, 'r'))
                            edge_type['dynamics_params'] = params_val
                            cached_dynamics[params_path] = params_val
                        except Exception:
                            # TODO: Check dynamics_params before
                            network.io.log_exception('Could not find edge dynamics_params file {}.'.format(params_path))

                # Split target_sections
                if 'target_sections' in edge_type:
                    trg_sec = edge_type['target_sections']
                    if trg_sec is not None:
                        try:
                            edge_type['target_sections'] = ast.literal_eval(trg_sec)
                        except Exception as exc:
                            network.io.log_warning('Unable to split target_sections list {}'.format(trg_sec))
                            edge_type['target_sections'] = None

                # Split target distances
                if 'distance_range' in edge_type:
                    dist_range = edge_type['distance_range']
                    if dist_range is not None:
                        try:
                            # TODO: Make the distance range has at most two values
                            edge_type['distance_range'] = json.loads(dist_range)
                        except Exception as e:
                            try:
                                edge_type['distance_range'] = [0.0, float(dist_range)]
                            except Exception as e:
                                network.io.log_warning('Unable to parse distance_range {}'.format(dist_range))
                                edge_type['distance_range'] = None

        @classmethod
        def create_adaptor(cls, edge_group, network):
            prop_map = cls(network)
            return cls.patch_adaptor(prop_map, edge_group)

        @staticmethod
        def patch_adaptor(adaptor, edge_group):
            # dynamics_params
            if edge_group.has_dynamics_params:
                adaptor.dynamics_params = types.MethodType(group_dynamics_params, adaptor)
            else:  # 'dynamics_params' in node_group.all_columns:
                adaptor.dynamics_params = types.MethodType(types_dynamics_params, adaptor)

            # For fetching/calculating synaptic weights
            if 'weight_function' in edge_group.all_columns:
                # Customized function for user to calculate the synaptic weight
                adaptor.weight_function = types.MethodType(weight_function, adaptor)
                adaptor.syn_weight = types.MethodType(syn_weight_function, adaptor)
            elif 'syn_weight' in edge_group.all_columns:
                # Just return the synaptic weight
                adaptor.weight_function = types.MethodType(ret_none_function, adaptor)
                adaptor.syn_weight = types.MethodType(syn_weight, adaptor)
            else:
                raise Exception('Could not find syn_weight or weight_function properties. Cannot create connections.')

            # For determining the synapse placement
            if 'sec_id' in edge_group.all_columns or 'afferent_section_id' in edge_group.all_columns:
                adaptor.preselected_targets = True
                adaptor.nsyns = types.MethodType(no_nsyns, adaptor)
            elif 'nsyns' in edge_group.all_columns:
                adaptor.preselected_targets = False
                adaptor.nsyns = types.MethodType(nsyns, adaptor)
            else:
                # It will get here for connections onto point neurons
                adaptor.preselected_targets = True
                adaptor.nsyns = types.MethodType(no_nsyns, adaptor)

            return adaptor

    def ret_none_function(self, edge):
        return None

    def weight_function(self, edge):
        return edge['weight_function']

    def syn_weight(self, edge, src_node, trg_node):
        return edge['syn_weight']

    def syn_weight_function(self, edge, src_node, trg_node):
        weight_fnc_name = edge.weight_function
        if weight_fnc_name is None:
            weight_fnc = self._func_caches.py_modules.synaptic_weight('default_weight_fnc')
            return weight_fnc(edge, src_node, trg_node)

        elif self._func_caches.py_modules.has_synaptic_weight(weight_fnc_name):
            weight_fnc = self._func_caches.py_modules.synaptic_weight(weight_fnc_name)
            return weight_fnc(edge, src_node, trg_node)

        else:
            self._network.io.log_exception('weight_function {} is not defined.'.format(weight_fnc_name))

    def nsyns(self, edge):
        return edge['nsyns']

    def no_nsyns(self, edge):
        return 1

    def types_dynamics_params(self, node):
        return node['dynamics_params']

    def group_dynamics_params(self, node):
        return node.dynamics_params