Feature Draft/ Use PCA to visualize cell 'upwards'

pyneuroml/plot/PlotMorphologyVispy.py

@@ -22,7 +22,7 @@
 from scipy.spatial.transform import Rotation
 
 from pyneuroml.pynml import read_neuroml2_file
-from pyneuroml.utils import extract_position_info
+from pyneuroml.utils import extract_position_info, rotate_cell, translate_cell_to_coords
 from pyneuroml.utils.plot import (
     DEFAULTS,
     get_cell_bound_box,
@@ -105,6 +105,7 @@ def create_new_vispy_canvas(
     axes_length: float = 100,
     axes_width: int = 2,
     theme=PYNEUROML_VISPY_THEME,
+    view_center: typing.Optional[typing.List[float]] = None,
 ):
     """Create a new vispy scene canvas with a view and optional axes lines
 
@@ -114,9 +115,11 @@ def create_new_vispy_canvas(
     :type view_min: [float, float, float]
     :param view_max: max view co-ordinates
     :type view_max: [float, float, float]
+    :param view_center: center view co-ordinates, calculated from view max/min if omitted
+    :type view_center: [float, float, float]
     :param title: title of plot
     :type title: str
-    :param axes_pos: position to draw axes at
+    :param axes_pos: None by default: axes omitted, otherwise axes centered at given position with colours red, green, blue for x,y,z axis respecitvely
     :type axes_pos: [float, float, float]
     :param axes_length: length of axes
     :type axes_length: float
@@ -144,13 +147,13 @@ def create_new_vispy_canvas(
 
     # create cameras
     # https://vispy.org/gallery/scene/flipped_axis.html
-    cam1 = scene.cameras.PanZoomCamera(parent=view.scene, name="PanZoom")
+    cam1 = scene.cameras.PanZoomCamera(parent=view.scene, name="PanZoom", up="y")
 
-    cam2 = scene.cameras.TurntableCamera(parent=view.scene, name="Turntable")
+    cam2 = scene.cameras.TurntableCamera(parent=view.scene, name="Turntable", up="y")
 
-    cam3 = scene.cameras.ArcballCamera(parent=view.scene, name="Arcball")
+    cam3 = scene.cameras.ArcballCamera(parent=view.scene, name="Arcball", up="y")
 
-    cam4 = scene.cameras.FlyCamera(parent=view.scene, name="Fly")
+    cam4 = scene.cameras.FlyCamera(parent=view.scene, name="Fly", up="y")
     # do not keep z up
     cam4.autoroll = False
 
@@ -161,13 +164,6 @@ def create_new_vispy_canvas(
     view.camera = cams[cam_index]
 
     if view_min is not None and view_max is not None:
-        view_center = (numpy.array(view_max) + numpy.array(view_min)) / 2
-        logger.debug(f"Center is {view_center}")
-        cam1.center = [view_center[0], view_center[1]]
-        cam2.center = view_center
-        cam3.center = view_center
-        cam4.center = view_center
-
         for acam in cams:
             x_width = abs(view_min[0] - view_max[0])
             y_width = abs(view_min[1] - view_max[1])
@@ -192,6 +188,15 @@ def create_new_vispy_canvas(
 
             acam.set_range(x=xrange, y=yrange, z=zrange)
 
+        # Calculate view center if it is None
+        if view_center is None:
+            view_center = (numpy.array(view_max) + numpy.array(view_min)) / 2
+        logger.debug(f"Center is {view_center}")
+        cam1.center = [view_center[0], view_center[1]]
+        cam2.center = view_center
+        cam3.center = view_center
+        cam4.center = view_center
+
     for acam in cams:
         acam.set_default_state()
 
@@ -204,10 +209,21 @@ def create_new_vispy_canvas(
             [axes_pos[0], axes_pos[1], axes_pos[2] + axes_length],
         ]
         scene.Line(
-            points,
-            connect=numpy.array([[0, 1], [0, 2], [0, 3]]),
+            [points[0], points[1]],
             parent=view.scene,
-            color=VISPY_THEME[theme]["fg"],
+            color="red",
+            width=axes_width,
+        )
+        scene.Line(
+            [points[0], points[2]],
+            parent=view.scene,
+            color="green",
+            width=axes_width,
+        )
+        scene.Line(
+            [points[0], points[3]],
+            parent=view.scene,
+            color="blue",
             width=axes_width,
         )
 
@@ -260,6 +276,7 @@ def plot_interactive_3D(
     ] = None,
     highlight_spec: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
     precision: typing.Tuple[int, int] = (4, 200),
+    upright: bool = False,
 ):
     """Plot interactive plots in 3D using Vispy
 
@@ -356,6 +373,11 @@ def plot_interactive_3D(
         If you have a good GPU, you can increase both these values to get more
         detailed visualizations
     :type precision: (int, int)
+    :param upright: bool only applicable for single cells: Makes cells "upright"
+    (along Y axis) by calculating it PCA and transforming cell co-ordinates to
+    align along the first principal component. Note that the original cell object
+    is unchanged, this is for visualization purposes only
+    :type upright: bool
     """
     if plot_type not in ["detailed", "constant", "schematic", "point"]:
         raise ValueError(
@@ -440,6 +462,8 @@ def plot_interactive_3D(
     # not used later, clear up
     del cell_id_vs_cell
 
+    singleCell = False
+
     if len(positions) > 1:
         only_pos = []
         for posdict in positions.values():
@@ -471,7 +495,9 @@ def plot_interactive_3D(
         logger.debug(f"center, view_min, max are {center}, {view_min}, {view_max}")
 
     else:
+        singleCell = True
         cell = list(pop_id_vs_cell.values())[0]
+
         if cell is not None:
             view_min, view_max = get_cell_bound_box(cell)
         else:
@@ -480,8 +506,17 @@ def plot_interactive_3D(
             view_min = list(numpy.array(pos))
             view_max = list(numpy.array(pos))
 
+    if singleCell and upright:
+        view_center = [0, 0, 0]
+    else:
+        view_center = None
     current_canvas, current_view = create_new_vispy_canvas(
-        view_min, view_max, title, theme=theme
+        view_min,
+        view_max,
+        title,
+        axes_pos=[0, 0, 0],
+        theme=theme,
+        view_center=view_center,
     )
 
     logger.debug(f"figure extents are: {view_min}, {view_max}")
@@ -587,6 +622,7 @@ def plot_interactive_3D(
                         nogui=True,
                         meshdata=meshdata,
                         mesh_precision=precision[0],
+                        upright=upright,
                     )
                 elif (
                     plot_type == "detailed"
@@ -614,6 +650,7 @@ def plot_interactive_3D(
                         meshdata=meshdata,
                         mesh_precision=precision[0],
                         highlight_spec=cell_highlight_spec,
+                        upright=upright,
                     )
 
             # if too many meshes, reduce precision and retry, recursively
@@ -634,6 +671,7 @@ def plot_interactive_3D(
                     plot_spec=plot_spec,
                     precision=precision,
                     highlight_spec=highlight_spec,
+                    upright=upright,
                 )
                 # break the recursion, don't plot in the calling method
                 return
@@ -647,6 +685,64 @@ def plot_interactive_3D(
         app.run()
 
 
+def make_cell_upright(
+    cell: Cell = None,
+    inplace: bool = False,
+) -> Cell:
+    """Use cell's PCA to make it upright
+
+    .. versionadded:: 1.2.13
+
+    :param cell: cell object to translate
+    :type cell: neuroml.Cell
+    :param inplace: toggle whether the cell object should be modified inplace
+        or a copy created (creates and returns a copy by default)
+    :type inplace: bool
+    :returns: new neuroml.Cell object
+    :rtype: neuroml.Cell
+    """
+
+    # Get all segments' distal points
+    segment_points = []
+    segments_all = cell.morphology.segments
+    for segment in segments_all:
+        segment_points.append([segment.distal.x, segment.distal.y, segment.distal.z])
+
+    coords = numpy.array(segment_points)
+    from sklearn.decomposition import PCA
+
+    # Get the PCA components
+    pca = PCA()
+    pca.fit(coords)
+
+    # Get the principal component axes
+    principal_axes = pca.components_
+    # Get the first principal component axis
+    first_pca = principal_axes[0]
+    # y angle needed to eliminate z component
+    y_angle = math.atan(first_pca[2] / first_pca[0])
+    rotation_y = numpy.array(
+        [
+            [math.cos(y_angle), 0, math.sin(y_angle)],
+            [0, 1, 0],
+            [-math.sin(y_angle), 0, math.cos(y_angle)],
+        ]
+    )
+    rotated_pca = numpy.dot(rotation_y, first_pca)
+
+    # z angle needed to eliminate x component
+    z_angle = -math.atan(rotated_pca[0] / rotated_pca[1])
+
+    if z_angle < 0:
+        z_angle += numpy.pi
+
+    cell = translate_cell_to_coords(cell, inplace=inplace, dest=[0, 0, 0])
+    cell = rotate_cell(
+        cell, 0, y_angle, z_angle, "yzx", relative_to_soma=False, inplace=inplace
+    )
+    return cell
+
+
 def plot_3D_cell_morphology(
     offset: typing.List[float] = [0, 0, 0],
     cell: Optional[Cell] = None,
@@ -663,6 +759,7 @@ def plot_3D_cell_morphology(
     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,
+    upright: bool = False,
 ):
     """Plot the detailed 3D morphology of a cell using vispy.
     https://vispy.org/
@@ -745,6 +842,11 @@ def plot_3D_cell_morphology(
           is used)
 
     :type highlight_spec: dict
+    :param upright: bool only applicable for single cells: Makes cells "upright"
+    (along Y axis) by calculating it PCA and transforming cell co-ordinates to
+    align along the first principal component. Note that the original cell object
+    is unchanged, this is for visualization purposes only
+    :type upright: bool
     :raises: ValueError if `cell` is None
 
     """
@@ -769,11 +871,16 @@ def plot_3D_cell_morphology(
         axon_segs = cell.get_all_segments_in_group("axon_group")
     except Exception:
         axon_segs = []
+    view_center = None
+    if upright:
+        cell = make_cell_upright(cell)
+        current_canvas = current_view = None
+        view_center = [0, 0, 0]
 
     if current_canvas is None or current_view is None:
         view_min, view_max = get_cell_bound_box(cell)
         current_canvas, current_view = create_new_vispy_canvas(
-            view_min, view_max, title, theme=theme
+            view_min, view_max, title, theme=theme, view_center=view_center
         )
 
     if color == "Groups":
@@ -1015,6 +1122,7 @@ def plot_3D_schematic(
     color: typing.Optional[str] = "Cell",
     meshdata: typing.Optional[typing.Dict[typing.Any, typing.Any]] = None,
     mesh_precision: int = 2,
+    upright: bool = False,
 ) -> None:
     """Plot a 3D schematic of the provided segment groups using vispy.
     layer..
@@ -1069,10 +1177,17 @@ def plot_3D_schematic(
           dendrites, and soma segments
 
     :type color: str
+    :param upright: bool only applicable for single cells: Makes cells "upright"
+    (along Y axis) by calculating it PCA and transforming cell co-ordinates to
+    align along the first principal component. Note that the original cell object
+    is unchanged, this is for visualization purposes only
+    :type upright: bool
     """
     if title == "":
         title = f"3D schematic of segment groups from {cell.id}"
 
+    if upright:
+        cell = make_cell_upright(cell)
     try:
         soma_segs = cell.get_all_segments_in_group("soma_group")
     except Exception:
@@ -1097,11 +1212,16 @@ def plot_3D_schematic(
         segment_groups, check_parentage=False
     )
 
-    # if no canvas is defined, define a new one
+    view_center = None
+    if upright:
+        cell = make_cell_upright(cell)
+        current_canvas = current_view = None
+        view_center = [0, 0, 0]
+
     if current_canvas is None or current_view is None:
         view_min, view_max = get_cell_bound_box(cell)
         current_canvas, current_view = create_new_vispy_canvas(
-            view_min, view_max, title, theme=theme
+            view_min, view_max, title, theme=theme, view_center=view_center
         )
 
     # colors for cell

setup.cfg

@@ -97,6 +97,7 @@ vispy-common =
     scipy
     pyopengl
     PyOpenGL-accelerate
+    scikit-learn
 
 vispy-qt5 =
     pyNeuroML[vispy-common]