replace_axon now allows the axon length and numbr of sections to be set

bluepyopt/ephys/morphologies.py

@@ -57,6 +57,8 @@ def __init__(
             do_set_nseg=True,
             comment='',
             replace_axon_hoc=None,
+            axon_stump_length=60,
+            axon_stump_nsec=2,
             nseg_frequency=40,
             morph_modifiers=None,
             morph_modifiers_hoc=None,
@@ -74,6 +76,8 @@ def __init__(
                 python, replace_axon is used instead. Must include
                 'proc replace_axon(){ ... }
                 If None,the default replace_axon is used
+            axon_stump_length (float): Length of replacement axon
+            axon_stump_nsec (int): Number of sections in replacement axon
             nseg_frequency (float): frequency of nseg
             do_set_nseg (bool): if True, it will use nseg_frequency
             morph_modifiers (list): list of functions to modify the icell
@@ -90,6 +94,8 @@ def __init__(
             morphology_path = str(morphology_path)
         self.morphology_path = morphology_path
         self.do_replace_axon = do_replace_axon
+        self.axon_stump_length = axon_stump_length
+        self.axon_stump_nsec = axon_stump_nsec
         self.do_set_nseg = do_set_nseg
         self.nseg_frequency = nseg_frequency
         self.morph_modifiers = morph_modifiers
@@ -155,7 +161,9 @@ def instantiate(self, sim=None, icell=None):
             self.set_nseg(icell)
 
         if self.do_replace_axon:
-            self.replace_axon(sim=sim, icell=icell)
+            self.replace_axon(sim=sim, icell=icell,
+                              axon_stump_length=self.axon_stump_length,
+                              n_sections=self.axon_stump_nsec)
 
         if self.morph_modifiers is not None:
             for morph_modifier in self.morph_modifiers:
@@ -172,43 +180,60 @@ def set_nseg(self, icell):
             section.nseg = 1 + 2 * int(section.L / self.nseg_frequency)
 
     @staticmethod
-    def replace_axon(sim=None, icell=None):
+    def replace_axon(sim=None, icell=None, axon_stump_length=60, n_sections=2):
         """Replace axon"""
 
         nsec = len([sec for sec in icell.axonal])
 
+        print(f"In replace_axon {nsec = }")
+
         if nsec == 0:
-            ais_diams = [1, 1]
+            ais_diams = [1] * n_sections
         elif nsec == 1:
-            ais_diams = [icell.axon[0].diam, icell.axon[0].diam]
+            sec_x = [min(1, x/n_sections * axon_stump_length/icell.axon[0].L)
+                     for x in range(0, n_sections)]
+            ais_diams = [icell.axon[0](sx).diam for sx in sec_x]
         else:
-            ais_diams = [icell.axon[0].diam, icell.axon[0].diam]
+            ais_diams = [icell.axon[0].diam] * n_sections
             # Define origin of distance function
             sim.neuron.h.distance(0, 0.5, sec=icell.soma[0])
 
+            end_point_dist = np.arange(1, n_sections+1) \
+                * axon_stump_length / n_sections
+
+            idx = 1
+
             for section in icell.axonal:
-                # If distance to soma is larger than 60, store diameter
-                if sim.neuron.h.distance(1, 0.5, sec=section) > 60:
-                    ais_diams[1] = section.diam
+                if idx >= len(ais_diams):
                     break
+                if sim.neuron.h.distance(1, 0.5, sec=section)\
+                        < end_point_dist[idx]:
+                    ais_diams[idx:] = section.diam
+                else:
+                    idx += 1
 
         for section in icell.axonal:
             sim.neuron.h.delete_section(sec=section)
 
         # Create new axon array
-        sim.neuron.h.execute('create axon[2]', icell)
+        sim.neuron.h.execute(f"create axon[{n_sections}]", icell)
 
         for index, section in enumerate(icell.axon):
             section.nseg = 1
-            section.L = 30
+            section.L = axon_stump_length / n_sections
             section.diam = ais_diams[index]
             icell.axonal.append(sec=section)
             icell.all.append(sec=section)
 
-        icell.axon[0].connect(icell.soma[0], 1.0, 0.0)
-        icell.axon[1].connect(icell.axon[0], 1.0, 0.0)
+        for index in range(len(icell.axon)):
+            if index == 0:
+                icell.axon[0].connect(icell.soma[0], 1.0, 0.0)
+            else:
+                icell.axon[index].connect(icell.axon[index - 1], 1.0, 0.0)
+
+        logger.debug(f"Replace axon with AIS {axon_stump_length = }, "
+                     f"{n_sections =}")
 
-        logger.debug('Replace axon with AIS')
 
     default_replace_axon_hoc = \
         '''