From 356518324a364a394160d18967f679de238573e6 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Mon, 29 Apr 2024 17:49:56 -0700
Subject: [PATCH 01/11] Making changes to automatically support ModelDB Hay
 model HOC Template

---
 bmtk/simulator/bionet/config.py               | 12 ++++-
 .../bionet/default_setters/cell_models.py     | 52 ++++++++++++++-----
 bmtk/simulator/bionet/nrn.py                  | 19 +++++--
 bmtk/simulator/bionet/sonata_adaptors.py      |  2 +
 4 files changed, 67 insertions(+), 18 deletions(-)

diff --git a/bmtk/simulator/bionet/config.py b/bmtk/simulator/bionet/config.py
index 5d8a9e4a1..1cb3c626f 100644
--- a/bmtk/simulator/bionet/config.py
+++ b/bmtk/simulator/bionet/config.py
@@ -39,7 +39,12 @@ def create_output_dir(self):
         io.setup_output_dir(self.output_dir, self.log_file)
 
     def load_nrn_modules(self):
-        nrn.load_neuron_modules(self.mechanisms_dir, self.templates_dir)
+        nrn.load_neuron_modules(
+            mechanisms_dir=self.mechanisms_dir, 
+            templates_dir=self.templates_dir, 
+            default_templates=self.use_default_templates,
+            use_old_import3d=self.use_old_import3d
+        )
 
     def build_env(self):
         self.io = io
@@ -52,3 +57,8 @@ def build_env(self):
 
         pc.barrier()
         self.load_nrn_modules()
+
+    def _set_class_props(self):
+        super(Config, self)._set_class_props()
+        self.use_old_import3d = self.run.get('use_old_import3d', False)
+        self.use_default_templates = self.run.get('use_old_import3d', True)
diff --git a/bmtk/simulator/bionet/default_setters/cell_models.py b/bmtk/simulator/bionet/default_setters/cell_models.py
index eff40594e..16d2342b0 100644
--- a/bmtk/simulator/bionet/default_setters/cell_models.py
+++ b/bmtk/simulator/bionet/default_setters/cell_models.py
@@ -23,6 +23,7 @@
 import os
 import numpy as np
 from neuron import h
+import inspect
 try:
     from sklearn.decomposition import PCA
 except Exception as e:
@@ -41,19 +42,46 @@
 """
 
 def loadHOC(cell, template_name, dynamics_params):
-    # Get template to instantiate
-    template_call = getattr(h, template_name)
-    if dynamics_params is not None and 'params' in dynamics_params:
-        template_params = dynamics_params['params']
-        if isinstance(template_params, list):
-            # pass in a list of parameters
-            hobj = template_call(*template_params)
+    """A Generic function for creating a cell object from a NEURON HOC Template (eg. a *.hoc file with 
+    `begintemplate template_name` in header). It essentially tries to guess the correct parameters that need to be
+    called so may not work the majority of the times and require to be overloaded.
+
+    :param cell: A SONATA node object, can be used as a dict to get individual properties of current cell.
+    :param template_name: name of HOCTemplate as stored in "model_template" attribute (hoc:<template_name>).
+    :param dynamics_params: Dictionary containing contents of cell['dynamics_params'] as loaded from a json file or hdf5. 
+        If cell does not have "dynamics_params" attributes then will be set to None.
+    """
+    try:
+        # Get template to instantiate
+        template_call = getattr(h, template_name)
+    except AttributeError as ae:
+        io.log_error(
+            f'loadHOC was unable to load in Neuron HOC Template "{template_name}, ' 
+            'Make sure appropiate .hoc file is stored in templates_dir.'
+        )
+        raise ae
+
+    try:
+        if dynamics_params is not None and 'params' in dynamics_params:
+            template_params = dynamics_params['params']
+            if isinstance(template_params, list):
+                # pass in a list of parameters
+                hobj = template_call(*template_params)
+            else:
+                # only a single parameter
+                hobj = template_call(template_params)
+        elif cell.morphology_file is not None:
+            # instantiate template with no parameters
+            hobj = template_call(cell.morphology_file)
         else:
-            # only a single parameter
-            hobj = template_call(template_params)
-    else:
-        # instantiate template with no parameters
-        hobj = template_call()
+            hobj = template_call()
+    except RuntimeError as rte:
+        io.log_error(
+            f'bmtk.simualtor.bionet.default_setters.cell_models.loadHOC function failed to load HOC template "{template_call}". '
+            'If Hoc Templates requires special call to be initialized consider using `bmtk.simulator.bionet.add_cell_model()` '
+            'to overwrite this function.'
+        )
+        raise rte
 
     # TODO: All "all" section if it doesn't exist
     # hobj.all = h.SectionList()
diff --git a/bmtk/simulator/bionet/nrn.py b/bmtk/simulator/bionet/nrn.py
index 4edf5f75b..cc97a95f4 100644
--- a/bmtk/simulator/bionet/nrn.py
+++ b/bmtk/simulator/bionet/nrn.py
@@ -46,18 +46,27 @@ def clear_gids():
     pc.barrier()
 
 
-def load_neuron_modules(mechanisms_dir, templates_dir, default_templates=True):
+def load_neuron_modules(mechanisms_dir, templates_dir, default_templates=True, use_old_import3d=False):
     """
 
     :param mechanisms_dir:
     :param templates_dir:
     :param default_templates:
     """
-    h.load_file('stdgui.hoc')
-
     bionet_dir = os.path.dirname(__file__)
-    h.load_file(os.path.join(bionet_dir, 'import3d.hoc').replace("\\","/"))  # customized import3d.hoc to supress warnings
-    h.load_file(os.path.join(bionet_dir,'default_templates',  'advance.hoc').replace("\\","/"))
+    
+    h.load_file('stdgui.hoc')   
+    
+    if use_old_import3d:
+        # Older versions of BMTK used a modified import3d.hoc that is saved in the current directory which
+        # due to being out-of-date can have issues with newer HOC models. Should be replaced fully, but 
+        # until we know the side-effects have a flag to use the old "import3d.hoc" file
+        h.load_file(os.path.join(bionet_dir, 'import3d.hoc').replace("\\","/"))  
+    else:
+        # This will load the import3d.hoc that is saved as a part of nrniv.
+        h.load_file('import3d.hoc')
+    
+    h.load_file(os.path.join(bionet_dir, 'default_templates',  'advance.hoc').replace("\\","/"))
 
     if isinstance(mechanisms_dir, list):
         for mdir in mechanisms_dir:
diff --git a/bmtk/simulator/bionet/sonata_adaptors.py b/bmtk/simulator/bionet/sonata_adaptors.py
index 93427b6df..ee17eaa45 100644
--- a/bmtk/simulator/bionet/sonata_adaptors.py
+++ b/bmtk/simulator/bionet/sonata_adaptors.py
@@ -34,6 +34,8 @@ def position(self):
 
     @property
     def morphology_file(self):
+        if 'morphology' not in self._node:
+            return None
         return self._node['morphology']
 
     @property

From 82adf09d1bc625a132422515b17928917febaabc Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Tue, 30 Apr 2024 17:25:13 -0700
Subject: [PATCH 02/11] small change to morphology to fix how sections are
 parsed

---
 bmtk/simulator/bionet/morphology.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bmtk/simulator/bionet/morphology.py b/bmtk/simulator/bionet/morphology.py
index b53f2779d..4f24b747c 100644
--- a/bmtk/simulator/bionet/morphology.py
+++ b/bmtk/simulator/bionet/morphology.py
@@ -63,7 +63,7 @@ def get(self):
 
             for sec_id, sec in enumerate(self._hobj.all):
                 fullsecname = sec.name()
-                sec_type = fullsecname.split(".")[1][:4]  # get sec name type without the cell name
+                sec_type = fullsecname.split(".")[-1][:4]  # get sec name type without the cell name
                 sec_type_swc = Morphology.sec_type_swc[sec_type]  # convert to swc code
                 x_range = 1.0 / (sec.nseg * 2)  # used to calculate [x0, x1]
 

From f1532371d0d12ebc38b8eeeb5fbf854cd807c586 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Wed, 1 May 2024 18:19:32 -0700
Subject: [PATCH 03/11] Found Hoc models that would fail calls to hobj.soma[0]
 since it cell had only one soma section and hobj.soma was not an iteratble
 list. Added Cell.soma property that can handle this case

---
 bmtk/simulator/bionet/biocell.py              | 14 +++++++---
 bmtk/simulator/bionet/cell.py                 |  3 +++
 bmtk/simulator/bionet/iclamp.py               | 16 ++++++++---
 .../bionet/modules/record_cellvars.py         |  6 +++--
 bmtk/simulator/bionet/morphology.py           | 27 ++++++++++++++-----
 bmtk/simulator/bionet/seclamp.py              |  1 +
 6 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/bmtk/simulator/bionet/biocell.py b/bmtk/simulator/bionet/biocell.py
index 71d8cef1c..187ff60bf 100644
--- a/bmtk/simulator/bionet/biocell.py
+++ b/bmtk/simulator/bionet/biocell.py
@@ -27,6 +27,7 @@
 from bmtk.simulator.bionet.morphology import Morphology
 import six
 
+import neuron
 from neuron import h
 
 pc = h.ParallelContext()    # object to access MPI methods
@@ -74,9 +75,6 @@ class BioCell(Cell):
     def __init__(self, node, population_name, bionetwork):
         super(BioCell, self).__init__(node=node, population_name=population_name, network=bionetwork)
 
-        # Set up netcon object that can be used to detect and communicate cell spikes.
-        self.set_spike_detector(bionetwork.spike_threshold)
-
         # Determine number of segments and store a list of all sections.
         self._secs = []
         self._secs_by_id = []
@@ -105,6 +103,10 @@ def __init__(self, node, population_name, bionetwork):
         self._seg_coords = None
         self.build_morphology()
 
+        # Set up netcon object that can be used to detect and communicate cell spikes.
+        self.set_spike_detector(bionetwork.spike_threshold)
+
+
     def build_morphology(self):
         morph_base = Morphology.load(hobj=self.hobj, morphology_file=self.morphology_file, cache_seg_props=True)
 
@@ -126,6 +128,10 @@ def morphology(self):
         """The actual Morphology object instanstiation"""
         return self._morphology
 
+    @property
+    def soma(self):
+        return self.morphology.soma
+
     @property
     def seg_coords(self):
         """Coordinates for segments/sections of the morphology, need to make public for ecp, xstim, and other
@@ -144,7 +150,7 @@ def seg_coords(self):
         return self.morphology.seg_coords
 
     def set_spike_detector(self, spike_threshold):
-        nc = h.NetCon(self.hobj.soma[0](0.5)._ref_v, None, sec=self.hobj.soma[0])  # attach spike detector to cell
+        nc = h.NetCon(self.soma[0](0.5)._ref_v, None, sec=self.soma[0])
         nc.threshold = spike_threshold
         pc.cell(self.gid, nc)  # associate gid with spike detector
 
diff --git a/bmtk/simulator/bionet/cell.py b/bmtk/simulator/bionet/cell.py
index 4236e8e08..837352cfe 100644
--- a/bmtk/simulator/bionet/cell.py
+++ b/bmtk/simulator/bionet/cell.py
@@ -106,6 +106,9 @@ def get_connection_info(self):
 
     def set_syn_connections(self, edge_prop, src_node, stim=None):
         raise NotImplementedError
+    
+    def get_section(self, sec_name, sec_index):
+        raise NotImplementedError
 
     def __getitem__(self, node_prop):
         return self._node[node_prop]
diff --git a/bmtk/simulator/bionet/iclamp.py b/bmtk/simulator/bionet/iclamp.py
index 758b5addb..eccf78a22 100644
--- a/bmtk/simulator/bionet/iclamp.py
+++ b/bmtk/simulator/bionet/iclamp.py
@@ -38,8 +38,10 @@ def __init__(self, amplitude, delay, duration, section_name='soma', section_inde
 
     def attach_current(self, cell):
         # self._stim = h.IClamp(cell.hobj.soma[0](0.5))
-        hboj_sec = getattr(cell.hobj, self._section_name)[self._section_index](self._section_dist)
-        self._stim = h.IClamp(hboj_sec)
+        # hboj_sec = getattr(cell.hobj, self._section_name)[self._section_index](self._section_dist)
+        hobj_sec = cell.get_section(sec_name=self._section_name, sec_index=self._section_index)(self._section_dist)
+
+        self._stim = h.IClamp(hobj_sec)
         self._stim.delay = self._iclamp_del
         self._stim.dur = self._iclamp_dur
         self._stim.amp = self._iclamp_amp
@@ -47,13 +49,19 @@ def attach_current(self, cell):
 
 
 class FileIClamp(object):
-    def __init__(self, amplitudes, dt):
+    def __init__(self, amplitudes, dt, section_name='soma', section_index=0, section_dist=0.5):
         self._iclamp_amps = amplitudes
         self._iclamp_dt = dt
         self._stim = None
 
+        self._section_name = section_name
+        self._section_index = section_index
+        self._section_dist = section_dist
+
     def attach_current(self, cell):
-        self._stim = h.IClamp(cell.hobj.soma[0](0.5))
+        hobj_sec = cell.get_section(sec_name=self._section_name, sec_index=self._section_index)(self._section_dist)
+
+        self._stim = h.IClamp(hobj_sec)
 
         # Listed as necessary values in the docs to use play() with an IClamp.
         self._stim.delay = 0
diff --git a/bmtk/simulator/bionet/modules/record_cellvars.py b/bmtk/simulator/bionet/modules/record_cellvars.py
index e407ca723..cbd11ed4c 100644
--- a/bmtk/simulator/bionet/modules/record_cellvars.py
+++ b/bmtk/simulator/bionet/modules/record_cellvars.py
@@ -314,7 +314,8 @@ def step(self, sim, tstep, rel_time=0.0):
             pop_id = self._gid_map.get_pool_id(gid)
             cell = sim.net.get_cell_gid(gid)
             for var_name in self._variables:
-                var_val = getattr(cell.hobj.soma[0](0.5), var_name)
+                var_val = getattr(cell.soma[0](0.5), var_name)
+                # var_val = getattr(cell.soma(0.5), var_name)
                 self._var_recorder.record_cell(
                     pop_id.node_id,
                     population=pop_id.population,
@@ -323,7 +324,8 @@ def step(self, sim, tstep, rel_time=0.0):
                 )
 
             for var_name, fnc in self._transforms.items():
-                var_val = getattr(cell.hobj.soma[0](0.5), var_name)
+                # var_val = getattr(cell.hobj.soma[0](0.5), var_name)
+                var_val = getattr(cell.soma(0.5), var_name)
                 new_val = fnc(var_val)
                 self._var_recorder.record_cell(
                     pop_id.node_id,
diff --git a/bmtk/simulator/bionet/morphology.py b/bmtk/simulator/bionet/morphology.py
index 4f24b747c..00a29f99c 100644
--- a/bmtk/simulator/bionet/morphology.py
+++ b/bmtk/simulator/bionet/morphology.py
@@ -26,6 +26,7 @@
 import numpy as np
 import pandas as pd
 from neuron import h
+import neuron
 
 from bmtk.simulator.bionet.io_tools import io
 
@@ -35,8 +36,9 @@
 
 class _LazySegmentProps(object):
     """A Struct for storing coordinate invariant properties of each segment of a cell."""
-    def __init__(self, hobj):
-        self._hobj = hobj
+    def __init__(self, morphology):
+        self._morphology = morphology
+        self._hobj = morphology.hobj
         self.sec_id = None
         self.type = None
         self.area = None
@@ -59,7 +61,7 @@ def get(self):
             seg_length = []
             seg_sec_id = []
 
-            h.distance(sec=self._hobj.soma[0])  # measure distance relative to the soma
+            h.distance(sec=self._morphology.soma[0])  # measure distance relative to the soma
 
             for sec_id, sec in enumerate(self._hobj.all):
                 fullsecname = sec.name()
@@ -99,8 +101,8 @@ def get(self):
 
 class _LazySegmentCoords(object):
     """A Struct for storing properties of each segment of a cell that vary by soma position and rotation."""
-    def __init__(self, hobj):
-        self._hobj = hobj
+    def __init__(self, morphology):
+        self._hobj = morphology.hobj
         self.p0 = None
         self.p1 = None
         self.p05 = None
@@ -235,15 +237,26 @@ def __init__(self, hobj, rng_seed=None, swc_path=None):
         self._prng = np.random.RandomState(self.rng_seed)
         self._sections = None
         self._segments = None
-        self._seg_props = _LazySegmentProps(self.hobj)  # None
-        self._seg_coords = _LazySegmentCoords(self.hobj)  # None
+        self._seg_props = _LazySegmentProps(self)  # None
+        self._seg_coords = _LazySegmentCoords(self)  # None
         self._nseg = None
         self._swc_map = None
+        self._soma = None # 
 
         # Used by find_sections() and other methods when building edges/synapses. Should make it faster to look-up
         # cooresponding segments for a large number of syns that target the same area of a cell
         self._trg_segs_cache = {}
 
+    @property
+    def soma(self):
+        if self._soma is None:
+            if isinstance(self.hobj.soma, (neuron.nrn.Section, neuron.nrn.Segment)):
+                self._soma = [self.hobj.soma]
+            else:
+                self._soma = self.hobj.soma
+
+        return self._soma
+
     def _copy(self):
         new_morph = Morphology(hobj=self.hobj, swc_path=self.swc_path)
         return new_morph
diff --git a/bmtk/simulator/bionet/seclamp.py b/bmtk/simulator/bionet/seclamp.py
index 1830fcbae..928eaaf30 100644
--- a/bmtk/simulator/bionet/seclamp.py
+++ b/bmtk/simulator/bionet/seclamp.py
@@ -31,6 +31,7 @@ def __init__(self, amplitudes, durations, rs=None):
         self._stim = None
 
     def attach_current(self, cell):
+        # TODO: Make sure to 
         self._stim = h.SEClamp(cell.hobj.soma[0](0.5))
         self._stim.dur1 = self._seclamp_durs[0]
         self._stim.dur2 = self._seclamp_durs[1]

From d66d82ef1774a99fcfc6ffdc51a03d1a65768cc9 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Thu, 2 May 2024 17:34:49 -0700
Subject: [PATCH 04/11] adding ability for pointnet to use cell_model decorator
 when loading cells in batch

---
 .../pointnet/default_setters/__init__.py      |  3 ++
 .../pointnet/default_setters/cell_models.py   | 10 ++++
 bmtk/simulator/pointnet/sonata_adaptors.py    | 52 +++++++++++++++++--
 bmtk/utils/sonata/group.py                    |  4 ++
 4 files changed, 64 insertions(+), 5 deletions(-)
 create mode 100644 bmtk/simulator/pointnet/default_setters/cell_models.py

diff --git a/bmtk/simulator/pointnet/default_setters/__init__.py b/bmtk/simulator/pointnet/default_setters/__init__.py
index b07cc2dcd..6e3037054 100644
--- a/bmtk/simulator/pointnet/default_setters/__init__.py
+++ b/bmtk/simulator/pointnet/default_setters/__init__.py
@@ -1,2 +1,5 @@
+# NOTE: It is important to keep these imports here so that the default functions for loading
+# cells and synapses are called. 
+from . import cell_models
 from . import synaptic_weights
 from . import synapse_models
\ No newline at end of file
diff --git a/bmtk/simulator/pointnet/default_setters/cell_models.py b/bmtk/simulator/pointnet/default_setters/cell_models.py
new file mode 100644
index 000000000..00f847456
--- /dev/null
+++ b/bmtk/simulator/pointnet/default_setters/cell_models.py
@@ -0,0 +1,10 @@
+import nest
+
+from bmtk.simulator.pointnet.pyfunction_cache import add_cell_model, add_cell_processor
+
+
+def loadNESTModel(cell, template_name, dynamics_params):
+    return nest.Create(template_name, cell.n_nodes, dynamics_params)
+
+
+add_cell_model(loadNESTModel, directive='nest', model_type='point_neuron', overwrite=False)
diff --git a/bmtk/simulator/pointnet/sonata_adaptors.py b/bmtk/simulator/pointnet/sonata_adaptors.py
index bb3296f95..68c1856da 100644
--- a/bmtk/simulator/pointnet/sonata_adaptors.py
+++ b/bmtk/simulator/pointnet/sonata_adaptors.py
@@ -24,12 +24,14 @@ def all_null(node_group, column_name):
 
 
 class PointNodeBatched(object):
-    def __init__(self, node_ids, gids, node_types_table, node_type_id):
+    def __init__(self, node_ids, gids, node_types_table, node_type_id, grp_indices, node_group):
         self._n_nodes = len(node_ids)
         self._node_ids = node_ids
         self._gids = gids
         self._nt_table = node_types_table
         self._nt_id = node_type_id
+        self._grp_indices = grp_indices
+        self._node_group = node_group
         self._nest_objs = []
         self._nest_ids = []
 
@@ -53,6 +55,18 @@ def nest_ids(self):
     def nest_model(self):
         return self._nt_table[self._nt_id]['model_template'].split(':')[1]
 
+    @property
+    def model_name(self):
+        return self.model_template.split(':')[1]
+    
+    @property
+    def model_schema(self):
+        return self.model_template.split(':')[0]
+    
+    @property
+    def model_template(self):
+        return self._nt_table[self._nt_id]['model_template'] # .split(':')
+
     @property
     def nest_params(self):
         return self._nt_table[self._nt_id]['dynamics_params']
@@ -62,9 +76,35 @@ def model_type(self):
         return self._nt_table[self._nt_id]['model_type']
 
     def build(self):
-        self._nest_objs = nest.Create(self.nest_model, self.n_nodes, self.nest_params)
+        model_name = self.nest_model
+        model_type = self.model_type
+
+        if py_modules.has_cell_model(self.model_template, self.model_type):
+            cell_fnc = py_modules.cell_model(self.model_template, model_type)
+            self._nest_objs = cell_fnc(self, model_name, self.nest_params)
+        elif py_modules.has_cell_model(self.model_schema, self.model_type):
+            cell_fnc = py_modules.cell_model(self.model_schema, model_type)
+            self._nest_objs = cell_fnc(self, model_name, self.nest_params)
+        else:
+            self._nest_objs = nest.Create(self.nest_model, self.n_nodes, self.nest_params)
+               
         self._nest_ids = self._nest_objs.tolist() if nest_version[0] >= 3 else self._nest_objs
 
+    def __getitem__(self, key):
+        if key in self._node_group or key in self._nt_table.columns:
+            # Since it's possible that this batch only contains a subset of the dataset for this
+            # group, we must filter-out and sort by the passed in indices.
+            grp_vals = self._node_group.get_values(key)
+            return grp_vals[self._grp_indices]
+        elif key in ['node_id', 'node_ids']:
+            return self.node_ids
+        elif key in ['node_type_id', 'node_type_ids']:
+            return np.full(self._n_nodes, self._nt_id, dtype=None)
+        else:
+            msg = f'Could not find property {key} in either nodes or node-types table for {self._node_group}'
+            io.log_error(msg)
+            raise KeyError(msg)
+
 
 class PointNode(SonataBaseNode):
     def __init__(self, node, prop_adaptor):
@@ -140,15 +180,17 @@ def get_batches(self, node_group):
 
         nid_groups = {nt_id: np.zeros(ntids_counter[nt_id], dtype=np.uint32) for nt_id in ntids_counter}
         gid_groups = {nt_id: np.zeros(ntids_counter[nt_id], dtype=np.uint32) for nt_id in ntids_counter}
+        grp_indices = {nt_id: np.zeros(ntids_counter[nt_id], dtype=np.uint32) for nt_id in ntids_counter}
         node_groups_counter = {nt_id: 0 for nt_id in ntids_counter}
 
-        for node_id, gid, node_type_id in zip(node_ids, node_gids, node_type_ids):
+        for node_id, gid, node_type_id, idx in zip(node_ids, node_gids, node_type_ids, range(len(node_ids))):
             grp_indx = node_groups_counter[node_type_id]
             nid_groups[node_type_id][grp_indx] = node_id
             gid_groups[node_type_id][grp_indx] = gid
+            grp_indices[node_type_id][grp_indx] = idx
             node_groups_counter[node_type_id] += 1
-
-        return [PointNodeBatched(nid_groups[nt_id], gid_groups[nt_id], node_group.parent.node_types_table, nt_id)
+        
+        return [PointNodeBatched(nid_groups[nt_id], gid_groups[nt_id], node_group.parent.node_types_table, nt_id, grp_indices[nt_id], node_group)
                 for nt_id in ntids_counter]
 
     @staticmethod
diff --git a/bmtk/utils/sonata/group.py b/bmtk/utils/sonata/group.py
index 589d5a6a8..4a7ab4189 100644
--- a/bmtk/utils/sonata/group.py
+++ b/bmtk/utils/sonata/group.py
@@ -145,6 +145,10 @@ def __contains__(self, prop_name):
         """Search that a column name exists in this group"""
         return prop_name in self._group_column_names
 
+    def __repr__(self):
+        rep = f'<SONATA population subgroup: {self.parent.name}/{self.group_id}>'
+        return rep
+
 
 class NodeGroup(Group):
     def __init__(self, group_id, h5_group, parent):

From 22a9bdc095b3d249a5465e7c9aa07da8ddcb1cea Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Fri, 3 May 2024 17:01:03 -0700
Subject: [PATCH 05/11] making updates to spontaneous syns inputs

---
 bmtk/simulator/bionet/biocell.py          | 31 ++++++++++++++++++++---
 bmtk/simulator/bionet/bionetwork.py       |  8 +++---
 bmtk/simulator/bionet/biosimulator.py     |  1 +
 bmtk/simulator/bionet/cell.py             |  3 +++
 bmtk/simulator/bionet/pointprocesscell.py | 19 +++++++++++++-
 5 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/bmtk/simulator/bionet/biocell.py b/bmtk/simulator/bionet/biocell.py
index 187ff60bf..ea5850915 100644
--- a/bmtk/simulator/bionet/biocell.py
+++ b/bmtk/simulator/bionet/biocell.py
@@ -443,18 +443,41 @@ def __init__(self, node, population_name, bionetwork):
         self._vecstim = h.VecStim()
         self._vecstim.play(self._spike_trains)
 
-        self._precell_filter = bionetwork.spont_syns_filter
+        self._precell_filter = bionetwork.spont_syns_filter_pre
+        self._postcell_filter = bionetwork.spont_syns_filter_post
         assert(isinstance(self._precell_filter, dict))
 
-    def _matches_filter(self, src_node):
+    def _matches_filter(self, src_node, trg_node=None):
         """Check to see if the presynaptic cell matches the criteria specified"""
         for k, v in self._precell_filter.items():
+            # Some key may not show up as node_variable
+            if k == 'population' and k not in src_node:
+                key_val = src_node.population_name
+            else:
+                key_val = src_node[k]
+
+            if isinstance(v, (list, tuple)):
+                if key_val not in v:
+                    return False
+            else:
+                if key_val != v:
+                    return False
+        
+        trg_node = trg_node or self
+        for k, v in self._postcell_filter.items():
+            # Some key may not show up as node_variable
+            if k == 'population' and k not in trg_node:
+                key_val = trg_node._node.population_name
+            else:
+                key_val = trg_node[k]
+
             if isinstance(v, (list, tuple)):
-                if src_node[k] not in v:
+                if key_val not in v:
                     return False
             else:
-                if src_node[k] != v:
+                if key_val != v:
                     return False
+        
         return True
 
     def _set_connections(self, edge_prop, src_node, syn_weight, stim=None):
diff --git a/bmtk/simulator/bionet/bionetwork.py b/bmtk/simulator/bionet/bionetwork.py
index 432c9133f..522b6b4da 100644
--- a/bmtk/simulator/bionet/bionetwork.py
+++ b/bmtk/simulator/bionet/bionetwork.py
@@ -77,7 +77,8 @@ def __init__(self):
         self._gid_pool = GidPool()
 
         self.has_spont_syns = False
-        self.spont_syns_filter = None
+        self.spont_syns_filter_pre = None
+        self.spont_syns_filter_post = None
         self.spont_syns_times = None
 
     @property
@@ -88,7 +89,7 @@ def gid_pool(self):
     def py_function_caches(self):
         return nrn
 
-    def set_spont_syn_activity(self, precell_filter, timestamps):
+    def set_spont_syn_activity(self, precell_filter, postcell_filter, timestamps):
         self._model_type_map = {
             'biophysical': BioCellSpontSyn,
             'point_process': PointProcessCellSpontSyns,
@@ -98,7 +99,8 @@ def set_spont_syn_activity(self, precell_filter, timestamps):
         }
 
         self.has_spont_syns = True
-        self.spont_syns_filter = precell_filter
+        self.spont_syns_filter_pre = precell_filter
+        self.spont_syns_filter_post = postcell_filter
         self.spont_syns_times = timestamps
 
     def get_node_id(self, population, node_id):
diff --git a/bmtk/simulator/bionet/biosimulator.py b/bmtk/simulator/bionet/biosimulator.py
index 9f848376e..1ec7cd65d 100644
--- a/bmtk/simulator/bionet/biosimulator.py
+++ b/bmtk/simulator/bionet/biosimulator.py
@@ -326,6 +326,7 @@ def from_config(cls, config, network, set_recordings=True):
             if sim_input.input_type == 'syn_activity':
                 network.set_spont_syn_activity(
                     precell_filter=sim_input.params['precell_filter'],
+                    postcell_filter=sim_input.params.get('postcell_filter', {}),
                     timestamps=sim_input.params['timestamps']
                 )
 
diff --git a/bmtk/simulator/bionet/cell.py b/bmtk/simulator/bionet/cell.py
index 837352cfe..261b1261a 100644
--- a/bmtk/simulator/bionet/cell.py
+++ b/bmtk/simulator/bionet/cell.py
@@ -110,5 +110,8 @@ def set_syn_connections(self, edge_prop, src_node, stim=None):
     def get_section(self, sec_name, sec_index):
         raise NotImplementedError
 
+    def __contains__(self, node_prop):
+        return node_prop in self._node
+    
     def __getitem__(self, node_prop):
         return self._node[node_prop]
diff --git a/bmtk/simulator/bionet/pointprocesscell.py b/bmtk/simulator/bionet/pointprocesscell.py
index e0b6a7b13..4cf38b254 100644
--- a/bmtk/simulator/bionet/pointprocesscell.py
+++ b/bmtk/simulator/bionet/pointprocesscell.py
@@ -132,7 +132,8 @@ def __init__(self, node, population_name, bionetwork):
         self._vecstim = h.VecStim()
         self._vecstim.play(self._spike_trains)
 
-        self._precell_filter = bionetwork.spont_syns_filter
+        self._precell_filter = bionetwork.spont_syns_filter_pre
+        self._postcell_filter = bionetwork.spont_syns_filter_post
         assert(isinstance(self._precell_filter, dict))
 
     def _matches_filter(self, src_node):
@@ -143,6 +144,22 @@ def _matches_filter(self, src_node):
             else:
                 if src_node[k] != v:
                     return False
+        
+        trg_node = trg_node or self
+        for k, v in self._postcell_filter.items():
+            # Some key may not show up as node_variable
+            if k == 'population' and k not in trg_node:
+                key_val = trg_node._node.population_name
+            else:
+                key_val = trg_node[k]
+
+            if isinstance(v, (list, tuple)):
+                if key_val not in v:
+                    return False
+            else:
+                if key_val != v:
+                    return False
+        
         return True
 
     def set_syn_connection(self, edge_prop, src_node, stim=None):

From 1f2914202f221af064d0c831810f31ee4a5c3f52 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Sat, 4 May 2024 12:52:01 -0700
Subject: [PATCH 06/11] adding ability in bionet and pointnet to handle custom
 spikes-train functions inputs

---
 bmtk/simulator/bionet/__init__.py             |  3 +-
 bmtk/simulator/bionet/bionetwork.py           | 10 ++---
 bmtk/simulator/bionet/biosimulator.py         | 21 ++++++++-
 .../bionet/modules/ecephys_module.py          |  2 +-
 bmtk/simulator/bionet/virtualcell.py          | 24 +++++++++-
 bmtk/simulator/core/node_sets.py              |  5 +++
 bmtk/simulator/core/pyfunction_cache.py       | 44 +++++++++++++++++++
 bmtk/simulator/pointnet/__init__.py           |  2 +-
 .../pointnet/modules/spikes_inputs.py         | 30 ++++++++++---
 bmtk/simulator/pointnet/pointsimulator.py     |  2 +-
 10 files changed, 125 insertions(+), 18 deletions(-)

diff --git a/bmtk/simulator/bionet/__init__.py b/bmtk/simulator/bionet/__init__.py
index 565de735a..0e24a44fb 100644
--- a/bmtk/simulator/bionet/__init__.py
+++ b/bmtk/simulator/bionet/__init__.py
@@ -20,7 +20,8 @@
 # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #
-from bmtk.simulator.bionet.pyfunction_cache import synapse_model, synaptic_weight, cell_model, add_weight_function, model_processing
+from bmtk.simulator.bionet.pyfunction_cache import synapse_model, synaptic_weight, cell_model, add_weight_function, model_processing, \
+    spikes_generator
 from bmtk.simulator.bionet.config import Config
 from bmtk.simulator.bionet.bionetwork import BioNetwork
 from bmtk.simulator.bionet.biosimulator import BioSimulator
diff --git a/bmtk/simulator/bionet/bionetwork.py b/bmtk/simulator/bionet/bionetwork.py
index 522b6b4da..8dea68954 100644
--- a/bmtk/simulator/bionet/bionetwork.py
+++ b/bmtk/simulator/bionet/bionetwork.py
@@ -136,12 +136,12 @@ def add_nodes(self, node_population):
         self._gid_pool.add_pool(node_population.name, node_population.n_nodes())
         super(BioNetwork, self).add_nodes(node_population)
 
-    def get_virtual_cells(self, population, node_id, spike_trains):
+    def get_virtual_cells(self, population, node_id, spike_trains, spikes_generator=None, sim=None):
         if node_id in self._virtual_nodes[population]:
             return self._virtual_nodes[population][node_id]
         else:
             node = self.get_node_id(population, node_id)
-            virt_cell = VirtualCell(node, population, spike_trains)
+            virt_cell = VirtualCell(node, population, spike_trains, spikes_generator, sim)
             self._virtual_nodes[population][node_id] = virt_cell
             return virt_cell
 
@@ -153,7 +153,7 @@ def get_disconnected_cell(self, population, node_id, spike_trains):
             virt_cell = self._disconnected_source_cells[population][node_id]
         else:
             node = self.get_node_id(population, node_id)
-            virt_cell = VirtualCell(node, population, spike_trains)
+            virt_cell = VirtualCell(node, population, spike_trains, self)
             self._disconnected_source_cells[population][node_id] = virt_cell
 
         return virt_cell
@@ -371,7 +371,7 @@ def find_edges(self, source_nodes=None, target_nodes=None):
 
         return selected_edges
 
-    def add_spike_trains(self, spike_trains, node_set):
+    def add_spike_trains(self, spike_trains, node_set, spikes_generator=None, sim=None):
         self._init_connections()
 
         src_nodes = [node_pop for node_pop in self.node_populations if node_pop.name in node_set.population_names()]
@@ -381,7 +381,7 @@ def add_spike_trains(self, spike_trains, node_set):
                 if edge_pop.virtual_connections:
                     for trg_nid, trg_cell in self._rank_node_ids[edge_pop.target_nodes].items():
                         for edge in edge_pop.get_target(trg_nid):
-                            src_cell = self.get_virtual_cells(source_population, edge.source_node_id, spike_trains)
+                            src_cell = self.get_virtual_cells(source_population, edge.source_node_id, spike_trains, spikes_generator, sim)
                             trg_cell.set_syn_connection(edge, src_cell, src_cell)
 
                 elif edge_pop.mixed_connections:
diff --git a/bmtk/simulator/bionet/biosimulator.py b/bmtk/simulator/bionet/biosimulator.py
index 1ec7cd65d..06f8bb4cd 100644
--- a/bmtk/simulator/bionet/biosimulator.py
+++ b/bmtk/simulator/bionet/biosimulator.py
@@ -347,13 +347,30 @@ def from_config(cls, config, network, set_recordings=True):
 
         # TODO: Need to create a gid selector
         for sim_input in inputs.from_config(config):
-            if sim_input.input_type == 'spikes' and sim_input.module in ['nwb', 'csv', 'sonata']:
+            if sim_input.input_type == 'spikes' and sim_input.module in ['nwb', 'csv', 'sonata', 'h5']:
                 io.log_info('Building virtual cell stimulations for {}'.format(sim_input.name))
                 path = sim_input.params['input_file']
                 spikes = SpikeTrains.load(path=path, file_type=sim_input.module, **sim_input.params)
                 # node_set_opts = sim_input.params.get('node_set', 'all')
                 node_set = network.get_node_set(sim_input.node_set)
-                network.add_spike_trains(spikes, node_set)
+                network.add_spike_trains(
+                    spike_trains=spikes, 
+                    node_set=node_set,
+                    spikes_generator=None,
+                    sim=sim
+                )
+
+            elif sim_input.input_type == 'spikes' and sim_input.module == 'function':
+                io.log_info('Building virtual cell stimulations for {}'.format(sim_input.name))
+                # path = sim_input.params.get['input_file']
+                spikes_generator = sim_input.params['spikes_function']
+                node_set = network.get_node_set(sim_input.node_set)
+                network.add_spike_trains(
+                    spike_trains=None, 
+                    node_set=node_set, 
+                    spikes_generator=spikes_generator,
+                    sim=sim
+                )
 
             elif sim_input.module == 'IClamp':
                 sim.add_mod(mods.IClampMod(input_type=sim_input.input_type, **sim_input.params))
diff --git a/bmtk/simulator/bionet/modules/ecephys_module.py b/bmtk/simulator/bionet/modules/ecephys_module.py
index 1b5e5d3a9..30057e3f2 100644
--- a/bmtk/simulator/bionet/modules/ecephys_module.py
+++ b/bmtk/simulator/bionet/modules/ecephys_module.py
@@ -23,7 +23,7 @@ def initialize(self, sim):
                             source_node_id = edge.source_node_id
                             spike_trains = self._mapping_strategy.get_spike_trains(source_node_id, source_population)
 
-                            src_cell = net.get_virtual_cells(source_population, source_node_id, spike_trains)
+                            src_cell = net.get_virtual_cells(source_population, source_node_id, spike_trains, sim)
                             trg_cell.set_syn_connection(edge, src_cell, src_cell)
 
                 elif edge_pop.mixed_connections:
diff --git a/bmtk/simulator/bionet/virtualcell.py b/bmtk/simulator/bionet/virtualcell.py
index 6350f4bbb..cb9bf266a 100644
--- a/bmtk/simulator/bionet/virtualcell.py
+++ b/bmtk/simulator/bionet/virtualcell.py
@@ -25,13 +25,14 @@
 import pandas as pd
 
 from bmtk.simulator.bionet.io_tools import io
+from bmtk.simulator.bionet.pyfunction_cache import py_modules
 from bmtk.utils.reports.spike_trains.spike_trains import SpikeTrains
 
 
 class VirtualCell(object):
     """Representation of a Virtual/External node"""
 
-    def __init__(self, node, population, spike_train_dataset):
+    def __init__(self, node, population, spike_train_dataset, spikes_generator, sim=None):
         # VirtualCell is currently not a subclass of bionet.Cell class b/c the parent has a bunch of properties that
         # just don't apply to a virtual cell. May want to make bionet.Cell more generic in the future.
         self._node_id = node.node_id
@@ -40,7 +41,15 @@ def __init__(self, node, population, spike_train_dataset):
         self._hobj = None
         self._spike_train_dataset = spike_train_dataset
         self._train_vec = []
-        self.set_stim(node, self._spike_train_dataset)
+        self._sim = sim
+        
+        if spike_train_dataset is not None:
+            self.set_stim(node, self._spike_train_dataset)
+        elif spikes_generator is not None:
+            self.set_stim_from_generator(node, spikes_generator)
+        else:
+            io.log_exception('Could not find source of spikes-trains (eg file or generator function)'
+                            f' for virtual cell #{self._node_id} from {self._population}')
         
     @property
     def node_id(self):
@@ -77,5 +86,16 @@ def set_stim(self, stim_prop, spike_train):
         vecstim.play(self._train_vec)
         self._hobj = vecstim
 
+    def set_stim_from_generator(self, node, spikes_generator):
+        if spikes_generator not in py_modules.spikes_generators:
+            io.log_exception(f'Could not find @spikes_generator function "{spikes_generator}". Unable to load spikes for virtual cell {self._node_id}')
+        spikes_func = py_modules.spikes_generator(name=spikes_generator)
+        spikes = spikes_func(self._node, self._sim)
+        
+        self._train_vec = h.Vector(spikes)
+        vecstim = h.VecStim()
+        vecstim.play(self._train_vec)
+        self._hobj = vecstim
+
     def __getitem__(self, item):
         return self._node[item]
diff --git a/bmtk/simulator/core/node_sets.py b/bmtk/simulator/core/node_sets.py
index 6ae65c35d..2c6140aa2 100644
--- a/bmtk/simulator/core/node_sets.py
+++ b/bmtk/simulator/core/node_sets.py
@@ -52,6 +52,11 @@ def node_ids(self):
     def nodes(self):
         return None
 
+    def fetch_nodes(self):
+        for pop in self._populations:
+            for node in pop.filter(self._filter):
+                yield node
+
 
 class NodeSetAll(NodeSet):
     def __init__(self, network):
diff --git a/bmtk/simulator/core/pyfunction_cache.py b/bmtk/simulator/core/pyfunction_cache.py
index 2af2b5c61..659230634 100644
--- a/bmtk/simulator/core/pyfunction_cache.py
+++ b/bmtk/simulator/core/pyfunction_cache.py
@@ -70,12 +70,14 @@ def __init__(self):
         self.__cell_models = {}
         self.__synapse_models = {}
         self.__cell_processors = {}
+        self.__spikes_generators = {}
 
     def clear(self):
         self.__syn_weights.clear()
         self.__cell_models.clear()
         self.__synapse_models.clear()
         self.__cell_processors.clear()
+        self.__spikes_generators.clear()
 
     def add_synaptic_weight(self, name, func, overwrite=True):
         """stores synaptic function for given name"""
@@ -156,6 +158,17 @@ def cell_processor(self, name):
     def add_cell_processor(self, name, func, overwrite=True):
         if overwrite or name not in self.__syn_weights:
             self.__cell_processors[name] = func
+   
+    @property
+    def spikes_generators(self):
+        return self.__spikes_generators.keys()
+    
+    def spikes_generator(self, name):
+        return self.__spikes_generators[name]
+    
+    def add_spikes_generator(self, name, func, overwrite=True):
+        if overwrite or name not in self.__spikes_generators:
+            self.__spikes_generators[name] = func
 
     def __repr__(self):
         rstr = '{}: {}\n'.format('cell_models', self.cell_models)
@@ -301,6 +314,31 @@ def func_wrapper(*args, **kwargs):
         return decorator
 
 
+def spikes_generator(*wargs, **wkwargs):
+    if len(wargs) == 1 and callable(wargs[0]):
+        # for the case without decorator arguments, grab the function object in wargs and create a decorator
+        func = wargs[0]
+        py_modules.add_spikes_generator(func.__name__, func)  # add function assigned to its original name
+
+        @wraps(func)
+        def func_wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+        return func_wrapper
+    else:
+        # for the case with decorator arguments
+        assert(all(k in ['name'] for k in wkwargs.keys()))
+
+        def decorator(func):
+            # store the function in py_modules but under the name given in the decorator arguments
+            py_modules.add_spikes_generator(wkwargs['name'], func)
+
+            @wraps(func)
+            def func_wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return func_wrapper
+        return decorator
+
+
 def add_weight_function(func, name=None, overwrite=True):
     assert(callable(func))
     func_name = name if name is not None else func.__name__
@@ -324,6 +362,12 @@ def add_synapse_model(func, name=None, overwrite=True):
     py_modules.add_synapse_model(func_name, func, overwrite)
 
 
+def add_spikes_generator(func, name=None, overwrite=True):
+    assert(callable(func))
+    func_name = name if name is not None else func.__name__
+    py_modules.add_spikes_generator(func_name, func, overwrite)
+
+
 def load_py_modules(cell_models=None, syn_models=None, syn_weights=None, cell_processors=None):
     # py_modules.clear()
     warnings.warn('Do not call this method directly', DeprecationWarning)
diff --git a/bmtk/simulator/pointnet/__init__.py b/bmtk/simulator/pointnet/__init__.py
index e66f20797..fe9e2429e 100644
--- a/bmtk/simulator/pointnet/__init__.py
+++ b/bmtk/simulator/pointnet/__init__.py
@@ -24,4 +24,4 @@
 from .config import Config
 from .pointnetwork import PointNetwork
 from .pointsimulator import PointSimulator
-from .pyfunction_cache import synapse_model, synaptic_weight, cell_model
+from .pyfunction_cache import synapse_model, synaptic_weight, cell_model, spikes_generator
diff --git a/bmtk/simulator/pointnet/modules/spikes_inputs.py b/bmtk/simulator/pointnet/modules/spikes_inputs.py
index eb96dcfe4..1e305fc25 100644
--- a/bmtk/simulator/pointnet/modules/spikes_inputs.py
+++ b/bmtk/simulator/pointnet/modules/spikes_inputs.py
@@ -24,6 +24,7 @@
 from bmtk.simulator.pointnet.modules.sim_module import SimulatorMod
 from bmtk.simulator.pointnet.io_tools import io
 from bmtk.utils.reports.spike_trains import SpikeTrains
+from bmtk.simulator.bionet.pyfunction_cache import py_modules
 
 
 class SpikesInputsMod(SimulatorMod):
@@ -38,9 +39,28 @@ def initialize(self, sim):
         io.log_info('Build virtual cell stimulations for {}'.format(self._name))
         
         node_set = sim.net.get_node_set(self._params['node_set'])
-        self._spike_trains = SpikeTrains.load(
-            path=self._params['input_file'], 
-            file_type=self._module, 
-            **self._params
-        )
+       
+        if self._module == 'function':
+            if 'spikes_function' not in self._params:
+                io.log_exception('missing parameter "spikes_function" for input {self._name}, module {self._module}')
+            spikes_generator = self._params['spikes_function']
+            if spikes_generator not in py_modules.spikes_generators:
+                io.log_exception(f'Could not find @spikes_generator function "{spikes_generator}" required for {self._name} inputs.')
+            spikes_func = py_modules.spikes_generator(name=spikes_generator)
+            
+            self._spike_trains = SpikeTrains(cace_to_disk=False)
+            for node in node_set.fetch_nodes():
+                timestamps = spikes_func(node, sim)
+                self._spike_trains.add_spikes(
+                    node_ids=node.node_id, 
+                    timestamps=timestamps, 
+                    population=node.population_name
+                )
+        else:
+            self._spike_trains = SpikeTrains.load(
+                path=self._params['input_file'], 
+                file_type=self._module, 
+                **self._params
+            )
+
         sim.net.add_spike_trains(self._spike_trains, node_set, sim.get_spike_generator_params())
diff --git a/bmtk/simulator/pointnet/pointsimulator.py b/bmtk/simulator/pointnet/pointsimulator.py
index 76482df89..b6a4af0ec 100644
--- a/bmtk/simulator/pointnet/pointsimulator.py
+++ b/bmtk/simulator/pointnet/pointsimulator.py
@@ -237,7 +237,7 @@ def from_config(cls, configure, graph, n_thread=None):
                 os.remove(gfile)
 
         for sim_input in inputs.from_config(config):
-            if sim_input.input_type == 'spikes' and sim_input.module in ['nwb', 'csv', 'sonata', 'h5', 'hdf5']:
+            if sim_input.input_type == 'spikes' and sim_input.module in ['nwb', 'csv', 'sonata', 'h5', 'hdf5', 'function']:
                 network.add_mod(mods.SpikesInputsMod(
                     name=sim_input.name,
                     input_type=sim_input.input_type,

From ad2b4663cae6433b11ef45b583c20417dd4fddec Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Sun, 5 May 2024 13:20:16 -0700
Subject: [PATCH 07/11] bionet iclamps can now support ranged section selection

---
 bmtk/simulator/bionet/modules/iclamp.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/bmtk/simulator/bionet/modules/iclamp.py b/bmtk/simulator/bionet/modules/iclamp.py
index c500715c5..4774e86c3 100644
--- a/bmtk/simulator/bionet/modules/iclamp.py
+++ b/bmtk/simulator/bionet/modules/iclamp.py
@@ -123,6 +123,9 @@ def __init__(self, input_type, **mod_args):
         self._section_index = mod_args.get('section_index', 0)
         self._section_dist = mod_args.get('section_dist', 0.5)
 
+        # Check f section_index is a range (ie. "section_index": [500.0, 1000.0])
+        self._ranged_index = isinstance(self._section_index, (list, tuple))
+
         # IClamp objects need to be saved in memory otherwise NRN will try to garbage collect them
         # prematurly
         self._iclamps = []
@@ -132,7 +135,23 @@ def initialize(self, sim):
         select_gids = list(sim.net.get_node_set(self._node_set).gids())
         gids_on_rank = list(set(select_gids) & set(select_gids))
 
-        for gid in gids_on_rank:
+        for gid in gids_on_rank:           
             cell = sim.net.get_cell_gid(gid)
-            hobj_sec = getattr(cell.hobj, self._section_name)[self._section_index](self._section_dist)
+            
+            if self._ranged_index:
+                hobj_sec = self._find_section(cell)
+            else:
+                hobj_sec = getattr(cell.hobj, self._section_name)[self._section_index](self._section_dist)
+            
             self._iclamps.extend(self._amp_reader.create_clamps(hobj_sec))
+
+
+    def _find_section(self, cell):
+        secs_list = self._section_name if isinstance(self._section_name, (list, tuple)) else [self._section_name]        
+        seg_idx, _ = cell.morphology.choose_sections(
+            section_names=secs_list,
+            distance_range=self._section_index,
+            n_sections=1
+        )
+
+        return cell.morphology.segments[seg_idx[0]]

From a6aac4311e94e5cd81a7a8dd293094e33b6d766a Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Sun, 5 May 2024 15:07:02 -0700
Subject: [PATCH 08/11] Adding ctdb download utils

---
 bmtk/utils/cell_types_db/__init__.py | 76 ++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)
 create mode 100644 bmtk/utils/cell_types_db/__init__.py

diff --git a/bmtk/utils/cell_types_db/__init__.py b/bmtk/utils/cell_types_db/__init__.py
new file mode 100644
index 000000000..8d613aa90
--- /dev/null
+++ b/bmtk/utils/cell_types_db/__init__.py
@@ -0,0 +1,76 @@
+from pathlib import Path
+import os
+import io
+
+import xml.etree.ElementTree as ET 
+import requests
+import zipfile
+
+
+def download_model(specimen_id, model_type='Biophysical - all active', base_dir='models', overwrite=False):
+    """Uses the Allen REST API to find and download the model files for a given cell. All relevant files (parameters json and morphology swc)
+    will be saved in a separate folder, the path to which will be returned.
+
+    Notes: Using the REST API because at the moment I don't think the AllenSDK is capable of downloading parameters_fit.json files, see 
+    https://community.brain-map.org/t/cell-types-database-api/3016 for more info on how to use the API.
+    """
+    # Set a request to get fetch model data available for a given speciment_id. It will return a xml string that needs to be parsed so
+    # that we can find the correct model_id.
+    api_url = f"http://api.brain-map.org/api/v2/data/query.xml?criteria=model::NeuronalModel,rma::critera,[specimen_id$eq{specimen_id}]"
+    response = requests.get(api_url)
+    xml_root = ET.fromstring(response.content)
+    model_id = None
+    for (model_name, model_id) in zip(xml_root.iter('name'), xml_root.iter('id')):
+        if 'Biophysical - all active' in model_name.text:
+            model_id = int(model_id.text)
+            break
+
+    if model_id is None:
+        raise ValueError(f'Could not find a "{model_type}" model for cell {specimen_id}')
+    
+    # Now that we have the model_id for the given cell we can download and unzip the files into the correct directory. To prevent downloading
+    # the zip everytime we'll check to see if the directory already exists.
+    model_dir = Path(f'{base_dir}/neuronal_model_{model_id}')
+    if model_dir.exists() and not overwrite:
+        print(f'> {model_dir} already exits, skipping donwloadng data')
+        return model_dir
+
+    zip_uri = f'http://api.brain-map.org/neuronal_model/download/{model_id}'
+    zip_response = requests.get(zip_uri)
+    zip_file = zipfile.ZipFile(io.BytesIO(zip_response.content))
+    zip_file.extractall(model_dir)
+    return model_dir
+
+
+def download_ephys_data(specimen_id, download_dir='ephys_inputs'):
+    """Download nwb file containing sweeps."""
+    api_url = f'http://api.brain-map.org/api/v2/data/query.xml?criteria=model::Specimen,rma::criteria,[id$eq{specimen_id}],rma::include,ephys_result(well_known_files(well_known_file_type[name$eqNWBDownload]))'
+    response = requests.get(api_url)
+    # print(response.content)
+    download_uri = None
+    xml_root = ET.fromstring(response.content)
+    for dl_link in xml_root.iter('download-link'):
+        download_uri = dl_link.text
+        break
+
+    ephys_id = None
+    for erid in xml_root.iter('ephys-result-id'):
+        ephys_id = erid.text
+        break
+
+    nwb_path = Path(f'{download_dir}/{ephys_id}_ephys.nwb')
+    if nwb_path.exists():
+        print(f'> {nwb_path} already exits, skipping donwload.')
+        return
+
+    if not Path(download_dir).exists():
+        os.makedirs(download_dir)
+
+    url_req = f'https://celltypes.brain-map.org/{download_uri}'
+    nwb_req = requests.get(url_req, stream=True)
+    with open(nwb_path, 'wb') as fh:
+        for chunk in nwb_req.iter_content():
+            fh.write(chunk)
+
+def test_fn():
+    print("HERE")
\ No newline at end of file

From 3973dc6e32c8514a5c6f1f66a59c3f3c5d0daa64 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Mon, 6 May 2024 17:38:20 -0700
Subject: [PATCH 09/11] Updating ECP Plotting function

---
 bmtk/analyzer/ecp.py                | 50 ++++++++++++++++++++++-------
 bmtk/simulator/bionet/morphology.py |  2 +-
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/bmtk/analyzer/ecp.py b/bmtk/analyzer/ecp.py
index 1179d0195..0c5835a41 100644
--- a/bmtk/analyzer/ecp.py
+++ b/bmtk/analyzer/ecp.py
@@ -2,10 +2,12 @@
 import h5py
 import matplotlib.pyplot as plt
 import numpy as np
+from decimal import Decimal
 
 from bmtk.utils.sonata.config import SonataConfig
 from bmtk.simulator.utils import simulation_reports
-# from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
+from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
+import matplotlib.font_manager as fm
 
 
 def _get_ecp_path(ecp_path=None, config=None, report_name=None):
@@ -55,30 +57,54 @@ def plot_ecp(config_file=None, report_name=None, ecp_path=None, title=None, show
     channels = ecp_h5['/ecp/channel_id'][()]
     fig, axes = plt.subplots(len(channels), 1)
     fig.text(0.04, 0.5, 'channel id', va='center', rotation='vertical')
+    v_min, v_max = ecp_h5['/ecp/data'][()].min(), ecp_h5['/ecp/data'][()].max()
+    # print(v_max - v_min)
+    # exit()
+    
     for idx, channel in enumerate(channels):
         data = ecp_h5['/ecp/data'][:, idx]
+        # print(channel, np.min(data), np.max(data))
         axes[idx].plot(time_traces, data)
         axes[idx].spines["top"].set_visible(False)
         axes[idx].spines["right"].set_visible(False)
         axes[idx].set_yticks([])
         axes[idx].set_ylabel(channel)
+        axes[idx].set_ylim([v_min, v_max])
 
         if idx+1 != len(channels):
             axes[idx].spines["bottom"].set_visible(False)
             axes[idx].set_xticks([])
         else:
             axes[idx].set_xlabel('timestamps (ms)')
-            # scalebar = AnchoredSizeBar(axes[idx].transData,
-            #                            2.0, '1 mV', 1,
-            #                            pad=0,
-            #                            borderpad=0,
-            #                            # color='b',
-            #                            frameon=True,
-            #                            # size_vertical=1.001,
-            #                            # fontproperties=fontprops
-            #                            )
-            #
-            # axes[idx].add_artist(scalebar)
+        
+        
+        if idx == 0:
+            scale_bar_size = (v_max-v_min)/2.0
+            scale_bar_label = f'{scale_bar_size:.2E}'
+            # print(scale_bar_label)
+            # exit()
+            fontprops = fm.FontProperties(size='x-small')
+
+            scalebar = AnchoredSizeBar(
+                axes[idx].transData,
+                size=scale_bar_size, 
+                label=scale_bar_label, 
+                loc='upper right',
+                pad=0.1,
+                borderpad=0.5,
+                sep=5,
+                # color='b',
+                frameon=False,
+                size_vertical=scale_bar_size,
+                # size_vertical=1.001,
+                fontproperties=fontprops
+            )            
+            axes[idx].add_artist(scalebar)
+
+            # label = scalebar.txt_label
+            # label.set_rotation(270.0)
+            # label.set_verticalalignment('bottom')
+            # label.set_horizontalalignment('left') 
 
     if title:
         fig.set_title(title)
diff --git a/bmtk/simulator/bionet/morphology.py b/bmtk/simulator/bionet/morphology.py
index 00a29f99c..2413bde8a 100644
--- a/bmtk/simulator/bionet/morphology.py
+++ b/bmtk/simulator/bionet/morphology.py
@@ -454,7 +454,7 @@ def move_and_rotate(self, soma_coords=None, rotation_angles=None, inplace=False)
             new_soma_pos = soma_coords
 
         # new_seg_coords = SegmentCoords(p0=new_p0, p1=new_p1, p05=new_p05, soma_pos=new_soma_pos)
-        new_seg_coords = _LazySegmentCoords(self.hobj)
+        new_seg_coords = _LazySegmentCoords(self)
         new_seg_coords.p0 = new_p0
         new_seg_coords.p05 = new_p05
         new_seg_coords.p1 = new_p1

From 20044d77778284e050d6132af4711ab9a1b17f37 Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Wed, 22 May 2024 13:26:39 -0700
Subject: [PATCH 10/11] updating way templates can be loaded to fix an ordering
 bug

---
 bmtk/simulator/bionet/nrn.py | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/bmtk/simulator/bionet/nrn.py b/bmtk/simulator/bionet/nrn.py
index cc97a95f4..1ee8e1978 100644
--- a/bmtk/simulator/bionet/nrn.py
+++ b/bmtk/simulator/bionet/nrn.py
@@ -23,6 +23,7 @@
 import sys
 import os
 import glob
+from pathlib import Path
 import neuron
 from neuron import h
 
@@ -88,17 +89,25 @@ def load_neuron_modules(mechanisms_dir, templates_dir, default_templates=True, u
         load_templates(templates_dir)
 
 
-def load_templates(template_dir):
+def load_templates(templates):
     """Load all templates to be available in the hoc namespace for instantiating cells"""
     cwd = os.getcwd()
-    os.chdir(template_dir)
-
-    hoc_templates = glob.glob("*.hoc")
-
-    for hoc_template in hoc_templates:
-        h.load_file(str(hoc_template))
-
-    os.chdir(cwd)
+    templates_list = templates if isinstance(templates, (list, tuple)) else [templates]
+    for templates_cont in templates_list:
+        if Path(templates_cont).is_dir():
+            # If string is a path to a directory, find all hoc files and load them in one at a time.
+            # TODO: Add option to sort before loading
+            os.chdir(templates_cont)
+            hoc_templates = glob.glob("*.hoc")
+            for hoc_template in hoc_templates:
+                h.load_file(str(hoc_template))
+            os.chdir(cwd)
+            
+        elif Path(templates_cont).is_file():
+            # Otherwise if just a file, eg .hoc, load it in separartly.
+            h.load_file(str(templates_cont))
+        else:
+            io.log_warning(f'Unable to find and load {templates_cont} templates. Please ensure it is a valid directory or hoc file!', display_once=True)
 
 
 def reset():

From cd294026d1c1a92b6b3d68f4e0f9dea334d6784c Mon Sep 17 00:00:00 2001
From: kaeldai <kael.dai@gmail.com>
Date: Fri, 14 Jun 2024 13:18:51 -0700
Subject: [PATCH 11/11] fixing issues with unit tests

---
 bmtk/simulator/bionet/virtualcell.py           |  2 +-
 bmtk/tests/simulator/bionet/test_morphology.py | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/bmtk/simulator/bionet/virtualcell.py b/bmtk/simulator/bionet/virtualcell.py
index cb9bf266a..c8e28e5f0 100644
--- a/bmtk/simulator/bionet/virtualcell.py
+++ b/bmtk/simulator/bionet/virtualcell.py
@@ -32,7 +32,7 @@
 class VirtualCell(object):
     """Representation of a Virtual/External node"""
 
-    def __init__(self, node, population, spike_train_dataset, spikes_generator, sim=None):
+    def __init__(self, node, population, spike_train_dataset, spikes_generator=None, sim=None):
         # VirtualCell is currently not a subclass of bionet.Cell class b/c the parent has a bunch of properties that
         # just don't apply to a virtual cell. May want to make bionet.Cell more generic in the future.
         self._node_id = node.node_id
diff --git a/bmtk/tests/simulator/bionet/test_morphology.py b/bmtk/tests/simulator/bionet/test_morphology.py
index 453516a24..30f8f6dec 100644
--- a/bmtk/tests/simulator/bionet/test_morphology.py
+++ b/bmtk/tests/simulator/bionet/test_morphology.py
@@ -56,7 +56,7 @@ def test_base():
     assert(len(morph.seg_props.dist) == 87)
     assert(len(morph.seg_props.length) == 87)
     assert(morph.seg_props.type[0] == 1 and morph.seg_props.type[1] == 2)
-    assert(np.isclose(morph.seg_props.area[0], 425.74, atol=1.0e-2))
+    assert(np.isclose(morph.seg_props.area[0], 425.756, atol=1.0e-2))
     assert(morph.seg_props.x[0] == 0.5)
     assert(np.isclose(morph.seg_props.dist[0], 5.82, atol=1.0e-2))
     assert(np.isclose(morph.seg_props.length[0], 11.64, atol=1.0e-2))
@@ -69,7 +69,7 @@ def test_base():
     assert(np.allclose(morph.seg_coords.p1[:, 0], [5.82, 0.0, 0.0], atol=1.0e-2))
     assert(np.allclose(morph.seg_coords.p0[:, 86], [-22.21, -18.18, -1.89], atol=1.0e-2))
     assert(np.allclose(morph.seg_coords.p05[:, 86], [-38.98, -12.29, -4.19], atol=1.0e-2))
-    assert(np.allclose(morph.seg_coords.p1[:, 86], [-51.45, 0.73, -6.35], atol=1.0e-2))
+    assert(np.allclose(morph.seg_coords.p1[:, 86], [-51.45, 0.740, -6.35], atol=1.0e-2))
 
 
 @pytest.mark.skipif(not has_mechanism, reason='Mechanisms has not been compiled, run nrnivmodl mechanisms.')
@@ -119,7 +119,7 @@ def test_full():
     assert(len(morph.seg_props.dist) == 206)
     assert(len(morph.seg_props.length) == 206)
     assert(morph.seg_props.type[0] == 1 and morph.seg_props.type[1] == 3 and morph.seg_props.type[-1] == 2)
-    assert(np.isclose(morph.seg_props.area[0], 425.74, atol=1.0e-2))
+    assert(np.isclose(morph.seg_props.area[0], 425.756, atol=1.0e-2))
     assert(morph.seg_props.x[0] == 0.5)
     assert(np.isclose(morph.seg_props.dist[0], 5.82, atol=1.0e-2))
     assert(np.isclose(morph.seg_props.length[0], 11.64, atol=1.0e-2))
@@ -158,7 +158,7 @@ def test_multicell():
     assert(len(morph1.seg_props.dist) == 206)
     assert(len(morph1.seg_props.length) == 206)
     assert(morph1.seg_props.type[0] == 1 and morph1.seg_props.type[1] == 3 and morph1.seg_props.type[-1] == 2)
-    assert(np.isclose(morph1.seg_props.area[0], 425.74, atol=1.0e-2))
+    assert(np.isclose(morph1.seg_props.area[0], 425.756, atol=1.0e-2))
     assert(morph1.seg_props.x[0] == 0.5)
     assert(np.isclose(morph1.seg_props.dist[0], 5.82, atol=1.0e-2))
     assert(np.isclose(morph1.seg_props.length[0], 11.64, atol=1.0e-2))
@@ -367,7 +367,7 @@ def test_move_and_rotate():
     assert(np.allclose(morph1.seg_coords.p1[:, 0], [5.82, 0.0, 0.0], atol=1.0e-2))
     assert(np.allclose(morph1.seg_coords.p0[:, 86], [-22.21, -18.18, -1.89], atol=1.0e-2))
     assert(np.allclose(morph1.seg_coords.p05[:, 86], [-38.98, -12.29, -4.19], atol=1.0e-2))
-    assert(np.allclose(morph1.seg_coords.p1[:, 86], [-51.45, 0.73, -6.35], atol=1.0e-2))
+    assert(np.allclose(morph1.seg_coords.p1[:, 86], [-51.45, 0.740, -6.35], atol=1.0e-2))
 
     assert(np.allclose(morph2.seg_coords.p0[:, 0], [103.48, -95.83, 2.08], atol=1.0e-2))
     assert(np.allclose(morph2.seg_coords.p05[:, 0], [100.0, -100.0, 0.0], atol=1.0e-2))