[MRG] ENH: Custom section colors in Cell.plot_morphology

hnn_core/cell.py

@@ -655,13 +655,20 @@ def parconnect_from_src(self, gid_presyn, nc_dict, postsyn,
 
         return nc
 
-    def plot_morphology(self, ax=None, cell_types=None, show=True):
+    def plot_morphology(self, ax=None, color=None, show=True):
         """Plot the cell morphology.
 
         Parameters
         ----------
         ax : instance of Axes3D
             Matplotlib 3D axis
+        color : str | dict | None
+            Color of cell. If str, entire cell plotted with
+            color indicated by str. If dict, colors of individual sections
+            can be specified. Must have a key for every section in cell as
+            defined in the `Cell.sections` attribute.
+        | Ex: ``{'apical_trunk': 'r', 'soma': 'b', ...}``
+
         show : bool
             If True, show the plot
 
@@ -670,7 +677,7 @@ def plot_morphology(self, ax=None, cell_types=None, show=True):
         axes : instance of Axes3D
             The matplotlib 3D axis handle.
         """
-        return plot_cell_morphology(self, ax=ax, show=show)
+        return plot_cell_morphology(self, ax=ax, color=color, show=show)
 
     def _update_end_pts(self):
         """"Create cell and copy coordinates to Section.end_pts"""

hnn_core/tests/test_viz.py

@@ -60,6 +60,24 @@ def test_network_visualization():
     with pytest.raises(ValueError, match='src_gid -1 not a valid cell ID'):
         plot_cell_connectivity(net, conn_idx, src_gid=-1)
 
+    # Test morphology plotting
+    for cell_type in net.cell_types.values():
+        cell_type.plot_morphology()
+        cell_type.plot_morphology(color='r')
+
+        sections = list(cell_type.sections.keys())
+        section_color = {sect_name: f'C{idx}' for
+                         idx, sect_name in enumerate(sections)}
+        cell_type.plot_morphology(color=section_color)
+
+    cell_type = net.cell_types['L2_basket']
+    with pytest.raises(ValueError):
+        cell_type.plot_morphology(color='z')
+    with pytest.raises(ValueError):
+        cell_type.plot_morphology(color={'soma': 'z'})
+    with pytest.raises(TypeError, match='color must be'):
+        cell_type.plot_morphology(color=123)
+
     plt.close('all')
 
     # test interactive clicking updates the position of src_cell in plot

hnn_core/viz.py

@@ -810,7 +810,7 @@ def _linewidth_from_data_units(ax, linewidth):
     return linewidth * (length / value_range)
 
 
-def plot_cell_morphology(cell, ax, show=True):
+def plot_cell_morphology(cell, ax, color=None, show=True):
     """Plot the cell morphology.
 
     Parameters
@@ -821,6 +821,13 @@ def plot_cell_morphology(cell, ax, show=True):
         Matplotlib 3D axis
     show : bool
         If True, show the plot
+    color : str | dict | None
+        Color of cell. If str, entire cell plotted with
+        color indicated by str. If dict, colors of individual sections
+        can be specified. Must have a key for every section in cell as
+        defined in the `Cell.sections` attribute.
+
+        | Ex: ``{'apical_trunk': 'r', 'soma': 'b', ...}``
 
     Returns
     -------
@@ -834,6 +841,15 @@ def plot_cell_morphology(cell, ax, show=True):
         plt.figure()
         ax = plt.axes(projection='3d')
 
+    _validate_type(color, (str, dict, None), 'color')
+
+    if color is None:
+        section_colors = {section: 'b' for section in cell.sections.keys()}
+    if isinstance(color, str):
+        section_colors = {section: color for section in cell.sections.keys()}
+    if isinstance(color, dict):
+        section_colors = color
+
     # Cell is in XZ plane
     ax.set_xlim((cell.pos[1] - 250, cell.pos[1] + 150))
     ax.set_zlim((cell.pos[2] - 100, cell.pos[2] + 1200))
@@ -849,7 +865,8 @@ def plot_cell_morphology(cell, ax, show=True):
             xs.append(pt[0] + dx)
             ys.append(pt[1] + dz)
             zs.append(pt[2] + dy)
-        ax.plot(xs, ys, zs, 'b-', linewidth=linewidth)
+        ax.plot(xs, ys, zs, '-', linewidth=linewidth,
+                color=section_colors[sec_name])
     ax.view_init(0, -90)
     ax.axis('off')