feat(2d-morph): allow plotting cells in different levels of detail

pyneuroml/plot/PlotMorphology.py

@@ -172,6 +172,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 +208,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,6 +218,20 @@ 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 the network 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"]:
@@ -268,9 +286,40 @@ def plot_2D(
     fig, ax = get_new_matplotlib_morph_plot(title, plane2d)
     axis_min_max = [float("inf"), -1 * float("inf")]
 
+    # process plot_spec
+    point_cells = []  # type: typing.List[int]
+    schematic_cells = []  # type: typing.List[int]
+    const_cells = []  # type: typing.List[int]
+    detailed_cells = []  # type: typing.List[int]
+    if plot_spec is not None:
+        cellids = [k for k in cell_id_vs_cell.keys()]  # type: typing.List[str]
+        try:
+            point_cells = random.sample(
+                cellids, int(len(cellids) * plot_spec["point_fraction"])
+            )
+        except KeyError:
+            pass
+        # override with explicit list of point cells
+        try:
+            point_cells = plot_spec["point_cells"]
+        except KeyError:
+            pass
+        try:
+            schematic_cells = plot_spec["schematic_cells"]
+        except KeyError:
+            pass
+        try:
+            const_cells = plot_spec["constant_cells"]
+        except KeyError:
+            pass
+        try:
+            detailed_cells = plot_spec["detailed_cells"]
+        except KeyError:
+            pass
+
     for pop_id in pop_id_vs_cell:
-        cell = pop_id_vs_cell[pop_id]
-        pos_pop = positions[pop_id]
+        cell = pop_id_vs_cell[pop_id]  # type: Cell
+        pos_pop = positions[pop_id]  # type: typing.List[typing.Any]
 
         for cell_index in pos_pop:
             pos = pos_pop[cell_index]
@@ -291,12 +340,32 @@ def plot_2D(
                     nogui=True,
                 )
             else:
-                if plot_type == "schematic":
+                if plot_type == "point" 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 +375,7 @@ def plot_2D(
                         autoscale=False,
                         square=False,
                     )
-                else:
+                elif plot_type == "detailed" or cell.id in detailed_cells:
                     plot_2D_cell_morphology(
                         offset=pos,
                         cell=cell,

tests/plot/test_morphology_plot.py

@@ -115,6 +115,54 @@ def test_2d_morphology_plotter_data_overlay(self):
             self.assertIsFile(filename)
             pl.Path(filename).unlink()
 
+    def test_2d_plotter_network_with_spec(self):
+        """Test plot_2D function with a network of a few cells with specs."""
+        nml_file = "tests/plot/L23-example/TestNetwork.net.nml"
+        ofile = pl.Path(nml_file).name
+        # percentage
+        for plane in ["xy", "yz", "xz"]:
+            filename = f"test_morphology_plot_2d_spec_{ofile.replace('.', '_', 100)}_{plane}.png"
+            # remove the file first
+            try:
+                pl.Path(filename).unlink()
+            except FileNotFoundError:
+                pass
+
+            plot_2D(
+                nml_file,
+                nogui=True,
+                plane2d=plane,
+                save_to_file=filename,
+                plot_spec={"point_fraction": 0.5},
+            )
+
+            self.assertIsFile(filename)
+            pl.Path(filename).unlink()
+
+        # more detailed plot_spec
+        for plane in ["xy", "yz", "xz"]:
+            filename = f"test_morphology_plot_2d_spec_{ofile.replace('.', '_', 100)}_{plane}.png"
+            # remove the file first
+            try:
+                pl.Path(filename).unlink()
+            except FileNotFoundError:
+                pass
+
+            plot_2D(
+                nml_file,
+                nogui=True,
+                plane2d=plane,
+                save_to_file=filename,
+                plot_spec={
+                    "point_cells": ["HL23VIP"],
+                    "detailed_cells": ["HL23PYR"],
+                    "schematic_cells": ["HL23PV"],
+                    "constant_cells": ["HL23SST"],
+                },
+            )
+            self.assertIsFile(filename)
+            pl.Path(filename).unlink()
+
     def test_2d_plotter_network(self):
         """Test plot_2D function with a network of a few cells."""
         nml_file = "tests/plot/L23-example/TestNetwork.net.nml"
@@ -132,7 +180,7 @@ def test_2d_plotter_network(self):
             plot_2D(nml_file, nogui=True, plane2d=plane, save_to_file=filename)
 
             self.assertIsFile(filename)
-            pl.Path(filename).unlink()
+            # pl.Path(filename).unlink()
 
     def test_2d_constant_plotter_network(self):
         """Test plot_2D_schematic function with a network of a few cells."""