Improve plot_collage

Change-Id: Ice742e539c409fb748170e19504edfd9d2e30927

src/synthesis_workflow/tasks/luigi_tools.py

@@ -186,26 +186,44 @@ def __init__(
         )
 
 
-class OptionalIntParameter(luigi.OptionalParameter, luigi.IntParameter):
-    """Class to parse optional int parameters"""
+class OptionalParameter(luigi.OptionalParameter):
+    """Mixin to make a parameter class optional"""
+
+    def __init__(self, *args, **kwargs):
+        self._cls = self.__class__
+        self._base_cls = self.__class__.__bases__[-1]
+        if OptionalParameter in (self._cls, self._base_cls):
+            raise TypeError(
+                "OptionalParameter can only be used as a mixin (must not be the rightmost "
+                "class in the class definition)"
+            )
+        super().__init__(*args, **kwargs)
 
     def parse(self, x):
-        if x:
-            return int(x)
+        if x and x.lower() != "null":
+            return self._base_cls.parse(self, x)
         else:
             return None
 
     def _warn_on_wrong_param_type(self, param_name, param_value):
-        if self.__class__ != OptionalIntParameter:
+        if self.__class__ != self._cls:
             return
         if not isinstance(param_value, int) and param_value is not None:
             warnings.warn(
-                'OptionalIntParameter "{}" with value "{}" is not of type int or None.'.format(
-                    param_name, param_value
+                '{} "{}" with value "{}" is not of type int or None.'.format(
+                    self._cls.__name__, param_name, param_value
                 )
             )
 
 
+class OptionalIntParameter(OptionalParameter, luigi.IntParameter):
+    """Class to parse optional int parameters"""
+
+
+class OptionalNumericalParameter(OptionalParameter, luigi.NumericalParameter):
+    """Class to parse optional int parameters"""
+
+
 class BoolParameter(luigi.BoolParameter):
     """Class to parse boolean parameters and set explicit parsing when default is True"""
 

src/synthesis_workflow/tasks/validation.py

@@ -6,8 +6,9 @@
 import pandas as pd
 import pkg_resources
 import yaml
-from voxcell import VoxelData
 from atlas_analysis.planes.planes import load_planes_centerline
+from neurom.view import view
+from voxcell import VoxelData
 
 from morphval import validation_main as morphval_validation
 from synthesis_workflow.tools import load_circuit
@@ -31,6 +32,7 @@
 from synthesis_workflow.tasks.config import ValidationLocalTarget
 from synthesis_workflow.tasks.luigi_tools import BoolParameter
 from synthesis_workflow.tasks.luigi_tools import copy_params
+from synthesis_workflow.tasks.luigi_tools import OptionalNumericalParameter
 from synthesis_workflow.tasks.luigi_tools import ParamLink
 from synthesis_workflow.tasks.luigi_tools import WorkflowError
 from synthesis_workflow.tasks.luigi_tools import WorkflowTask
@@ -203,11 +205,33 @@ class PlotCollage(WorkflowTask):
         nb_jobs (int) : number of joblib workers
         joblib_verbose (int) verbose level of joblib
         dpi (int): dpi for pdf rendering (rasterized)
-
+        realistic_diameters (bool): set or unset realistic diameter when NeuroM plot neurons
+        linewidth (float): linewidth used by NeuroM to plot neurons
+        diameter_scale (float): diameter scale used by NeuroM to plot neurons
     """
 
     collage_base_path = luigi.Parameter(default="collages")
     dpi = luigi.IntParameter(default=1000)
+    realistic_diameters = BoolParameter(
+        default=True,
+        description="Set or unset realistic diameter when NeuroM plot neurons",
+    )
+    linewidth = luigi.NumericalParameter(
+        default=0.1,
+        var_type=float,
+        min_value=0,
+        max_value=float("inf"),
+        left_op=luigi.parameter.operator.lt,
+        description="Linewidth used by NeuroM to plot neurons",
+    )
+    diameter_scale = OptionalNumericalParameter(
+        default=view._DIAMETER_SCALE,  # pylint: disable=protected-access
+        var_type=float,
+        min_value=0,
+        max_value=float("inf"),
+        left_op=luigi.parameter.operator.lt,
+        description="Diameter scale used by NeuroM to plot neurons",
+    )
 
     def requires(self):
         """"""
@@ -228,6 +252,9 @@ def run(self):
                 nb_jobs=self.nb_jobs,
                 joblib_verbose=self.joblib_verbose,
                 dpi=self.dpi,
+                realistic_diameters=self.realistic_diameters,
+                linewidth=self.linewidth,
+                diameter_scale=self.diameter_scale,
             )
 
     def output(self):
@@ -240,6 +267,10 @@ def output(self):
     joblib_verbose=ParamLink(RunnerConfig),
     collage_base_path=ParamLink(PlotCollage),
     sample=ParamLink(ValidationConfig),
+    dpi=ParamLink(PlotCollage),
+    realistic_diameters=ParamLink(PlotCollage),
+    linewidth=ParamLink(PlotCollage),
+    diameter_scale=ParamLink(PlotCollage),
 )
 class PlotSingleCollage(WorkflowTask):
     """Plot collage for single mtype.
@@ -254,7 +285,6 @@ class PlotSingleCollage(WorkflowTask):
     """
 
     mtype = luigi.Parameter()
-    dpi = luigi.IntParameter(default=1000)
 
     def requires(self):
         """"""
@@ -266,16 +296,27 @@ def requires(self):
 
     def run(self):
         """"""
+        mvd3_path = self.input()["synthesis"]["out_mvd3"].path
+        morphologies_path = self.input()["synthesis"]["out_morphologies"].path
+        atlas_path = CircuitConfig().atlas_path
+        L.debug("Load circuit mvd3 from %s", mvd3_path)
+        L.debug("Load circuit morphologies from %s", morphologies_path)
+        L.debug("Load circuit atlas from %s", atlas_path)
         circuit = load_circuit(
-            path_to_mvd3=self.input()["synthesis"]["out_mvd3"].path,
-            path_to_morphologies=self.input()["synthesis"]["out_morphologies"].path,
-            path_to_atlas=CircuitConfig().atlas_path,
+            path_to_mvd3=mvd3_path,
+            path_to_morphologies=morphologies_path,
+            path_to_atlas=atlas_path,
         )
 
-        planes = load_planes_centerline(self.input()["planes"].path)["planes"]
-        layer_annotation = VoxelData.load_nrrd(
-            self.input()["layers"]["annotations"].path
-        )
+        planes_path = self.input()["planes"].path
+        L.debug("Load planes from %s", planes_path)
+        planes = load_planes_centerline(planes_path)["planes"]
+
+        layer_annotation_path = self.input()["layers"]["annotations"].path
+        L.debug("Load layer annotations from %s", layer_annotation_path)
+        layer_annotation = VoxelData.load_nrrd(layer_annotation_path)
+
+        L.debug("Plot single collage")
         plot_collage(
             circuit,
             planes,
@@ -286,6 +327,11 @@ def run(self):
             nb_jobs=self.nb_jobs,
             joblib_verbose=self.joblib_verbose,
             dpi=self.dpi,
+            plot_neuron_kwargs={
+                "realistic_diameters": self.realistic_diameters,
+                "linewidth": self.linewidth,
+                "diameter_scale": self.diameter_scale,
+            },
         )
 
     def output(self):

src/synthesis_workflow/validation.py

@@ -393,16 +393,39 @@ def plot_cells(
     mtype=None,
     sample=10,
     atlas=None,
+    plot_neuron_kwargs=None,
 ):
     """Plot cells for collage."""
-    cells = circuit.cells.get({"mtype": mtype})
+
     if mtype is not None:
-        cells = cells[cells.mtype == mtype]
+        cells = circuit.cells.get({"mtype": mtype})
+    else:
+        cells = circuit.cells.get()
 
     if len(cells) == 0:
         raise Exception("no cells of that mtype")
-    gids = get_cells_between_planes(cells, plane_left, plane_right).index
-    for gid in gids[:sample]:
+
+    if plot_neuron_kwargs is None:
+        plot_neuron_kwargs = {}
+
+    gids = get_cells_between_planes(cells, plane_left, plane_right).index[:sample]
+
+    if atlas is not None:
+        vec = [0, 1, 0]
+        all_pos_orig = cells.loc[gids, ["x", "y", "z"]].values
+        all_orientations = atlas.orientations.lookup(all_pos_orig)
+        all_lookups = np.einsum("ijk, k", all_orientations, vec)
+        all_pos_final = all_pos_orig + all_lookups * 300
+        all_dist_plane_orig = all_pos_orig - plane_left.point
+        all_dist_plane_final = all_pos_final - plane_left.point
+        all_pos_orig_plane_coord = np.tensordot(
+            all_dist_plane_orig, rotation_matrix.T, axes=1
+        )
+        all_pos_final_plane_coord = np.tensordot(
+            all_dist_plane_final, rotation_matrix.T, axes=1
+        )
+
+    for num, gid in enumerate(gids):
         morphology = circuit.morph.get(gid, transform=True, source="ascii")
 
         def _to_plane_coord(p):
@@ -412,22 +435,17 @@ def _to_plane_coord(p):
         morphology = morphology.transform(_to_plane_coord)
 
         if atlas is not None:
-            pos_orig = circuit.cells.positions(gid).to_numpy()
-            pos_final = pos_orig + atlas.lookup_orientation(pos_orig) * 300
-            pos_orig = _to_plane_coord(pos_orig)
-            pos_final = _to_plane_coord(pos_final)
             plt.plot(
-                [pos_orig[0], pos_final[0]],
-                [pos_orig[1], pos_final[1]],
+                [all_pos_orig_plane_coord[num, 0], all_pos_final_plane_coord[num, 0]],
+                [all_pos_orig_plane_coord[num, 1], all_pos_final_plane_coord[num, 1]],
                 c="0.5",
                 lw=0.2,
             )
 
-        viewer.plot_neuron(
-            ax, morphology, plane="xy", realistic_diameters=True, linewidth=0.1
-        )
+        viewer.plot_neuron(ax, morphology, plane="xy", **plot_neuron_kwargs)
 
 
+# pylint: disable=too-many-arguments
 def _plot_collage(
     planes,
     layer_annotation=None,
@@ -439,6 +457,7 @@ def _plot_collage(
     n_pixels_y=64,
     with_y_field=True,
     with_cells=True,
+    plot_neuron_kwargs=None,
 ):
     """Internal plot collage for multiprocessing."""
     if with_y_field and atlas is None:
@@ -459,6 +478,7 @@ def _plot_collage(
         cmap=matplotlib.colors.ListedColormap(["C0", "C1", "C2", "C3", "C4", "C5"]),
         alpha=0.3,
     )
+
     plt.colorbar()
     if with_cells:
         plot_cells(
@@ -470,7 +490,9 @@ def _plot_collage(
             mtype=mtype,
             sample=sample,
             atlas=atlas,
+            plot_neuron_kwargs=plot_neuron_kwargs,
         )
+
     if with_y_field:
         # note: some of these parameters are harcoded for NCx plot, adjust as needed
         X_y, Y_y, orientation_u, orientation_v = get_y_info(
@@ -486,10 +508,12 @@ def _plot_collage(
             scale_units="xy",
             scale=1,
         )
+
     ax = plt.gca()
     ax.set_aspect("equal")
     ax.set_rasterized(True)
     ax.set_title("")
+
     return fig
 
 
@@ -507,6 +531,7 @@ def plot_collage(
     n_pixels=1024,
     with_y_field=True,
     n_pixels_y=64,
+    plot_neuron_kwargs=None,
 ):
     """Plot collage of an mtype and a list of planes.
 
@@ -539,6 +564,7 @@ def plot_collage(
             n_pixels=n_pixels,
             n_pixels_y=n_pixels_y,
             with_y_field=with_y_field,
+            plot_neuron_kwargs=plot_neuron_kwargs,
         )
         for fig in Parallel(nb_jobs, verbose=joblib_verbose)(
             delayed(f)(planes) for planes in zip(planes[:-1:3], planes[2::3])

tests/data/in_small_O1/luigi.cfg

@@ -2,7 +2,7 @@
 logging_conf_file = logging.conf
 
 [RunnerConfig]
-nb_jobs = 1
+nb_jobs = 2
 
 [DiametrizerConfig]
 model = generic
@@ -59,6 +59,11 @@ debug_region_grower_scales = true
 [BuildAxonMorphologies]
 axon_cells_path = <path_to_clones>
 
+[PlotCollage]
+realistic_diameters = False
+linewidth = 0.04
+diameter_scale = null
+
 ####################################
 # ########## Validation ########## #
 ####################################