Merge pull request #248 from NeuroML/feat/improve-schematic-morph-plots

feat(morph-plots): improve colouring of cells/groups in schematic plot

pyneuroml/plot/PlotMorphologyVispy.py

@@ -453,7 +453,7 @@ def plot_interactive_3D(
                         offset=pos,
                         cell=cell,
                         segment_groups=None,
-                        labels=True,
+                        color=color,
                         verbose=verbose,
                         current_scene=current_scene,
                         current_view=current_view,
@@ -525,8 +525,9 @@ def plot_3D_cell_morphology(
     :type offset: [float, float]
     :param cell: cell to plot
     :type cell: neuroml.Cell
-    :param color: color to use for segments:
+    :param color: color to use for segments, with some special values:
 
+        - if a color string is given, use that
         - if None, each segment is given a new unique color
         - if "Groups", each unbranched segment group is given a unique color,
           and segments that do not belong to an unbranched segment group are in
@@ -721,6 +722,7 @@ def plot_3D_schematic(
     current_scene: scene.SceneCanvas = None,
     current_view: scene.ViewBox = None,
     theme: str = "light",
+    color: typing.Optional[str] = "Cell",
 ) -> None:
     """Plot a 3D schematic of the provided segment groups using vispy.
     layer..
@@ -767,10 +769,31 @@ def plot_3D_schematic(
     :type current_view: ViewBox
     :param theme: theme to use (light/dark)
     :type theme: str
+    :param color: color to use for segment groups with some special values:
+
+        - if None, each unbranched segment group is given a unique color,
+        - if "Cell", each cell is given a unique color
+        - if "Default Groups", each cell is given unique colors for all axons,
+          dendrites, and soma segments
+
+    :type color: str
     """
     if title == "":
         title = f"3D schematic of segment groups from {cell.id}"
 
+    try:
+        soma_segs = cell.get_all_segments_in_group("soma_group")
+    except Exception:
+        soma_segs = []
+    try:
+        dend_segs = cell.get_all_segments_in_group("dendrite_group")
+    except Exception:
+        dend_segs = []
+    try:
+        axon_segs = cell.get_all_segments_in_group("axon_group")
+    except Exception:
+        axon_segs = []
+
     # if no segment groups are given, do them all
     if segment_groups is None:
         segment_groups = []
@@ -794,6 +817,10 @@ def plot_3D_schematic(
     colors = []
     text = []
     textpoints = []
+    # colors for cell
+    cell_color_soma = get_next_hex_color()
+    cell_color_axon = get_next_hex_color()
+    cell_color_dendrites = get_next_hex_color()
 
     for sgid, segs in ord_segs.items():
         sgobj = cell.get_segment_group(sgid)
@@ -821,8 +848,29 @@ def plot_3D_schematic(
                 offset[2] + last_seg.distal.z,
             ]
         )
-        colors.append(get_next_hex_color())
-        colors.append(get_next_hex_color())
+        if color is None:
+            colors.append(get_next_hex_color())
+            colors.append(get_next_hex_color())
+        elif color == "Cell":
+            colors.append(cell_color_soma)
+            colors.append(cell_color_soma)
+        elif color == "Default Groups":
+            if first_seg.id in soma_segs:
+                colors.append(cell_color_soma)
+                colors.append(cell_color_soma)
+            elif first_seg.id in axon_segs:
+                colors.append(cell_color_axon)
+                colors.append(cell_color_axon)
+            elif first_seg.id in dend_segs:
+                colors.append(cell_color_dendrites)
+                colors.append(cell_color_dendrites)
+            else:
+                colors.append(get_next_hex_color())
+                colors.append(get_next_hex_color())
+        else:
+            colors.append(color)
+            colors.append(color)
+
         toconnect.append([len(points) - 2, len(points) - 1])
 
         # TODO: needs fixing to show labels
@@ -853,7 +901,7 @@ def plot_3D_schematic(
         connect=numpy.array(toconnect),
     )
     if labels:
-        print("Text rendering")
+        logger.debug("Text rendering")
         scene.Text(
             text, pos=textpoints, font_size=30, color="black", parent=current_view.scene
         )