Feat: Add major_termination_length parameter (#78)

BlueBrain · Sep 19, 2023 · ed9f4f0 · ed9f4f0
1 parent effca74
commit ed9f4f0
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/neurots/generate/tree.py b/neurots/generate/tree.py
@@ -248,8 +248,21 @@ def next_point(self):
         for section_grower in ordered_sections:
             # the current section_grower is generated
             # In here the stop criterion can be modified accordingly
+
+            if section_grower.process == "major" and "major_termination_length" in self.params:
+                # this makes an early termination of major branch
+                # it needs to revert the value after getting the state to preserve
+                # the original topology
+                _term = section_grower.stop_criteria["TMD"].term
+                section_grower.stop_criteria["TMD"].term = self.params["major_termination_length"]
+            else:
+                _term = None
+
             state = self.growth_algo.extend(section_grower)
 
+            if _term is not None:
+                section_grower.stop_criteria["TMD"].term = _term
+
             if state != "continue":
                 section = self.append_section(section_grower)
 

diff --git a/neurots/schemas/parameters.json b/neurots/schemas/parameters.json
@@ -45,6 +45,12 @@
                     ],
                     "type": "string"
                 },
+                "major_termination_length": {
+                    "description": "Termination length of major branches, to use mostly for main axon to connect to long-range axon",
+                    "type": "number",
+                    "minimum": 0
+                },
+
                 "has_apical_tuft": {
                     "description": "Select True if the tree is an apical and is expected to have a tuft",
                     "type": "boolean"

diff --git a/tests/test_generate_tree.py b/tests/test_generate_tree.py
@@ -88,6 +88,30 @@ def test_TreeGrower():
         )
 
 
+def test_TreeGrower_termination_length():
+    np.random.seed(0)
+    with open(os.path.join(_path, "bio_distribution.json"), encoding="utf-8") as f:
+        distributions = json.load(f)
+
+    with open(os.path.join(_path, "bio_path_params.json"), encoding="utf-8") as f:
+        params = json.load(f)
+
+    params["apical_dendrite"]["major_termination_length"] = 200
+    grower = NeuronGrower(input_distributions=distributions, input_parameters=params)
+    grower.grow()
+
+    # Test order_per_process()
+    sections = [i.active_sections[0] for i in grower.active_neurites]
+    for num, i in enumerate(sections):
+        i.process = str(len(sections) - num - 1)
+    res = TreeGrower.order_per_process(sections)
+    for num, i in enumerate(res):
+        assert i == sections[len(sections) - num - 1], (
+            i,
+            sections[len(sections) - num - 1],
+        )
+
+
 def test_TreeGrower_3d_angles():
     np.random.seed(0)
     with open(os.path.join(_path, "bio_distribution_3d_angles.json"), encoding="utf-8") as f: