Feat: Fix default diam distr (#83)

BlueBrain · Sep 13, 2023 · 39d6026 · 39d6026
1 parent 4249ee9
commit 39d6026
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 8 deletions.
diff --git a/neurots/extract_input/input_distributions.py b/neurots/extract_input/input_distributions.py
@@ -97,11 +97,10 @@ def distributions(
     elif (
         isinstance(diameter_model, str) and diameter_model == "default"
     ) or diameter_model is None:
-        input_distributions["diameter"] = {}
-        input_distributions["diameter"]["method"] = "default"
-        input_distributions["diameter"]["diameter"] = build_diameter_models(
+        input_distributions["diameter"] = build_diameter_models(
             morphology, config={"models": ["simpler"], "neurite_types": neurite_types}
         )
+        input_distributions["diameter"]["method"] = "default"
     else:
         raise NotImplementedError(f"Diameter model {diameter_model} not understood")
 

diff --git a/neurots/extract_input/input_parameters.py b/neurots/extract_input/input_parameters.py
@@ -104,12 +104,21 @@ def merged_params(data):
             input_parameters["apical_dendrite"]["growth_method"] = "tmd_apical"
 
     input_parameters["diameter_params"] = {}
-    if diameter_parameters is None:
+    if isinstance(diameter_parameters, str):
+        input_parameters["diameter_params"]["method"] = diameter_parameters
+    elif diameter_parameters is None or (
+        isinstance(diameter_parameters, dict)
+        and list(diameter_parameters.keys()) == ["neurite_types"]
+    ):
         input_parameters["diameter_params"]["method"] = "default"
         input_parameters["diameter_params"]["models"] = ["simpler"]
-    elif isinstance(diameter_parameters, str):
-        input_parameters["diameter_params"]["method"] = diameter_parameters
-    elif isinstance(diameter_parameters, dict):
+        if diameter_parameters is not None:
+            input_parameters["diameter_params"]["neurite_types"] = diameter_parameters[
+                "neurite_types"
+            ]
+    elif isinstance(diameter_parameters, dict) and list(diameter_parameters.keys()) != [
+        "neurite_types"
+    ]:
         input_parameters["diameter_params"] = diameter_parameters
         input_parameters["diameter_params"]["method"] = "external"
     else:

diff --git a/neurots/schemas/distributions.json b/neurots/schemas/distributions.json
@@ -210,12 +210,41 @@
                     "properties": {
                         "method": {
                             "enum": [
-                                "external", "default"
+                                "external"
                             ],
                             "type": "string"
                         }
                     }
                 },
+                {
+                    "additionalProperties": false,
+                    "properties": {
+                        "method": {
+                            "enum": [
+                                "default"
+                            ],
+                            "type": "string"
+                        },
+                        "apical_dendrite": {
+                            "items": {
+                                "type": "number"
+                            },
+                            "type": "array"
+                        },
+                        "basal_dendrite": {
+                            "items": {
+                                "type": "number"
+                            },
+                            "type": "array"
+                        },
+                        "axon": {
+                            "items": {
+                                "type": "number"
+                            },
+                            "type": "array"
+                        }
+                    }
+                },
                 {
                     "additionalProperties": false,
                     "properties": {

diff --git a/tests/test_extract_input.py b/tests/test_extract_input.py
@@ -705,6 +705,13 @@ def test_parameters():
         diameter_parameters="some_diametrizer",
     )
 
+    # test default input with custom neurite_types
+    extract_input.parameters(
+        neurite_types=["axon"],
+        method="trunk",
+        diameter_parameters={"neurite_types": "axon"},
+    )
+
 
 def test_from_TMD():
     files = sorted([os.path.join(POP_PATH, neuron_dir) for neuron_dir in os.listdir(POP_PATH)])