Feature/cc 434

frontend/packages/data-portal/app/components/Viewer/ViewerPage.tsx

@@ -3,6 +3,7 @@ import { SnackbarCloseReason } from '@mui/material/Snackbar'
 import {
   currentNeuroglancerState,
   NeuroglancerAwareIframe,
+  NeuroglancerState,
   NeuroglancerWrapper,
   updateState,
 } from 'neuroglancer'
@@ -19,6 +20,7 @@ import { useEffectOnce } from 'app/hooks/useEffectOnce'
 import { useI18n } from 'app/hooks/useI18n'
 import { useTour } from 'app/hooks/useTour'
 import { cns } from 'app/utils/cns'
+import { getTomogramName } from 'app/utils/tomograms'
 
 import { ReusableSnackbar } from '../common/ReusableSnackbar/ReusableSnackbar'
 import {
@@ -43,8 +45,12 @@ import {
   isCurrentLayout,
   isDepositionActivated,
   isDimensionPanelVisible,
+  isTomogramActivated,
+  isTomogramActivatedFromConfig,
   isTopBarVisible,
   panelsDefaultValues,
+  replaceOnlyTomogram,
+  replaceOnlyTomogramSource,
   resolveStateBool,
   setCurrentLayout,
   setTopBarVisibleFromSuperState,
@@ -68,13 +74,26 @@ import { Tour } from './Tour'
 import styles from './ViewerPage.module.css'
 
 type Run = GetRunByIdV2Query['runs'][number]
+type Tomograms = GetRunByIdV2Query['tomograms']
+type Tomogram = Tomograms[number]
 type Annotations = Run['annotations']
 type Annotation = Annotations['edges'][number]['node']
 interface AnnotationUIConfig {
   name?: string
   annotation: Annotation
 }
 
+const toZarr = (httpsMrcFile: string | undefined | null) => {
+  if (!httpsMrcFile) return httpsMrcFile
+  return `zarr://${httpsMrcFile.replace('.mrc', '.zarr')}`
+}
+
+const selectedTomogram = (tomogram: Tomogram) => {
+  return tomogram.neuroglancerConfig
+    ? isTomogramActivatedFromConfig(tomogram.neuroglancerConfig)
+    : isTomogramActivated(toZarr(tomogram.httpsMrcFile))
+}
+
 const buildDepositionsConfig = (
   annotations: Annotations,
 ): Record<number, AnnotationUIConfig[]> => {
@@ -118,9 +137,11 @@ const isSmallScreen = () => {
 
 export function ViewerPage({
   run,
+  tomograms,
   shouldStartTour = false,
 }: {
   run: Run
+  tomograms: Tomograms
   shouldStartTour?: boolean
 }) {
   const { t } = useI18n()
@@ -141,12 +162,15 @@ export function ViewerPage({
   const iframeRef = useRef<NeuroglancerAwareIframe>(null)
   const hashReady = useRef<boolean>(false)
   const helpMenuRef = useRef<MenuDropdownRef>(null)
+  const voxelSpacing = useRef<number>(0)
+  const alignmentId = useRef<number>(0)
 
   const shareSnackbar = useAutoHideSnackbar()
   const snapSnackbar = useAutoHideSnackbar()
 
   const depositionConfigs = buildDepositionsConfig(run.annotations)
   const shouldShowAnnotationDropdown = Object.keys(depositionConfigs).length > 0
+  const shouldShowTomogramDropdown = tomograms.length > 1
 
   const scheduleRefresh = () => {
     setRenderVersion(renderVersion + 1)
@@ -202,6 +226,14 @@ export function ViewerPage({
     }
     hashReady.current = true
 
+    const currentlyActiveTomogram = tomograms.find(
+      (tomogram) =>
+        isTomogramActivatedFromConfig(tomogram.neuroglancerConfig) ||
+        isTomogramActivated(toZarr(tomogram.httpsMrcFile)),
+    )
+    voxelSpacing.current = currentlyActiveTomogram?.voxelSpacing || 0
+    alignmentId.current = currentlyActiveTomogram?.alignment?.id || 0
+
     window.addEventListener('keydown', keyDownHandler)
     setTourRunning(shouldStartTour)
     return () => {
@@ -215,6 +247,15 @@ export function ViewerPage({
     }
   }, [tourRunning])
 
+  const unsupportedTomogramSwitch = (tomogram: Tomogram) => {
+    const hasFullState = !!tomogram.neuroglancerConfig
+    const hasSourceInSameSpace =
+      !!tomogram.s3OmezarrDir &&
+      voxelSpacing.current === tomogram.voxelSpacing &&
+      alignmentId.current === (tomogram.alignment?.id || 0)
+    return !(hasFullState || hasSourceInSameSpace)
+  }
+
   const handleOnStateChange = (state: ViewerPageSuperState) => {
     scheduleRefresh()
     setTopBarVisibleFromSuperState()
@@ -347,6 +388,54 @@ export function ViewerPage({
         <div className="basis-sds-xxl flex-grow md:mr-sds-xxl" />
         <div className="flex basis-auto flex-shrink-0">
           <div className="flex items-center pt-1 gap-[1px] sm:gap-1 sm:pt-0">
+            {shouldShowTomogramDropdown && (
+              <NeuroglancerDropdown title="Tomograms" variant="outlined">
+                {tomograms.map((tomogram) => {
+                  return (
+                    <NeuroglancerDropdownOption
+                      key={tomogram.id.toString()}
+                      selected={selectedTomogram(tomogram)}
+                      disabled={unsupportedTomogramSwitch(tomogram)}
+                      onSelect={() => {
+                        if (selectedTomogram(tomogram)) return
+                        voxelSpacing.current = tomogram.voxelSpacing
+                        alignmentId.current = tomogram.alignment?.id || 0
+                        updateState((state) => {
+                          return {
+                            ...state,
+                            neuroglancer: tomogram.neuroglancerConfig
+                              ? replaceOnlyTomogram(
+                                  state.neuroglancer,
+                                  JSON.parse(
+                                    tomogram.neuroglancerConfig,
+                                  ) as NeuroglancerState,
+                                )
+                              : replaceOnlyTomogramSource(
+                                  state.neuroglancer,
+                                  toZarr(tomogram.httpsMrcFile)!,
+                                ),
+                          }
+                        })
+                      }}
+                    >
+                      <span className="line-clamp-3">
+                        {getTomogramName(tomogram)}
+                      </span>
+                      <span className="text-sds-body-xxxs-400-narrow text-light-sds-color-primitive-gray-600">
+                        {[
+                          `${IdPrefix.Tomogram}-${tomogram.id}`,
+                          `${t('unitAngstrom', { value: tomogram.voxelSpacing })} (${tomogram.sizeX}, ${tomogram.sizeY}, ${tomogram.sizeZ}) px`,
+                          tomogram.alignment?.id != null &&
+                            `${IdPrefix.Alignment}-${tomogram.alignment.id}`,
+                        ]
+                          .filter(Boolean)
+                          .join(' · ')}
+                      </span>
+                    </NeuroglancerDropdownOption>
+                  )
+                })}
+              </NeuroglancerDropdown>
+            )}
             {shouldShowAnnotationDropdown && (
               <NeuroglancerDropdown title="Annotations" variant="outlined">
                 <MenuDropdownSection title="Show annotations for deposition">

frontend/packages/data-portal/app/components/Viewer/state.ts

@@ -5,7 +5,10 @@ import {
   currentNeuroglancer,
   currentNeuroglancerState,
   currentState,
+  DimensionValue,
+  getLayerSourceUrl,
   NeuroglancerLayout,
+  NeuroglancerState,
   ResolvedSuperState,
   updateState,
 } from 'neuroglancer'
@@ -301,6 +304,123 @@ export function toggleOrMakeDimensionPanel() {
   else updateState(toggleDimensionPanelVisible)
 }
 
+export function isTomogramActivatedFromConfig(
+  tomogramConfig: string | undefined | null,
+) {
+  if (!tomogramConfig) return false
+  const layers = currentNeuroglancerState().layers || []
+  const jsonConfig = JSON.parse(tomogramConfig) as NeuroglancerState
+  const newLayers = jsonConfig.layers || []
+  const tomogramLayer = newLayers.find((l) => l.type === 'image')
+  if (!tomogramLayer) return false
+  return layers.some(
+    (l) =>
+      l.type === 'image' &&
+      getLayerSourceUrl(l) === getLayerSourceUrl(tomogramLayer),
+  )
+}
+
+export function isTomogramActivated(tomogramPath: string | undefined | null) {
+  if (!tomogramPath) return false
+  const layers = currentNeuroglancerState().layers || []
+  return layers.some(
+    (l) => l.type === 'image' && getLayerSourceUrl(l) === tomogramPath,
+  )
+}
+
+function inferVoxelSpacingFromState(state: NeuroglancerState) {
+  const { dimensions } = state
+  if (dimensions === undefined) {
+    throw new Error('Cannot infer voxel spacing without dimensions')
+  }
+  // Get the average of all dims, usually isotropic but just in case
+  const dimensionValues = Object.values(dimensions)
+  const averageUnit =
+    dimensionValues.reduce((a: number, b: DimensionValue) => a + b[0], 0) /
+    dimensionValues.length
+  return averageUnit
+}
+
+export function replaceOnlyTomogramSource(
+  incomingState: NeuroglancerState,
+  newPath: string,
+) {
+  const newState = incomingState
+  const tomogramLayer = newState.layers?.find((l) => l.type === 'image')
+  if (tomogramLayer) {
+    // Replace either the source directly or the url inside the source object
+    if (typeof tomogramLayer.source === 'string') {
+      tomogramLayer.source = newPath
+    } else {
+      tomogramLayer.source.url = newPath
+    }
+  }
+  return newState
+}
+
+export function replaceOnlyTomogram(
+  incomingState: NeuroglancerState,
+  newState: NeuroglancerState,
+) {
+  // The first image layer is always the tomogram -- we can completely replace that layer
+  // For the other layers, we only need to adjust the "source" because they can have
+  // different transforms needed
+  if (!newState.layers) return incomingState
+  const incomingLayers = newState.layers
+  const newTomogramLayer = incomingLayers.find((l) => l.type === 'image')
+  if (!newTomogramLayer) return incomingState // No tomogram layer in the new state
+  newTomogramLayer.visible = true
+  newTomogramLayer.archived = false
+  const newLayers = incomingState.layers || []
+
+  // First, let's check for the tomogram layer in the current state
+  const tomogramLayerIndex = newLayers.findIndex((l) => l.type === 'image')
+  if (tomogramLayerIndex === -1) {
+    newLayers.unshift(newTomogramLayer)
+  } else {
+    const currentTomogram = newLayers[tomogramLayerIndex]
+    const openedTab = currentTomogram.tab
+    if (openedTab) {
+      newTomogramLayer.tab = openedTab
+    }
+    newLayers[tomogramLayerIndex] = newTomogramLayer
+  }
+
+  // For the other layers, we need to update their sources if they exist in both states
+  for (const newLayer of incomingLayers) {
+    if (newLayer.type === 'image') continue // Skip the tomogram layer
+    const matchingLayer = newLayers.find(
+      (l) => getLayerSourceUrl(l) === getLayerSourceUrl(newLayer),
+    )
+    if (matchingLayer) {
+      matchingLayer.source = newLayer.source
+    } else {
+      newLayers.push(newLayer)
+    }
+  }
+
+  // Adjust the zoom levels and position to keep view consistent when switching
+  const currentSpacing = inferVoxelSpacingFromState(incomingState)
+  const newSpacing = inferVoxelSpacingFromState(newState)
+  const voxelRatio = newSpacing / currentSpacing
+  const newCrossSectionScale = incomingState.crossSectionScale
+    ? incomingState.crossSectionScale * voxelRatio
+    : undefined
+  const newProjectionScale = incomingState.projectionScale
+    ? incomingState.projectionScale * voxelRatio
+    : undefined
+  const newPosition = incomingState.position
+    ? incomingState.position.map((x) => x * voxelRatio)
+    : undefined
+  return {
+    ...incomingState,
+    layers: newLayers,
+    crossSectionScale: newCrossSectionScale,
+    projectionScale: newProjectionScale,
+    position: newPosition,
+  }
+}
+
 export function isDepositionActivated(
   depositionEntries: (string | undefined)[],
 ) {

frontend/packages/data-portal/app/routes/view.runs.$id.tsx

@@ -54,13 +54,17 @@ const ViewerPage = lazy(() =>
 )
 
 export default function RunByIdViewerPage() {
-  const { run } = useRunById()
+  const { run, tomograms } = useRunById()
   const [searchParams] = useSearchParams()
   const shouldStartTour = searchParams.get(QueryParams.ShowTour) === 'true'
 
   return (
     <Suspense fallback={<div>Loading...</div>}>
-      <ViewerPage run={run} shouldStartTour={shouldStartTour} />
+      <ViewerPage
+        run={run}
+        tomograms={tomograms}
+        shouldStartTour={shouldStartTour}
+      />
     </Suspense>
   )
 }

frontend/packages/neuroglancer/src/utils.ts

@@ -18,8 +18,11 @@ export interface NeuroglancerState
   selection?: PanelState
   toolPalettes?: Record<string, ToolPaletteState>
   layers?: LayerWithSource[]
+  dimensions?: { [key: string]: DimensionValue }
 }
 
+export type DimensionValue = [number, string]
+
 export interface SuperState extends Record<string, unknown> {
   neuroglancer: string
 }
@@ -79,14 +82,27 @@ interface ToolPaletteState extends PanelState {
 }
 
 interface LayerWithSource extends LayerElement {
-  source: string | { url?: string }
+  source:
+    | string
+    | {
+        url: string
+        transform?: { outputDimensions: unknown; inputDimensions: unknown }
+      }
   archived?: boolean
+  tab: string
 }
 
 interface WatchableBoolean {
   value: boolean
 }
 
+export function getLayerSourceUrl(layer: LayerWithSource): string {
+  if (typeof layer.source === 'string') {
+    return layer.source
+  }
+  return layer.source.url
+}
+
 const emptySuperState = (config: string): SuperState => {
   return {
     neuroglancer: config.length > 0 ? decompressHash(config) : '',