From 123c36a9812466dd9cc165ac9a29765def6f43d4 Mon Sep 17 00:00:00 2001
From: Adrien Berchet <adrien.berchet@epfl.ch>
Date: Fri, 23 Oct 2020 10:07:13 +0200
Subject: [PATCH] Add diametrizer in vacuum synthesis

Change-Id: I3ee112796dd081e5ab31eb79cd09f5f5cd276433
---
 examples/luigi_cfg/luigi_vacuum.cfg          |  4 +--
 synthesis_workflow/tasks/vacuum_synthesis.py |  4 +++
 synthesis_workflow/vacuum_synthesis.py       | 38 ++++++++++++++++----
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/examples/luigi_cfg/luigi_vacuum.cfg b/examples/luigi_cfg/luigi_vacuum.cfg
index d72917b..eaa413d 100644
--- a/examples/luigi_cfg/luigi_vacuum.cfg
+++ b/examples/luigi_cfg/luigi_vacuum.cfg
@@ -1,6 +1,6 @@
 # luigi parameters
-[core]
-logging_conf_file = logging.conf
+# [core]
+# logging_conf_file = logging.conf
 
 # global parameters
 [SynthesisConfig]
diff --git a/synthesis_workflow/tasks/vacuum_synthesis.py b/synthesis_workflow/tasks/vacuum_synthesis.py
index 4bd84d8..b34d59f 100644
--- a/synthesis_workflow/tasks/vacuum_synthesis.py
+++ b/synthesis_workflow/tasks/vacuum_synthesis.py
@@ -35,6 +35,9 @@ class VacuumSynthesize(WorkflowTask):
 
     vacuum_synth_morphology_path = luigi.Parameter(default="vacuum_synth_morphologies")
     vacuum_synth_morphs_df_path = luigi.Parameter(default="vacuum_synth_morphs_df.csv")
+    diametrizer = luigi.ChoiceParameter(
+        default="external", choices=["external"] + [f"M{i}" for i in range(1, 6)]
+    )
     n_cells = luigi.IntParameter(default=10)
 
     def requires(self):
@@ -66,6 +69,7 @@ def run(self):
             tmd_parameters,
             tmd_distributions,
             morphology_base_path,
+            diametrizer=self.diametrizer,
             joblib_verbose=self.joblib_verbose,
             nb_jobs=self.nb_jobs,
         )
diff --git a/synthesis_workflow/vacuum_synthesis.py b/synthesis_workflow/vacuum_synthesis.py
index 27425bd..4866e71 100644
--- a/synthesis_workflow/vacuum_synthesis.py
+++ b/synthesis_workflow/vacuum_synthesis.py
@@ -11,6 +11,7 @@
 from neurom import load_neuron
 from neurom import viewer
 from tns import NeuronGrower
+from diameter_synthesis import build_diameters
 
 from . import STR_TO_TYPES
 from .synthesis import get_max_len
@@ -18,7 +19,12 @@
 
 
 def _grow_morphology(
-    gid, mtype, tmd_parameters, tmd_distributions, morphology_base_path
+    gid,
+    mtype,
+    tmd_parameters,
+    tmd_distributions,
+    morphology_base_path,
+    external_diametrizer=None,
 ):
     """Grow single morphology for parallel computations."""
 
@@ -30,7 +36,7 @@ def _grow_morphology(
     grower = NeuronGrower(
         input_parameters=tmd_parameters,
         input_distributions=tmd_distributions,
-        external_diametrizer=None,
+        external_diametrizer=external_diametrizer,
     )
     grower.grow()
     grower.neuron.write(morphology_path)
@@ -49,17 +55,36 @@ def grow_vacuum_morphologies(
     tmd_parameters,
     tmd_distributions,
     morphology_base_path,
+    diametrizer="external",
     joblib_verbose=0,
     nb_jobs=1,
 ):
-    """Grow morphologies in vacuum."""
+    """Grow morphologies in vacuum.
+
+    With diametrizer='external', we will use diameter-synthesis,
+    otherwise 'M1-M5' from TNS are allowed.
+    """
 
     global_gid = 0
     vacuum_synth_morphs_df = pd.DataFrame()
     for mtype in tqdm(mtypes):
-        # no need to realistic diameters here, using internal TNS diametrizer
-        tmd_parameters[mtype]["diameter_params"]["method"] = "M1"
-        tmd_distributions["mtypes"][mtype]["diameter"]["method"] = "M1"
+        tmd_parameters[mtype]["diameter_params"]["method"] = diametrizer
+        tmd_distributions["mtypes"][mtype]["diameter"]["method"] = diametrizer
+
+        if diametrizer == "external":
+
+            def external_diametrizer(neuron, model, neurite_type):
+                return build_diameters.build(
+                    neuron,
+                    model,
+                    [neurite_type],
+                    tmd_parameters[mtype][  # pylint: disable=cell-var-from-loop
+                        "diameter_params"
+                    ],
+                )
+
+        else:
+            external_diametrizer = None
 
         gids = range(global_gid, global_gid + n_cells)
         global_gid += n_cells
@@ -70,6 +95,7 @@ def grow_vacuum_morphologies(
                 tmd_parameters[mtype],
                 tmd_distributions["mtypes"][mtype],
                 morphology_base_path,
+                external_diametrizer=external_diametrizer,
             )
             for gid in gids
         ):