Merge pull request #253 from NeuroML/feat/hybrid-plots

Feat/hybrid plots: allow users to specify how they want various cells in networks to be plotted

pyneuroml/plot/PlotMorphology.py

@@ -75,10 +75,17 @@ def process_args():
         help="Plane to plot on for 2D plot",
     )
 
+    parser.add_argument(
+        "-pointFraction",
+        type=str,
+        metavar="<fraction of each population to plot as point cells>",
+        default=DEFAULTS["pointFraction"],
+        help="Fraction of network to plot as point cells",
+    )
     parser.add_argument(
         "-plotType",
         type=str,
-        metavar="<type: detailed, constant, or schematic>",
+        metavar="<type: detailed, constant, schematic, or point>",
         default=DEFAULTS["plotType"],
         help="Level of detail to plot in",
     )
@@ -147,6 +154,7 @@ def plot_from_console(a: typing.Optional[typing.Any] = None, **kwargs: str):
             verbose=a.v,
             plot_type=a.plot_type,
             theme=a.theme,
+            plot_spec={"point_fraction": a.point_fraction},
         )
     else:
         plot_2D(
@@ -158,6 +166,7 @@ def plot_from_console(a: typing.Optional[typing.Any] = None, **kwargs: str):
             a.save_to_file,
             a.square,
             a.plot_type,
+            plot_spec={"point_fraction": a.point_fraction},
         )
 
 
@@ -172,6 +181,9 @@ def plot_2D(
     plot_type: str = "detailed",
     title: typing.Optional[str] = None,
     close_plot: bool = False,
+    plot_spec: typing.Optional[
+        typing.Dict[str, typing.Union[str, typing.List[int], float]]
+    ] = None,
 ):
     """Plot cells in a 2D plane.
 
@@ -205,6 +217,7 @@ def plot_2D(
         - "constant": show morphology, but use constant line widths
         - "schematic": only plot each unbranched segment group as a straight
           line, not following each segment
+        - "point": show all cells as points
 
         This is only applicable for neuroml.Cell cells (ones with some
         morphology)
@@ -214,20 +227,36 @@ def plot_2D(
     :type title: str
     :param close_plot: call pyplot.close() to close plot after plotting
     :type close_plot: bool
+    :param plot_spec: dictionary that allows passing some specifications that
+        control how a plot is generated. This is mostly useful for large
+        network plots where one may want to have a mix of full morphology and
+        schematic, and point representations of cells. Possible keys are:
+
+        - point_fraction: what fraction of each population to plot as point cells:
+          these cells will be randomly selected
+        - points_cells: list of cell ids to plot as point cells
+        - schematic_cells: list of cell ids to plot as schematics
+        - constant_cells: list of cell ids to plot as constant widths
+
+        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.
     """
 
-    if plot_type not in ["detailed", "constant", "schematic"]:
+    if plot_type not in ["detailed", "constant", "schematic", "point"]:
         raise ValueError(
-            "plot_type must be one of 'detailed', 'constant', or 'schematic'"
+            "plot_type must be one of 'detailed', 'constant', 'schematic', 'point'"
         )
 
     if verbose:
         print("Plotting %s" % nml_file)
 
-    if type(nml_file) == str:
+    # do not recursive read the file, the extract_position_info function will
+    # do that for us, from a copy of the model
+    if type(nml_file) is str:
         nml_model = read_neuroml2_file(
             nml_file,
-            include_includes=True,
+            include_includes=False,
             check_validity_pre_include=False,
             verbose=False,
             optimized=True,
@@ -250,7 +279,9 @@ def plot_2D(
         positions,
         pop_id_vs_color,
         pop_id_vs_radii,
-    ) = extract_position_info(nml_model, verbose)
+    ) = extract_position_info(
+        nml_model, verbose, nml_file if type(nml_file) is str else ""
+    )
 
     if title is None:
         if len(nml_model.networks) > 0:
@@ -268,12 +299,45 @@ def plot_2D(
     fig, ax = get_new_matplotlib_morph_plot(title, plane2d)
     axis_min_max = [float("inf"), -1 * float("inf")]
 
-    for pop_id in pop_id_vs_cell:
-        cell = pop_id_vs_cell[pop_id]
-        pos_pop = positions[pop_id]
+    # process plot_spec
+    point_cells = []  # type: typing.List[int]
+    schematic_cells = []  # type: typing.List[int]
+    constant_cells = []  # type: typing.List[int]
+    detailed_cells = []  # type: typing.List[int]
+    if plot_spec is not None:
+        try:
+            point_cells = plot_spec["point_cells"]
+        except KeyError:
+            pass
+        try:
+            schematic_cells = plot_spec["schematic_cells"]
+        except KeyError:
+            pass
+        try:
+            constant_cells = plot_spec["constant_cells"]
+        except KeyError:
+            pass
+        try:
+            detailed_cells = plot_spec["detailed_cells"]
+        except KeyError:
+            pass
+
+    for pop_id, cell in pop_id_vs_cell.items():
+        pos_pop = positions[pop_id]  # type: typing.Dict[typing.Any, typing.List[float]]
+
+        # reinit point_cells for each loop
+        point_cells_pop = []
+        if len(point_cells) == 0 and plot_spec is not None:
+            cell_indices = list(pos_pop.keys())
+            try:
+                point_cells_pop = random.sample(
+                    cell_indices,
+                    int(len(cell_indices) * float(plot_spec["point_fraction"])),
+                )
+            except KeyError:
+                pass
 
-        for cell_index in pos_pop:
-            pos = pos_pop[cell_index]
+        for cell_index, pos in pos_pop.items():
             radius = pop_id_vs_radii[pop_id] if pop_id in pop_id_vs_radii else 10
             color = pop_id_vs_color[pop_id] if pop_id in pop_id_vs_color else None
 
@@ -291,12 +355,36 @@ def plot_2D(
                     nogui=True,
                 )
             else:
-                if plot_type == "schematic":
+                if (
+                    plot_type == "point"
+                    or cell_index in point_cells_pop
+                    or cell.id in point_cells
+                ):
+                    # assume that soma is 0, plot point at where soma should be
+                    soma_x_y_z = cell.get_actual_proximal(0)
+                    pos1 = [
+                        pos[0] + soma_x_y_z.x,
+                        pos[1] + soma_x_y_z.y,
+                        pos[2] + soma_x_y_z.z,
+                    ]
+                    plot_2D_point_cells(
+                        offset=pos1,
+                        plane2d=plane2d,
+                        color=color,
+                        soma_radius=radius,
+                        verbose=verbose,
+                        ax=ax,
+                        fig=fig,
+                        autoscale=False,
+                        scalebar=False,
+                        nogui=True,
+                    )
+                elif plot_type == "schematic" or cell.id in schematic_cells:
                     plot_2D_schematic(
                         offset=pos,
                         cell=cell,
                         segment_groups=None,
-                        labels=True,
+                        labels=False,
                         plane2d=plane2d,
                         verbose=verbose,
                         fig=fig,
@@ -306,7 +394,12 @@ def plot_2D(
                         autoscale=False,
                         square=False,
                     )
-                else:
+                elif (
+                    plot_type == "detailed"
+                    or cell.id in detailed_cells
+                    or plot_type == "constant"
+                    or cell.id in constant_cells
+                ):
                     plot_2D_cell_morphology(
                         offset=pos,
                         cell=cell,

pyneuroml/plot/PlotMorphologyVispy.py

@@ -11,22 +11,17 @@
 
 
 import logging
-import typing
-import numpy
+import random
 import textwrap
-from vispy import scene, app
+import typing
 
-from pyneuroml.utils.plot import (
-    DEFAULTS,
-    get_cell_bound_box,
-    get_next_hex_color,
-)
+import numpy
+from neuroml import Cell, NeuroMLDocument, Segment, SegmentGroup
+from neuroml.neuro_lex_ids import neuro_lex_ids
 from pyneuroml.pynml import read_neuroml2_file
 from pyneuroml.utils import extract_position_info
-
-from neuroml import Cell, NeuroMLDocument, SegmentGroup, Segment
-from neuroml.neuro_lex_ids import neuro_lex_ids
-
+from pyneuroml.utils.plot import DEFAULTS, get_cell_bound_box, get_next_hex_color
+from vispy import app, scene
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -296,6 +291,9 @@ def plot_interactive_3D(
     title: typing.Optional[str] = None,
     theme: str = "light",
     nogui: bool = False,
+    plot_spec: typing.Optional[
+        typing.Dict[str, typing.Union[str, typing.List[int], float]]
+    ] = None,
 ):
     """Plot interactive plots in 3D using Vispy
 
@@ -316,6 +314,7 @@ def plot_interactive_3D(
         - "constant": show morphology, but use constant line widths
         - "schematic": only plot each unbranched segment group as a straight
           line, not following each segment
+        - "point": show all cells as points
 
         This is only applicable for neuroml.Cell cells (ones with some
         morphology)
@@ -327,19 +326,33 @@ def plot_interactive_3D(
     :type theme: str
     :param nogui: toggle showing gui (for testing only)
     :type nogui: bool
+    :param plot_spec: dictionary that allows passing some specifications that
+        control how a plot is generated. This is mostly useful for large
+        network plots where one may want to have a mix of full morphology and
+        schematic, and point representations of cells. Possible keys are:
+
+        - point_fraction: what fraction of each population to plot as point cells:
+          these cells will be randomly selected
+        - points_cells: list of cell ids to plot as point cells
+        - schematic_cells: list of cell ids to plot as schematics
+        - constant_cells: list of cell ids to plot as constant widths
+
+        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.
     """
-    if plot_type not in ["detailed", "constant", "schematic"]:
+    if plot_type not in ["detailed", "constant", "schematic", "point"]:
         raise ValueError(
-            "plot_type must be one of 'detailed', 'constant', or 'schematic'"
+            "plot_type must be one of 'detailed', 'constant', 'schematic', 'point'"
         )
 
     if verbose:
         print(f"Plotting {nml_file}")
 
-    if type(nml_file) == str:
+    if type(nml_file) is str:
         nml_model = read_neuroml2_file(
             nml_file,
-            include_includes=True,
+            include_includes=False,
             check_validity_pre_include=False,
             verbose=False,
             optimized=True,
@@ -360,7 +373,9 @@ def plot_interactive_3D(
         positions,
         pop_id_vs_color,
         pop_id_vs_radii,
-    ) = extract_position_info(nml_model, verbose)
+    ) = extract_position_info(
+        nml_model, verbose, nml_file if type(nml_file) is str else ""
+    )
 
     # Collect all markers and only plot one markers object
     # this is more efficient than multiple markers, one for each point.
@@ -429,12 +444,45 @@ def plot_interactive_3D(
 
     logger.debug(f"figure extents are: {view_min}, {view_max}")
 
-    for pop_id in pop_id_vs_cell:
-        cell = pop_id_vs_cell[pop_id]
-        pos_pop = positions[pop_id]
+    # process plot_spec
+    point_cells = []  # type: typing.List[int]
+    schematic_cells = []  # type: typing.List[int]
+    constant_cells = []  # type: typing.List[int]
+    detailed_cells = []  # type: typing.List[int]
+    if plot_spec is not None:
+        try:
+            point_cells = plot_spec["point_cells"]
+        except KeyError:
+            pass
+        try:
+            schematic_cells = plot_spec["schematic_cells"]
+        except KeyError:
+            pass
+        try:
+            constant_cells = plot_spec["constant_cells"]
+        except KeyError:
+            pass
+        try:
+            detailed_cells = plot_spec["detailed_cells"]
+        except KeyError:
+            pass
+
+    for pop_id, cell in pop_id_vs_cell.items():
+        pos_pop = positions[pop_id]  # type: typing.Dict[typing.Any, typing.List[float]]
 
-        for cell_index in pos_pop:
-            pos = pos_pop[cell_index]
+        # reinit point_cells for each loop
+        point_cells_pop = []
+        if len(point_cells) == 0 and plot_spec is not None:
+            cell_indices = list(pos_pop.keys())
+            try:
+                point_cells_pop = random.sample(
+                    cell_indices,
+                    int(len(cell_indices) * float(plot_spec["point_fraction"])),
+                )
+            except KeyError:
+                pass
+
+        for cell_index, pos in pos_pop.items():
             radius = pop_id_vs_radii[pop_id] if pop_id in pop_id_vs_radii else 10
             color = pop_id_vs_color[pop_id] if pop_id in pop_id_vs_color else None
 
@@ -448,7 +496,24 @@ def plot_interactive_3D(
                 marker_sizes.extend([radius])
                 marker_colors.extend([color])
             else:
-                if plot_type == "schematic":
+                if (
+                    plot_type == "point"
+                    or cell_index in point_cells_pop
+                    or cell.id in point_cells
+                ):
+                    # assume that soma is 0, plot point at where soma should be
+                    soma_x_y_z = cell.get_actual_proximal(0)
+                    pos1 = [
+                        pos[0] + soma_x_y_z.x,
+                        pos[1] + soma_x_y_z.y,
+                        pos[2] + soma_x_y_z.z,
+                    ]
+                    marker_points.extend([pos1])
+                    # larger than the default soma width, which would be too
+                    # small
+                    marker_sizes.extend([25])
+                    marker_colors.extend([color])
+                elif plot_type == "schematic" or cell.id in schematic_cells:
                     plot_3D_schematic(
                         offset=pos,
                         cell=cell,
@@ -459,7 +524,12 @@ def plot_interactive_3D(
                         current_view=current_view,
                         nogui=True,
                     )
-                else:
+                elif (
+                    plot_type == "detailed"
+                    or cell.id in detailed_cells
+                    or plot_type == "constant"
+                    or cell.id in constant_cells
+                ):
                     pts, sizes, colors = plot_3D_cell_morphology(
                         offset=pos,
                         cell=cell,