feat: Fix/clean up tomogram_stats query (#1033)

#993 - `object_name` and `annotation_software` are now distinct in the GQL response. - `object_shape` is now guaranteed distinct across `tomogram_voxel_spacings` in the GQL response. - The `tomogram`s query used for downloads is now separate from the query for distinct `voxel_spacing`s.
chanzuckerberg · Aug 20, 2024 · 2d54a0c · 2d54a0c
1 parent 3d8c0d5
commit 2d54a0c
Show file tree

Hide file tree

Showing 10 changed files with 107 additions and 130 deletions.
diff --git a/frontend/packages/data-portal/app/components/Download/ConfigureTomogramDownloadContent.tsx b/frontend/packages/data-portal/app/components/Download/ConfigureTomogramDownloadContent.tsx
@@ -32,17 +32,17 @@ export function ConfigureTomogramDownloadContent() {
 
   const {
     allTomogramProcessing = [],
-    allTomogramResolutions = [],
+    allTomograms = [],
     runId,
   } = useDownloadModalContext()
 
   const tomogramSamplingOptions = useMemo<SelectOption[]>(
     () =>
-      allTomogramResolutions.map((tomogram) => ({
+      allTomograms.map((tomogram) => ({
         key: t('unitAngstrom', { value: tomogram.voxel_spacing }),
         value: `(${tomogram.size_x}, ${tomogram.size_y}, ${tomogram.size_z})px`,
       })),
-    [allTomogramResolutions, t],
+    [allTomograms, t],
   )
 
   const tomogramProcessingOptions = useMemo<SelectOption[]>(
@@ -54,18 +54,18 @@ export function ConfigureTomogramDownloadContent() {
     [allTomogramProcessing],
   )
 
-  const activeTomogram = allTomogramResolutions.find(
+  const activeTomogram = allTomograms.find(
     (tomogram) => `${tomogram.voxel_spacing}` === tomogramSampling,
   )
 
   const setTomogramConfigWithInitialValues = useCallback(() => {
-    const tomogram = allTomogramResolutions.at(0)
+    const tomogram = allTomograms.at(0)
     const processing = allTomogramProcessing.at(0)
 
     if (tomogram && processing) {
       setTomogramConfig(`${tomogram.voxel_spacing}`, processing)
     }
-  }, [allTomogramProcessing, allTomogramResolutions, setTomogramConfig])
+  }, [allTomogramProcessing, allTomograms, setTomogramConfig])
 
   const { logPlausibleCopyEvent } = useLogPlausibleCopyEvent()
 

diff --git a/frontend/packages/data-portal/app/components/Download/DownloadOptionsContent.tsx b/frontend/packages/data-portal/app/components/Download/DownloadOptionsContent.tsx
@@ -34,7 +34,7 @@ export function DownloadOptionsContent() {
     fileFormat,
     objectShapeType,
   } = useDownloadModalQueryParamState()
-  const { activeTomogramResolution } = useDownloadModalContext()
+  const { activeTomogram } = useDownloadModalContext()
 
   const downloadTabs = useMemo<TabData<DownloadTab>[]>(
     () => [
@@ -73,14 +73,12 @@ export function DownloadOptionsContent() {
       {objectShapeType && (
         <ModalSubtitle label={t('objectShapeType')} value={objectShapeType} />
       )}
-      {tomogramSampling && activeTomogramResolution && (
+      {tomogramSampling && activeTomogram && (
         <ModalSubtitle
           label={t('tomogramSampling')}
           value={`${t('unitAngstrom', { value: tomogramSampling })}, (${
-            activeTomogramResolution.size_x
-          }, ${activeTomogramResolution.size_y}, ${
-            activeTomogramResolution.size_z
-          })px`}
+            activeTomogram.size_x
+          }, ${activeTomogram.size_y}, ${activeTomogram.size_z})px`}
         />
       )}
       {tomogramProcessing && (

diff --git a/frontend/packages/data-portal/app/components/Run/RunHeader.tsx b/frontend/packages/data-portal/app/components/Run/RunHeader.tsx
@@ -44,8 +44,14 @@ function FileSummary({ data }: { data: FileSummaryData[] }) {
 
 export function RunHeader() {
   const multipleTomogramsEnabled = useFeatureFlag('multipleTomograms')
-  const { run, processingMethods, annotationFilesAggregates, tomogramsCount } =
-    useRunById()
+  const {
+    run,
+    processingMethods,
+    objectNames,
+    resolutions,
+    annotationFilesAggregates,
+    tomogramsCount,
+  } = useRunById()
   const { toggleDrawer } = useMetadataDrawer()
   const { t } = useI18n()
 
@@ -230,11 +236,9 @@ export function RunHeader() {
                   {
                     label: i18n.resolutionsAvailable,
                     inline: true,
-                    values: run.tomogram_stats
-                      .flatMap((stats) => stats.tomogram_resolutions)
-                      .map((resolutions) =>
-                        t('unitAngstrom', { value: resolutions.voxel_spacing }),
-                      ),
+                    values: resolutions.map((resolution) =>
+                      t('unitAngstrom', { value: resolution }),
+                    ),
                   },
                   {
                     label: i18n.tomogramProcessing,
@@ -244,13 +248,7 @@ export function RunHeader() {
                   {
                     label: i18n.annotatedObjects,
                     inline: true,
-                    values: Array.from(
-                      new Set(
-                        run.tomogram_stats
-                          .flatMap((stats) => stats.annotations)
-                          .map((annotation) => annotation.object_name),
-                      ),
-                    ),
+                    values: objectNames,
                     renderValues: (values: TableDataValue[]) => (
                       <CollapsibleList
                         entries={values.map((value) => ({

diff --git a/frontend/packages/data-portal/app/context/DownloadModal.context.ts b/frontend/packages/data-portal/app/context/DownloadModal.context.ts
@@ -5,14 +5,13 @@ import { BaseAnnotation } from 'app/state/annotation'
 
 export type DownloadModalType = 'dataset' | 'runs' | 'annotation'
 
-export type TomogramResolution =
-  GetRunByIdQuery['runs'][number]['tomogram_stats'][number]['tomogram_resolutions'][number]
+export type Tomogram = GetRunByIdQuery['tomograms_for_download'][number]
 
 export interface DownloadModalContextValue {
-  activeAnnotation?: BaseAnnotation | null
-  activeTomogramResolution?: TomogramResolution | null
+  activeAnnotation?: BaseAnnotation
+  activeTomogram?: Tomogram
   allTomogramProcessing?: string[]
-  allTomogramResolutions?: TomogramResolution[]
+  allTomograms?: Tomogram[]
   datasetId?: number
   datasetTitle?: string
   fileSize?: number

diff --git a/frontend/packages/data-portal/app/graphql/getRunById.server.ts b/frontend/packages/data-portal/app/graphql/getRunById.server.ts
@@ -138,29 +138,6 @@ const GET_RUN_BY_ID_QUERY = gql(`
         }
       }
 
-      tomogram_stats: tomogram_voxel_spacings {
-        annotations {
-          object_name
-          annotation_software
-
-          files(distinct_on: shape_type) {
-            shape_type
-          }
-        }
-
-        tomogram_resolutions: tomograms(distinct_on: voxel_spacing) {
-          https_mrc_scale0
-          id
-          processing
-          s3_mrc_scale0
-          s3_omezarr_dir
-          size_x
-          size_y
-          size_z
-          voxel_spacing
-        }
-      }
-
       tiltseries_aggregate {
         aggregate {
           avg {
@@ -292,7 +269,50 @@ const GET_RUN_BY_ID_QUERY = gql(`
       }
     }
 
-    # Distinct tomogram processing methods:
+    # Tomograms download:
+    tomograms_for_download: tomograms(
+      where: { tomogram_voxel_spacing: { run_id: { _eq: $id } } }
+    ) {
+      https_mrc_scale0
+      id
+      processing
+      reconstruction_method
+      s3_mrc_scale0
+      s3_omezarr_dir
+      size_x
+      size_y
+      size_z
+      voxel_spacing
+    }
+
+    # Annotation metadata:
+    annotations_for_softwares: annotations(
+      where: { tomogram_voxel_spacing: { run_id: { _eq: $id } } }
+      distinct_on: annotation_software
+    ) {
+      annotation_software
+    }
+    annotations_for_object_names: annotations(
+      where: { tomogram_voxel_spacing: { run_id: { _eq: $id } } }
+      distinct_on: object_name
+    ) {
+      object_name
+    }
+    annotation_files_for_shape_types: annotation_files(
+      where: { annotation: { tomogram_voxel_spacing: { run_id: { _eq: $id } } } }
+      distinct_on: shape_type
+    ) {
+      shape_type
+    }
+
+    # Tomogram metadata:
+    tomograms_for_resolutions: tomograms(
+      where: { tomogram_voxel_spacing: { run_id: { _eq: $id } } }
+      distinct_on: voxel_spacing
+      order_by: { voxel_spacing: asc  }
+    ) {
+      voxel_spacing
+    }
     tomograms_for_distinct_processing_methods: tomograms(
       where: { tomogram_voxel_spacing: { run_id: { _eq: $id } } }
       distinct_on: processing

diff --git a/frontend/packages/data-portal/app/hooks/useRunById.ts b/frontend/packages/data-portal/app/hooks/useRunById.ts
@@ -1,7 +1,7 @@
-import { useMemo } from 'react'
 import { useTypedLoaderData } from 'remix-typedjson'
 
 import { GetRunByIdQuery } from 'app/__generated__/graphql'
+import { isNotNullish } from 'app/utils/nullish'
 
 export function useRunById() {
   const data = useTypedLoaderData<GetRunByIdQuery>()
@@ -12,50 +12,26 @@ export function useRunById() {
 
   const { tomograms } = data
 
+  const tomogramsForDownload = data.tomograms_for_download
+
   const processingMethods = data.tomograms_for_distinct_processing_methods.map(
     (tomogram) => tomogram.processing,
   )
 
-  const objectNames = useMemo(
-    () =>
-      Array.from(
-        new Set(
-          run.tomogram_stats.flatMap((voxelSpacing) =>
-            voxelSpacing.annotations.map(
-              (annotation) => annotation.object_name,
-            ),
-          ),
-        ),
-      ),
-    [run],
+  const objectNames = data.annotations_for_object_names.map(
+    (annotation) => annotation.object_name,
   )
 
-  const objectShapeTypes = useMemo(
-    () =>
-      Array.from(
-        new Set(
-          run.tomogram_stats.flatMap((voxelSpacing) =>
-            voxelSpacing.annotations.flatMap((annotation) =>
-              annotation.files.map((file) => file.shape_type),
-            ),
-          ),
-        ),
-      ),
-    [run],
+  const objectShapeTypes = data.annotation_files_for_shape_types.map(
+    (file) => file.shape_type,
   )
 
-  const annotationSoftwares = useMemo(
-    () =>
-      Array.from(
-        new Set(
-          run.tomogram_stats.flatMap((voxelSpacing) =>
-            voxelSpacing.annotations
-              .filter((annotation) => annotation.annotation_software)
-              .map((annotation) => annotation.annotation_software as string),
-          ),
-        ),
-      ),
-    [run],
+  const annotationSoftwares = data.annotations_for_softwares
+    .map((annotation) => annotation.annotation_software)
+    .filter(isNotNullish)
+
+  const resolutions = data.tomograms_for_resolutions.map(
+    (tomogram) => tomogram.voxel_spacing,
   )
 
   const annotationFilesAggregates = {
@@ -73,10 +49,12 @@ export function useRunById() {
     run,
     annotationFiles,
     tomograms,
+    tomogramsForDownload,
     processingMethods,
     objectNames,
     objectShapeTypes,
     annotationSoftwares,
+    resolutions,
     annotationFilesAggregates,
     tomogramsCount,
   }

diff --git a/frontend/packages/data-portal/app/routes/runs.$id.tsx b/frontend/packages/data-portal/app/routes/runs.$id.tsx
@@ -2,7 +2,7 @@
 
 import { ShouldRevalidateFunctionArgs } from '@remix-run/react'
 import { json, LoaderFunctionArgs } from '@remix-run/server-runtime'
-import { isNumber, toNumber } from 'lodash-es'
+import { toNumber } from 'lodash-es'
 import { useMemo } from 'react'
 import { match } from 'ts-pattern'
 
@@ -81,14 +81,11 @@ export default function RunByIdPage() {
     run,
     processingMethods,
     annotationFiles,
+    tomogramsForDownload,
     annotationFilesAggregates,
     tomogramsCount,
   } = useRunById()
 
-  const allTomogramResolutions = run.tomogram_stats.flatMap((stats) =>
-    stats.tomogram_resolutions.map((tomogram) => tomogram),
-  )
-
   const {
     downloadConfig,
     tomogramProcessing,
@@ -99,13 +96,13 @@ export default function RunByIdPage() {
   } = useDownloadModalQueryParamState()
 
   const activeTomogram =
-    (downloadConfig === DownloadConfig.Tomogram &&
-      allTomogramResolutions.find(
-        (tomogram) =>
-          `${tomogram.voxel_spacing}` === tomogramSampling &&
-          tomogram.processing === tomogramProcessing,
-      )) ||
-    null
+    downloadConfig === DownloadConfig.Tomogram
+      ? tomogramsForDownload.find(
+          (tomogram) =>
+            `${tomogram.voxel_spacing}` === tomogramSampling &&
+            tomogram.processing === tomogramProcessing,
+        )
+      : undefined
 
   const tomogram = run.tomogram_voxel_spacings.at(0)
   const { t } = useI18n()
@@ -138,16 +135,6 @@ export default function RunByIdPage() {
     enabled: fileFormat !== 'zarr',
   })
 
-  const activeTomogramResolution = useMemo(
-    () =>
-      allTomogramResolutions.find((resolution) =>
-        tomogramSampling !== null && isNumber(+tomogramSampling)
-          ? resolution.voxel_spacing === +tomogramSampling
-          : false,
-      ),
-    [allTomogramResolutions, tomogramSampling],
-  )
-
   return (
     <TablePageLayout
       header={<RunHeader />}
@@ -178,9 +165,9 @@ export default function RunByIdPage() {
       downloadModal={
         <DownloadModal
           activeAnnotation={activeAnnotation}
-          activeTomogramResolution={activeTomogramResolution}
+          activeTomogram={activeTomogram}
           allTomogramProcessing={processingMethods}
-          allTomogramResolutions={allTomogramResolutions}
+          allTomograms={tomogramsForDownload}
           datasetId={run.dataset.id}
           datasetTitle={run.dataset.title}
           fileSize={fileSize}

diff --git a/frontend/packages/data-portal/app/utils/nullish.ts b/frontend/packages/data-portal/app/utils/nullish.ts
@@ -0,0 +1,4 @@
+/** Checks nullish and updates the type. */
+export function isNotNullish<T>(x: T): x is NonNullable<T> {
+  return x != null
+}