Merge pull request #294 from NeuroML/feat/morphplots-mark-segments

feat: morphology plotting: allow marking/highlighting individual segments

Author: sanjayankur31

pyneuroml/plot/PlotMorphology.py

@@ -182,6 +182,7 @@ def plot_2D(
     plot_spec: typing.Optional[
         typing.Dict[str, typing.Union[str, typing.List[int], float]]
     ] = None,
+    highlight_spec: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
 ):
     """Plot cells in a 2D plane.
 
@@ -192,6 +193,10 @@ def plot_2D(
 
     This method uses matplotlib.
 
+    .. versionadded:: 1.1.12
+        The hightlight_spec parameter
+
+
     :param nml_file: path to NeuroML cell file, or a NeuroMLDocument object
     :type nml_file: str or :py:class:`neuroml.NeuroMLDocument` or
         :py:class:`neuroml.Cell`
@@ -239,13 +244,43 @@ def plot_2D(
         The last three lists override the point_fraction setting. If a cell id
         is not included in the spec here, it will follow the plot_type provided
         before.
+    :type plot_spec: dict
+    :param highlight_spec: dictionary that allows passing some
+        specifications to allow highlighting of particular elements.  Only used
+        when plotting multi-compartmental cells for marking segments on them
+        ("plot_type" is either "constant" or "detailed")
+
+        Each key in the dictionary will be of the cell id and the values will
+        be more dictionaries, with the segment id as key and the following keys
+        in it:
+
+        - marker_color: color of the marker
+        - marker_size: width of the marker
+
+        E.g.:
+
+        .. code-block:: python
+
+        {
+            "cell id1": {
+                "seg id1": {
+                    "marker_color": "blue",
+                    "marker_size": 10
+                }
+            }
+        }
+
+    :type highlight_spec: dict
     """
 
     if plot_type not in ["detailed", "constant", "schematic", "point"]:
         raise ValueError(
             "plot_type must be one of 'detailed', 'constant', 'schematic', 'point'"
         )
 
+    if highlight_spec is None:
+        highlight_spec = {}
+
     if verbose:
         print("Plotting %s" % nml_file)
 
@@ -412,6 +447,11 @@ def plot_2D(
                     or plot_type == "constant"
                     or cell.id in constant_cells
                 ):
+                    cell_highlight_spec = {}
+                    try:
+                        cell_highlight_spec = highlight_spec[cell.id]
+                    except KeyError:
+                        pass
                     plot_2D_cell_morphology(
                         offset=pos,
                         cell=cell,
@@ -427,6 +467,7 @@ def plot_2D(
                         nogui=True,
                         autoscale=False,
                         square=False,
+                        highlight_spec=cell_highlight_spec,
                     )
 
     add_scalebar_to_matplotlib_plot(axis_min_max, ax)
@@ -467,6 +508,7 @@ def plot_2D_cell_morphology(
     datamin: typing.Optional[float] = None,
     datamax: typing.Optional[float] = None,
     colormap_name: str = "viridis",
+    highlight_spec: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
 ):
     """Plot the detailed 2D morphology of a cell in provided plane.
 
@@ -534,6 +576,16 @@ def plot_2D_cell_morphology(
     :type datamin: float
     :param datamax: max limits of data (useful to compare different plots)
     :type datamax: float
+    :param highlight_spec: dictionary that allows passing some
+        specifications to allow highlighting of particular elements. Mostly
+        only helpful for marking segments on multi-compartmental cells. In the
+        main dictionary are more dictionaries, one for each segment id which
+        will be the key:
+
+        - marker_color: color of the marker
+        - marker_size: width of the marker
+
+    :type highlight_spec: dict
 
     :raises: ValueError if `cell` is None
 
@@ -543,6 +595,10 @@ def plot_2D_cell_morphology(
             "No cell provided. If you would like to plot a network of point neurons, consider using `plot_2D_point_cells` instead"
         )
 
+    if highlight_spec is None:
+        highlight_spec = {}
+    logging.debug("highlight_spec is " + str(highlight_spec))
+
     try:
         soma_segs = cell.get_all_segments_in_group("soma_group")
     except Exception:
@@ -599,12 +655,27 @@ def plot_2D_cell_morphology(
         d = seg.distal
         width = (p.diameter + d.diameter) / 2
 
+        segment_spec = {
+            "marker_size": None,
+            "marker_color": None,
+        }
+        try:
+            segment_spec.update(highlight_spec[str(seg.id)])
+        # if there's no spec for this segment
+        except KeyError:
+            logger.debug("No segment highlight spec found for segment" + str(seg.id))
+
+        logger.debug("segment_spec for " + str(seg.id) + " is" + str(segment_spec))
+
         if width < min_width:
             width = min_width
 
         if plot_type == "constant":
             width = min_width
 
+        if segment_spec["marker_size"] is not None:
+            width = float(segment_spec["marker_size"])
+
         if overlay_data and acolormap and norm:
             try:
                 seg_color = acolormap(norm(overlay_data[seg.id]))
@@ -617,6 +688,9 @@ def plot_2D_cell_morphology(
             elif seg.id in axon_segs:
                 seg_color = "r"
 
+        if segment_spec["marker_color"] is not None:
+            seg_color = segment_spec["marker_color"]
+
         spherical = (
             p.x == d.x and p.y == d.y and p.z == d.z and p.diameter == d.diameter
         )

pyneuroml/plot/PlotMorphologyVispy.py

@@ -251,12 +251,17 @@ def plot_interactive_3D(
     plot_spec: typing.Optional[
         typing.Dict[str, typing.Union[str, typing.List[int], float]]
     ] = None,
+    highlight_spec: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
     precision: typing.Tuple[int, int] = (4, 200),
 ):
     """Plot interactive plots in 3D using Vispy
 
     https://vispy.org
 
+    .. versionadded:: 1.1.12
+        The hightlight_spec parameter
+
+
     :param nml_file: path to NeuroML cell file or
         :py:class:`neuroml.NeuroMLDocument` or :py:class:`neuroml.Cell` object
     :type nml_file: str or neuroml.NeuroMLDocument or neuroml.Cell
@@ -298,6 +303,34 @@ def plot_interactive_3D(
         The last three lists override the point_fraction setting. If a cell id
         is not included in the spec here, it will follow the plot_type provided
         before.
+    :type plot_spec: dict
+    :param highlight_spec: dictionary that allows passing some
+        specifications to allow highlighting of particular elements.  Only used
+        when plotting multi-compartmental cells for marking segments on them
+        ("plot_type" is either "constant" or "detailed")
+
+        Each key in the dictionary will be of the cell id and the values will
+        be more dictionaries, with the segment id as key and the following keys
+        in it:
+
+        - marker_color: color of the marker
+        - marker_size: [diameter 1, diameter 2] (in case of sphere, the first value
+          is used)
+
+        E.g.:
+
+        .. code-block:: python
+
+        {
+            "cell id1": {
+                "seg id1": {
+                    "marker_color": "blue",
+                    "marker_size": [0.1, 0.1]
+                }
+            }
+        }
+
+    :type highlight_spec: dict
     :param precision: tuple containing two values: (number of decimal places,
         maximum number of meshes). The first is used to group segments into
         meshes to create instances. More precision means fewer segments will be
@@ -315,6 +348,9 @@ def plot_interactive_3D(
             "plot_type must be one of 'detailed', 'constant', 'schematic', 'point'"
         )
 
+    if highlight_spec is None:
+        highlight_spec = {}
+
     if verbose:
         logger.info(f"Visualising {nml_file}")
 
@@ -525,6 +561,7 @@ def plot_interactive_3D(
                     logger.debug(f"meshdata added: {key}: {meshdata[key]}")
 
                 elif plot_type == "schematic" or cell.id in schematic_cells:
+                    logger.debug(f"Cell for 3d schematic is: {cell.id}")
                     plot_3D_schematic(
                         offset=pos,
                         cell=cell,
@@ -544,6 +581,12 @@ def plot_interactive_3D(
                     or cell.id in constant_cells
                 ):
                     logger.debug(f"Cell for 3d is: {cell.id}")
+                    cell_highlight_spec = {}
+                    try:
+                        cell_highlight_spec = highlight_spec[cell.id]
+                    except KeyError:
+                        pass
+
                     plot_3D_cell_morphology(
                         offset=pos,
                         cell=cell,
@@ -556,6 +599,7 @@ def plot_interactive_3D(
                         nogui=True,
                         meshdata=meshdata,
                         mesh_precision=precision[0],
+                        highlight_spec=cell_highlight_spec,
                     )
 
             # if too many meshes, reduce precision and retry, recursively
@@ -566,15 +610,16 @@ def plot_interactive_3D(
                     f"More meshes than threshold ({len(meshdata.keys())}/{precision[1]}), reducing precision to {precision[0]} and re-calculating."
                 )
                 plot_interactive_3D(
-                    nml_model,
-                    min_width,
-                    verbose,
-                    plot_type,
-                    title,
-                    theme,
-                    nogui,
-                    plot_spec,
-                    precision,
+                    nml_file=nml_model,
+                    min_width=min_width,
+                    verbose=verbose,
+                    plot_type=plot_type,
+                    title=title,
+                    theme=theme,
+                    nogui=nogui,
+                    plot_spec=plot_spec,
+                    precision=precision,
+                    highlight_spec=highlight_spec,
                 )
                 # break the recursion, don't plot in the calling method
                 return
@@ -603,12 +648,16 @@ def plot_3D_cell_morphology(
     theme: str = "light",
     meshdata: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
     mesh_precision: int = 2,
+    highlight_spec: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
 ):
     """Plot the detailed 3D morphology of a cell using vispy.
     https://vispy.org/
 
     .. versionadded:: 1.0.0
 
+    .. versionadded:: 1.1.12
+        The hightlight_spec parameter
+
     .. seealso::
 
         :py:func:`plot_2D`
@@ -671,6 +720,17 @@ def plot_3D_cell_morphology(
         instances: more precision means more detail (meshes), means less
         performance
     :type mesh_precision: int
+    :param highlight_spec: dictionary that allows passing some
+        specifications to allow highlighting of particular elements. Mostly
+        only helpful for marking segments on multi-compartmental cells. In the
+        main dictionary are more dictionaries, one for each segment id which
+        will be the key:
+
+        - marker_color: color of the marker
+        - marker_size: [diameter 1, diameter 2] (in case of sphere, the first value
+          is used)
+
+    :type highlight_spec: dict
     :raises: ValueError if `cell` is None
 
     """
@@ -679,6 +739,10 @@ def plot_3D_cell_morphology(
             "No cell provided. If you would like to plot a network of point neurons, consider using `plot_2D_point_cells` instead"
         )
 
+    if highlight_spec is None:
+        highlight_spec = {}
+    logging.debug("highlight_spec is " + str(highlight_spec))
+
     try:
         soma_segs = cell.get_all_segments_in_group("soma_group")
     except Exception:
@@ -726,6 +790,25 @@ def plot_3D_cell_morphology(
         # round up to precision
         r1 = round(p.diameter / 2, mesh_precision)
         r2 = round(d.diameter / 2, mesh_precision)
+
+        segment_spec = {
+            "marker_size": None,
+            "marker_color": None,
+        }
+        try:
+            segment_spec.update(highlight_spec[str(seg.id)])
+        # if there's no spec for this segment
+        except KeyError:
+            logger.debug("No segment highlight spec found for segment" + str(seg.id))
+
+        logger.debug("segment_spec for " + str(seg.id) + " is" + str(segment_spec))
+
+        if segment_spec["marker_size"] is not None:
+            if type(segment_spec["marker_size"]) is not list:
+                raise RuntimeError("The marker size must be a list")
+            r1 = round(float(segment_spec["marker_size"][0]) / 2, mesh_precision)
+            r2 = round(float(segment_spec["marker_size"][1]) / 2, mesh_precision)
+
         key = (
             f"{r1:.{mesh_precision}f}",
             f"{r2:.{mesh_precision}f}",
@@ -756,6 +839,9 @@ def plot_3D_cell_morphology(
         else:
             seg_color = color
 
+        if segment_spec["marker_color"] is not None:
+            seg_color = segment_spec["marker_color"]
+
         try:
             meshdata[key].append((p, d, seg_color, offset))
         except KeyError: