Add NESTML support to BMTK

bmtk/simulator/pointnet/pointnetwork.py

@@ -76,7 +76,6 @@ def __init__(self, **properties):
         self._params_cache = {}
 
         self._virtual_ids_map = {}
-
         self._batch_nodes = True
 
         # self._nest_id_map = {}
@@ -90,6 +89,8 @@ def __init__(self, **properties):
         self._gid_map = GidPool()
         self._virtual_gids = GidPool()
 
+        self._nestml_models = []
+
     @property
     def py_function_caches(self):
         return pyfunction_cache
@@ -168,6 +169,9 @@ def build_recurrent_edges(self, force_resolution=False):
             for edge in edge_pop.get_edges():
                 nest_srcs = self.gid_map.get_nestids(edge_pop.source_nodes, edge.source_node_ids)
                 nest_trgs = self.gid_map.get_nestids(edge_pop.target_nodes, edge.target_node_ids)
+                if np.isscalar(edge.nest_params['weight']):
+                    edge.nest_params['weight'] = np.full(shape=len(nest_srcs),
+                                                         fill_value=edge.nest_params['weight'])
                 self._nest_connect(nest_srcs, nest_trgs, conn_spec='one_to_one', syn_spec=edge.nest_params)
 
     def find_edges(self, source_nodes=None, target_nodes=None):
@@ -219,6 +223,9 @@ def add_spike_trains(self, spike_trains, node_set, sg_params={'precise_times': T
                 for edge in edge_pop.get_edges():
                     nest_trgs = self.gid_map.get_nestids(edge_pop.target_nodes, edge.target_node_ids)
                     nest_srcs = virt_gid_map.get_nestids(edge_pop.source_nodes, edge.source_node_ids)
+                    if np.isscalar(edge.nest_params['weight']):
+                        edge.nest_params['weight'] = np.full(shape=len(nest_srcs),
+                                                             fill_value=edge.nest_params['weight'])
                     self._nest_connect(nest_srcs, nest_trgs, conn_spec='one_to_one', syn_spec=edge.nest_params)
 
     def _nest_connect(self, nest_srcs, nest_trgs, conn_spec='one_to_one', syn_spec=None):
@@ -230,7 +237,7 @@ def _nest_connect(self, nest_srcs, nest_trgs, conn_spec='one_to_one', syn_spec=N
             # An occuring issue is when dt > delay, add some extra messaging in log to help users fix problem.
             res_kernel = nest.GetKernelStatus().get('resolution', 'NaN')
             delay_edges = syn_spec.get('delay', 'NaN')
-            msg = 'synaptic "delay" value in edges ({}) is not compatable with simulator resolution/"dt" ({})'.format(
+            msg = 'synaptic "delay" value in edges ({}) is not compatible with simulator resolution/"dt" ({})'.format(
                 delay_edges, res_kernel
             )
             self.io.log_error('{}{}'.format(bde.errorname, bde.errormessage))
@@ -241,3 +248,25 @@ def _nest_connect(self, nest_srcs, nest_trgs, conn_spec='one_to_one', syn_spec=N
             # Record exception to log file.
             self.io.log_error(str(e))
             raise
+
+    def add_nestml_models(self, modelname):
+        if modelname not in self._nestml_models:
+            self._nestml_models.append(modelname)
+
+    def initialize_nestml(self, rebuild_nestml=True):
+        from pynestml.frontend.pynestml_frontend import generate_nest_target
+        models_dir = os.path.commonpath([self.get_component('point_neuron_models_dir'), self.get_component('synaptic_models_dir')])
+        nestml_model_name_suffix = '_nestml'
+
+        if self._nestml_models:
+            if rebuild_nestml or not (os.path.exists(os.path.join(models_dir, 'nestml_target'))):
+                if not rebuild_nestml:
+                    self.io.log_info('Cannot find nestml target, must rebuild')
+                generate_nest_target(input_path=self._nestml_models,
+                                     logging_level='ERROR',
+                                     suffix=nestml_model_name_suffix,
+                                     target_path=os.path.join(models_dir, 'nestml_target'))
+                self.io.log_info('Generating NESTML models: ', self._nestml_models)
+                # a module by the same name can only be loaded once; do this at the very end of the function
+            nest.Install('nestmlmodule')
+

bmtk/simulator/pointnet/pointsimulator.py

@@ -47,6 +47,7 @@ def __init__(self, graph, dt=0.001, overwrite=True, print_time=False):
 
         self._cells_built = False
         self._internal_connections_built = False
+        self._rebuild_nestml = True
 
         self._graph = graph
         self._external_cells = {}  # dict-of-dict of external pointnet cells with keys [network_name][cell_id]
@@ -228,6 +229,11 @@ def from_config(cls, configure, graph):
         if run_dict.get('allow_offgrid_spikes', False):
             network.set_spike_generator_params(allow_offgrid_spikes=True)
 
+        if 'rebuild_nestml' in config:
+            network._rebuild_nestml = config['rebuild_nestml']
+
+        graph.initialize_nestml(network._rebuild_nestml)
+
         # Create the output-directory, or delete existing files if it already exists
         graph.io.log_info('Setting up output directory')
         if not os.path.exists(config['output']['output_dir']):
@@ -276,7 +282,7 @@ def from_config(cls, configure, graph):
                 mod = mods.SpikesMod(**report.params)
 
             elif isinstance(report, reports.MembraneReport):
-                # For convience and for compliance with SONATA format. "membrane_report" and "multimeter_report is the
+                # For convenience and for compliance with SONATA format. "membrane_report" and "multimeter_report is the
                 # same in pointnet.
                 mod = mods.MultimeterMod(**report.params)
 

bmtk/simulator/pointnet/sonata_adaptors.py

@@ -1,7 +1,9 @@
+import logging
 import numpy as np
 from collections import Counter
 import numbers
 import nest
+import os
 import types
 import pandas as pd
 
@@ -11,6 +13,8 @@
 from bmtk.simulator.pointnet.glif_utils import convert_aibs2nest
 from bmtk.simulator.pointnet.nest_utils import nest_version
 
+logger = logging.getLogger(__name__)
+
 NEST_SYNAPSE_MODEL_PROP = 'model' if nest_version[0] == 2 else 'synapse_model'
 
 
@@ -155,6 +159,7 @@ def get_batches(self, node_group):
     def preprocess_node_types(network, node_population):
         NodeAdaptor.preprocess_node_types(network, node_population)
         node_types_table = node_population.types_table
+
         if 'model_template' in node_types_table.columns and 'dynamics_params' in node_types_table.columns:
             node_type_ids = np.unique(node_population.type_ids)
             for nt_id in node_type_ids:
@@ -166,6 +171,20 @@ def preprocess_node_types(network, node_population):
                     node_type_attrs['model_template'] = model_template
                     node_type_attrs['dynamics_params'] = dynamics_params
 
+                if mtemplate.startswith('nestml:'):
+                    # NESTML model was requested
+
+                    model_name = mtemplate.split(':')[1]
+
+                    models_dir = network.get_component('point_neuron_models_dir')
+                    model_fn = os.path.join(models_dir, model_name + '.nestml')
+                    network.add_nestml_models(model_fn)
+
+                    # replace model_template name with suffixed version
+
+                    node_types_table[nt_id]['model_template'] = 'nest:' + model_name + '_nestml'
+                    node_types_table._df_cache = None   # clear the table internal cache
+
     @staticmethod
     def patch_adaptor(adaptor, node_group, network):
         node_adaptor = NodeAdaptor.patch_adaptor(adaptor, node_group, network)
@@ -256,14 +275,24 @@ def preprocess_edge_types(network, edge_population):
         # Fix for sonata/300_pointneurons
         EdgeAdaptor.preprocess_edge_types(network, edge_population)
         edge_types_table = edge_population.types_table
-        edge_type_ids = np.unique(edge_population.type_ids)
-
-        for et_id in edge_type_ids:
-            edge_type = edge_types_table[et_id]
-            if 'model_template' in edge_types_table.columns:
-                model_template = edge_type['model_template']
-                if model_template.startswith('nest'):
-                    edge_type['model_template'] = model_template[5:]
+
+        if 'model_template' in edge_types_table.columns and 'dynamics_params' in edge_types_table.columns:
+            edge_type_ids = np.unique(edge_population.type_ids)
+            for et_id in edge_type_ids:
+                edge_type_attrs = edge_types_table[et_id]
+                mtemplate = edge_type_attrs['model_template']
+                if mtemplate.startswith('nest:'):
+                    edge_type['model_template'] = mtemplate[5:]
+
+                if mtemplate.startswith('nestml:'):
+                    model_name = mtemplate.split(':')[1]
+
+                    models_dir = network.get_component('synaptic_models_dir')
+                    model_fn = os.path.join(models_dir, model_name + '.nestml')
+                    network.add_nestml_models(model_fn)
+
+                    edge_types_table[et_id]['model_template'] = model_name + '_nestml'
+                    edge_types_table._df_cache = None  # clear the table internal cache
 
     def get_batches(self, edge_group):
         src_ids = {}

examples/point_120cells_nestml/README.md

@@ -0,0 +1,69 @@
+# 120 point neuron network
+
+A small network of 120 point-neurons. Neuron and synapse models are specified as NESTML models to demonstrate how to incorporate NESTML models. Uses PointNet and will require NEST and NESTML to run.
+
+
+## Running:
+To run a simulation, install bmtk and run the following:
+```
+$ python run_pointnet.py config.simulation.json
+```
+If successful, will create a *output* directory containing log, spike trains and recorded cell variables.
+
+## The Network:
+The network files have already been built and stored as SONATA files in the *network/* directory. The bmtk Builder
+script used to create the network files is *build_network.py*. To adjust the parameters and/or topology of the network
+change this file and run:
+```
+$ python build_network.py
+```
+This will overwrite the existing files in the network directory. Note that there is some randomness in how the network
+is built, so expect (slightly) different simulation results everytime the network is rebuilt.
+
+## Simulation Parameters
+Parameters to run the simulation, including run-time, inputs, recorded variables, and networks are stored in config.json
+and can modified with a text editor.
+
+## Plotting results
+```
+$ python plot_output.py
+```
+
+## Perturbation simulations
+The file ```config_pertubration.json``` uses the same network (recurrent connections + feedforward inputs) as before. However
+is designed to also simulate the effects of optogenetic or current clamp inhibition and excitation on a subset of the cells. 
+To run this example:
+```bash
+$ python run_pointnet.py config_perturbations.json
+``` 
+
+The only difference between this simulation and the original is in the **inputs** section in the config. To simulate the 
+perturbations we use a current clamp input to make cells 20 through 39 overly excitatory:
+```json
+"exc_perturbation": {
+   "input_type": "current_clamp",
+   "module": "IClamp",
+   "node_set": {
+     "population": "cortex",
+     "node_ids": [20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]
+   },
+   "amp": 230.0,
+   "delay": 1.0,
+   "duration": 3000.0
+},
+```
+
+For cells 40 through 49 we can use a negative current to depress the cell activity
+```json
+"inh_perturbation": {
+  "input_type": "current_clamp",
+  "module": "IClamp",
+  "node_set": {
+    "population": "cortex",
+    "node_ids": [40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
+  },
+  "amp": -230.0,
+  "delay": 1.0,
+  "duration": 3000.0
+}
+```

examples/point_120cells_nestml/build_network.py

@@ -0,0 +1,116 @@
+import numpy as np
+
+from bmtk.builder import NetworkBuilder
+from bmtk.builder.auxi.node_params import positions_columinar
+
+
+lif_models = {
+    'LIF_exc': {
+        'N': 80,
+        'ei': 'e',
+        'pop_name': 'LIF_exc',
+        'model_type': 'point_neuron',
+        'model_template': 'nestml:iaf_psc_alpha',
+        'dynamics_params': 'iaf_psc_delta_exc.json',
+    },
+    'LIF_inh': {
+        'N': 40,
+        'ei': 'i',
+        'pop_name': 'LIF_inh',
+        'model_type': 'point_neuron',
+        'model_template': 'nestml:iaf_psc_delta',
+        'dynamics_params': 'iaf_psc_delta_inh.json',
+    }
+}
+
+input_network_model = {
+    'input_network': {
+        'N': 100,
+        'ei': 'e',
+        'pop_name': 'input_network',
+        'model_type': 'virtual'
+    }
+}
+
+
+def random_connections(source, target, p=0.1):
+    sid = source['node_id']  # Get source id
+    tid = target['node_id']  # Get target id
+
+    # Avoid self-connections.
+    if sid == tid:
+        return None
+
+    return np.random.binomial(1, p)  # nsyns
+
+
+def build_cortex_network():
+    cortex = NetworkBuilder('cortex')
+    for model in lif_models:
+        params = lif_models[model].copy()
+        positions = positions_columinar(N=lif_models[model]['N'], center=[0, 10.0, 0], max_radius=50.0, height=200.0)
+        cortex.add_nodes(
+            x=positions[:, 0],
+            y=positions[:, 1],
+            z=positions[:, 2],
+            **params
+        )
+
+    cortex.add_edges(
+        source={'ei': 'e'},
+        connection_rule=random_connections,
+        connection_params={'p': 0.1},
+        syn_weight=2.0,  # 2.0
+        delay=1.5,
+        dynamics_params='ExcToInh.json',
+        model_template='static_synapse'
+    )
+
+    cortex.add_edges(
+        source={'ei': 'i'},
+        connection_rule=random_connections,
+        connection_params={'p': 0.1},
+        syn_weight=-1.5,
+        delay=1.5,
+        dynamics_params='InhToExc.json',
+        model_template='static_synapse'
+    )
+
+    cortex.build()
+    cortex.save(output_dir='network')
+
+    return cortex
+
+
+def build_thalamus_network(cortex):
+    thalamus = NetworkBuilder('thalamus')
+    thalamus.add_nodes(**input_network_model['input_network'])
+
+    thalamus.add_edges(
+        target=cortex.nodes(ei='e'),
+        connection_rule=random_connections,
+        connection_params={'p': 0.1},
+        syn_weight=220.0, 
+        delay=1.5,
+        dynamics_params='ExcToExc.json',
+        model_template='nestml:static_synapse'
+    )
+
+    thalamus.add_edges(
+        target=cortex.nodes(ei='i'),
+        connection_rule=random_connections,
+        connection_params={'p': 0.1},
+        syn_weight=5.0,
+        delay=1.5,
+        dynamics_params='ExcToExc.json',
+        model_template='nestml:static_synapse'
+    )
+    thalamus.build()
+    thalamus.save(output_dir='network')
+
+    return thalamus
+
+
+if __name__ == '__main__':
+    cortex_net = build_cortex_network()
+    thalamus_net = build_thalamus_network(cortex_net)