CELE-107 Fix last issues with selection

MetaCell · Nov 5, 2024 · 1afe2eb · 1afe2eb
1 parent 38afec8
commit 1afe2eb
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 24 deletions.
diff --git a/applications/visualizer/frontend/src/components/viewers/EM/EMStackTilesViewer.tsx b/applications/visualizer/frontend/src/components/viewers/EM/EMStackTilesViewer.tsx
@@ -51,8 +51,7 @@ function isNeuronActive(neuronId: string, workspace: Workspace): boolean {
 }
 
 function isNeuronSelected(neuronId: string, workspace: Workspace): boolean {
-  const emViewerSelectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.EM);
-  return emViewerSelectedNeurons.includes(neuronId) || emViewerSelectedNeurons.includes(workspace.getNeuronClass(neuronId));
+  return workspace.getSelection(ViewerType.EM).includes(neuronId);
 }
 
 function isNeuronVisible(neuronId: string, workspace: Workspace): boolean {
@@ -89,12 +88,19 @@ function onNeuronSelect(position: Coordinate, source: VectorSource<Feature> | un
   const feature = features[0];
   const neuronName = neuronFeatureName(feature);
 
-  if (!isNeuronVisible(neuronName, workspace)) {
+  if (isNeuronSelected(neuronName, workspace)) {
+    workspace.removeSelection(neuronName, ViewerType.EM);
+    // Is there a neuron in the selection that comes from the same class. If not, we can remove the class from the selection
+    const removeClass = !workspace
+      .getSelection(ViewerType.ThreeD)
+      .some((e) => workspace.getNeuronClass(e) !== e && workspace.getNeuronClass(e) === workspace.getNeuronClass(neuronName));
+    if (removeClass) {
+      workspace.removeSelection(workspace.getNeuronClass(neuronName), ViewerType.ThreeD);
+    }
     return;
   }
 
-  if (isNeuronSelected(neuronName, workspace)) {
-    workspace.removeSelection(neuronName, ViewerType.EM);
+  if (!isNeuronVisible(neuronName, workspace)) {
     return;
   }
 
@@ -175,7 +181,7 @@ const EMStackViewer = () => {
     neuronsStyleRef.current = (feature: Feature) => neuronsStyle(feature, currentWorkspace);
     onNeuronSelectRef.current = (position) => onNeuronSelect(position, currSegLayer.current.getSource(), currentWorkspace);
     currSegLayer.current.getSource().changed();
-  }, [currentWorkspace.getVisibleNeuronsInEM(), currentWorkspace.visibilities, currentWorkspace.getViewerSelectedNeurons(ViewerType.EM), segSlice]);
+  }, [currentWorkspace.getVisibleNeuronsInEM(), currentWorkspace.visibilities, currentWorkspace.getSelection(ViewerType.EM), segSlice]);
 
   useEffect(() => {
     if (mapRef.current) {

diff --git a/applications/visualizer/frontend/src/components/viewers/ThreeD/STLMesh.tsx b/applications/visualizer/frontend/src/components/viewers/ThreeD/STLMesh.tsx
@@ -1,6 +1,6 @@
 import { Outlines } from "@react-three/drei";
 import type { ThreeEvent } from "@react-three/fiber";
-import { useCallback, type FC } from "react";
+import { type FC, useCallback, useMemo } from "react";
 import { useSelector } from "react-redux";
 import { type BufferGeometry, DoubleSide, NormalBlending } from "three";
 import { useGlobalContext } from "../../../contexts/GlobalContext";
@@ -23,16 +23,29 @@ const STLMesh: FC<Props> = ({ id, color, opacity, renderOrder, isWireframe, stl
   const { workspaces } = useGlobalContext();
   const workspaceId = useSelector((state: RootState) => state.workspaceId);
   const workspace: Workspace = workspaces[workspaceId];
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.ThreeD);
-  const isSelected = selectedNeurons.includes(id);
+
+  const isSelected = useMemo(() => {
+    const selectedNeurons = workspace.getSelection(ViewerType.ThreeD);
+    return selectedNeurons.includes(id);
+  }, [workspace.getSelection(ViewerType.ThreeD)]);
 
   const onClick = useCallback(
     (event: ThreeEvent<MouseEvent>) => {
       const clicked = getFurthestIntersectedObject(event);
+      if (!clicked) {
+        return;
+      }
       const { id } = clicked.userData;
       if (clicked) {
         if (isSelected) {
           workspace.removeSelection(id, ViewerType.ThreeD);
+          // Is there a neuron in the selection that comes from the same class. If not, we can remove the class from the selection
+          const removeClass = !workspace
+            .getSelection(ViewerType.ThreeD)
+            .some((e) => workspace.getNeuronClass(e) !== e && workspace.getNeuronClass(e) === workspace.getNeuronClass(id));
+          if (removeClass) {
+            workspace.removeSelection(workspace.getNeuronClass(id), ViewerType.ThreeD);
+          }
         } else {
           workspace.addSelection(id, ViewerType.ThreeD);
         }

diff --git a/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx b/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx
@@ -38,7 +38,7 @@ interface ContextMenuProps {
 const ContextMenu: React.FC<ContextMenuProps> = ({ open, onClose, position, setSplitJoinState, openGroups, setOpenGroups, cy }) => {
   const workspace = useSelectedWorkspace();
   const [submenuAnchorEl, setSubmenuAnchorEl] = useState<null | HTMLElement>(null);
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+  const selectedNeurons = workspace.getSelection(ViewerType.Graph);
 
   const submenuOpen = Boolean(submenuAnchorEl);
 

diff --git a/applications/visualizer/frontend/src/components/viewers/TwoD/TwoDViewer.tsx b/applications/visualizer/frontend/src/components/viewers/TwoD/TwoDViewer.tsx
@@ -70,7 +70,7 @@ const TwoDViewer = () => {
     unreportedNeurons: new Set<string>(),
   });
 
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+  const selectedNeurons = workspace.getSelection(ViewerType.Graph);
 
   const visibleActiveNeurons = useMemo(() => {
     return getVisibleActiveNeuronsIn2D(workspace);
@@ -236,23 +236,47 @@ const TwoDViewer = () => {
     };
   }, []);
 
+  useEffect(() => {
+    for (const node of cyRef.current.nodes()) {
+      const neuronId = node.id();
+      const isSelected = selectedNeurons.includes(neuronId) || selectedNeurons.some((e) => workspace.getNeuronCellsByClass(neuronId).includes(e));
+
+      if (isSelected) {
+        node.addClass(SELECTED_CLASS);
+      } else {
+        node.removeClass(SELECTED_CLASS);
+      }
+    }
+  }, [selectedNeurons]);
+
   // Add event listener for node clicks to toggle neuron selection and right-click context menu
   useEffect(() => {
     if (!cyRef.current) return;
 
     const cy = cyRef.current;
 
     const handleNodeClick = (event) => {
-      const neuronId = event.target.id();
-      const selectedNeurons = workspace.getSelection(ViewerType.Graph);
-      const isSelected = selectedNeurons.includes(neuronId);
+      const node = event.target;
+      const neuronId = node.id();
+      const isSelected = selectedNeurons.includes(neuronId) || selectedNeurons.some((e) => workspace.getNeuronCellsByClass(neuronId).includes(e));
 
       if (isSelected) {
         workspace.removeSelection(neuronId, ViewerType.Graph);
-        event.target.removeClass(SELECTED_CLASS);
+
+        if (workspace.getNeuronClass(neuronId) === neuronId) {
+          const relatedNeurons = workspace.getNeuronCellsByClass(neuronId);
+          for (const neuron of relatedNeurons) {
+            workspace.forceRemoveSelection(neuron, ViewerType.EM);
+            workspace.forceRemoveSelection(neuron, ViewerType.ThreeD);
+          }
+        }
       } else {
         workspace.addSelection(neuronId, ViewerType.Graph);
-        event.target.addClass(SELECTED_CLASS);
+        const relatedNeurons = workspace.getNeuronCellsByClass(neuronId);
+        for (const neuron of relatedNeurons) {
+          workspace.forceInjectSelection(neuron, ViewerType.EM);
+          workspace.forceInjectSelection(neuron, ViewerType.ThreeD);
+        }
       }
     };
 
@@ -270,7 +294,7 @@ const TwoDViewer = () => {
 
       const cyEvent = event as any; // Cast to any to access originalEvent
       const originalEvent = cyEvent.originalEvent as MouseEvent;
-      const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+      const selectedNeurons = workspace.getSelection(ViewerType.Graph);
       if (selectedNeurons.length > 0) {
         setMousePosition({
           mouseX: originalEvent.clientX,

diff --git a/applications/visualizer/frontend/src/helpers/twoD/graphRendering.ts b/applications/visualizer/frontend/src/helpers/twoD/graphRendering.ts
@@ -27,7 +27,7 @@ export const computeGraphDifferences = (
   includePostEmbryonic: boolean,
 ) => {
   const visibleActiveNeurons = getVisibleActiveNeuronsIn2D(workspace);
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+  const selectedNeurons = workspace.getSelection(ViewerType.Graph);
 
   // Current nodes and edges in the Cytoscape instance
   const currentNodes = new Set(cy.nodes().map((node) => node.id()));

diff --git a/applications/visualizer/frontend/src/helpers/twoD/splitJoinHelper.ts b/applications/visualizer/frontend/src/helpers/twoD/splitJoinHelper.ts
@@ -10,7 +10,7 @@ interface SplitJoinState {
 export const processNeuronSplit = (workspace: Workspace, splitJoinState: SplitJoinState): SplitJoinState => {
   const newSplit = new Set(splitJoinState.split);
   const newJoin = new Set(splitJoinState.join);
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+  const selectedNeurons = workspace.getSelection(ViewerType.Graph);
 
   const newSelectedNeurons = new Set(selectedNeurons);
   const graphViewDataUpdates: Record<string, Partial<GraphViewerData>> = {};
@@ -83,7 +83,7 @@ export const processNeuronSplit = (workspace: Workspace, splitJoinState: SplitJo
 export const processNeuronJoin = (workspace: Workspace, splitJoinState: SplitJoinState): SplitJoinState => {
   const newJoin = new Set(splitJoinState.join);
   const newSplit = new Set(splitJoinState.split);
-  const selectedNeurons = workspace.getViewerSelectedNeurons(ViewerType.Graph);
+  const selectedNeurons = workspace.getSelection(ViewerType.Graph);
 
   const newSelectedNeurons = new Set(selectedNeurons);
   const graphViewDataUpdates: Record<string, Partial<GraphViewerData>> = {};

diff --git a/applications/visualizer/frontend/src/models/synchronizer.ts b/applications/visualizer/frontend/src/models/synchronizer.ts
@@ -155,6 +155,18 @@ export class SynchronizerOrchestrator {
     }
   }
 
+  public forceInjectSelection(selection: string, target: ViewerType) {
+    this.contexts[target].push(selection);
+  }
+
+  public forceRemoveSelection(selection: string, target: ViewerType) {
+    const selected = this.contexts[target];
+    const index = selected.indexOf(selection);
+    if (index > -1) {
+      selected.splice(index, 1);
+    }
+  }
+
   public selectNeuron(selection: string, initiator: ViewerType) {
     const synchronizers = this.getConnectedViewers(initiator);
     for (const synchronizer of synchronizers) {

diff --git a/applications/visualizer/frontend/src/models/workspace.ts b/applications/visualizer/frontend/src/models/workspace.ts
@@ -220,6 +220,16 @@ export class Workspace {
     this.updateContext(updated);
   }
 
+  @triggerUpdate
+  forceInjectSelection(selection: string, target: ViewerType) {
+    this.syncOrchestrator.forceInjectSelection(selection, target);
+  }
+
+  @triggerUpdate
+  forceRemoveSelection(selection: string, target: ViewerType) {
+    this.syncOrchestrator.forceRemoveSelection(selection, target);
+  }
+
   @triggerUpdate
   setSelection(selection: Array<string>, initiator: ViewerType) {
     this.syncOrchestrator.select(selection, initiator);
@@ -245,10 +255,6 @@ export class Workspace {
     return this.syncOrchestrator.getSelection(viewerType);
   }
 
-  getViewerSelectedNeurons(viewerType: ViewerType): string[] {
-    return this.syncOrchestrator.getSelection(viewerType);
-  }
-
   getNeuronCellsByClass(neuronClassId: string): string[] {
     return Object.values(this.availableNeurons)
       .filter((neuron) => neuron.nclass === neuronClassId && neuron.nclass !== neuron.name)