Merge pull request #2324 from broadinstitute/ew-fast-dot-plots

Speed up dot plots: integrate preprocessed data on frontend (SCP-5992)

Author: eweitz

app/javascript/components/visualization/DotPlot.jsx

@@ -11,13 +11,15 @@ import { withErrorBoundary } from '~/lib/ErrorBoundary'
 import LoadingSpinner, { morpheusLoadingSpinner } from '~/lib/LoadingSpinner'
 import { fetchServiceWorkerCache } from '~/lib/service-worker-cache'
 import { getSCPContext } from '~/providers/SCPContextProvider'
+import { getFeatureFlagsWithDefaults } from '~/providers/UserProvider'
+import '~/lib/dot-plot-precompute-patch'
 
 export const dotPlotColorScheme = {
   // Blue, purple, red.  These red and blue hues are accessible, per WCAG.
   colors: ['#0000BB', '#CC0088', '#FF0000'],
-
   // TODO: Incorporate expression units, once such metadata is available.
-  values: [0, 0.5, 1]
+  values: [0, 0.5, 1],
+  scalingMode: 'relative'
 }
 
 /**
@@ -199,14 +201,19 @@ function RawDotPlot({
   useEffect(() => {
     /** Fetch Morpheus data for dot plot */
     async function getDataset() {
+      const flags = getFeatureFlagsWithDefaults()
+      const usePrecomputed = flags?.dot_plot_preprocessing_frontend || false
+
       const [dataset, perfTimes] = await fetchMorpheusJson(
         studyAccession,
         genes,
         cluster,
         annotation.name,
         annotation.type,
         annotation.scope,
-        subsample
+        subsample,
+        false, // mock
+        usePrecomputed // isPrecomputed
       )
       logFetchMorpheusDataset(perfTimes, cluster, annotation, genes)
 
@@ -265,8 +272,19 @@ export function renderDotPlot({
   const $target = $(target)
   $target.empty()
 
-  // Collapse by mean
-  const tools = [{
+  // Check if dataset is pre-computed dot plot data
+  // Pre-computed data has structure: { annotation_name, values, genes }
+  let processedDataset = dataset
+  let isPrecomputed = false
+
+  if (dataset && dataset.annotation_name && dataset.values && dataset.genes) {
+    // This is pre-computed dot plot data - convert it using the patch
+    processedDataset = window.createMorpheusDotPlot(dataset)
+    isPrecomputed = true
+  }
+
+  // Collapse by mean (only for non-precomputed data)
+  const tools = isPrecomputed ? [] : [{
     name: 'Collapse',
     params: {
       collapse_method: 'Mean',
@@ -275,30 +293,37 @@ export function renderDotPlot({
       collapse_to_fields: [annotationName],
       pass_expression: '>',
       pass_value: '0',
-      percentile: '100',
+      percentile: '75',
       compute_percent: true
     }
   }]
 
   const config = {
     shape: 'circle',
-    dataset,
+    dataset: processedDataset,
     el: $target,
     menu: null,
     error: morpheusErrorHandler($target, setShowError, setErrorContent),
-    colorScheme: {
-      scalingMode: 'relative'
-    },
     focus: null,
     tabManager: morpheusTabManager($target),
     tools,
     loadedCallback: () => logMorpheusPerfTime(target, 'dotplot', genes)
   }
 
+  // For pre-computed data, tell Morpheus to display series 0 for color
+  // and use series 1 for sizing (which happens automatically with shape: 'circle')
+  if (isPrecomputed) {
+    config.symmetricColorScheme = false
+    // Tell Morpheus which series to use for coloring the heatmap
+    config.seriesIndex = 0 // Display series 0 (Mean Expression) for colors
+    // Explicitly set the size series
+    config.sizeBySeriesIndex = 1 // Use series 1 (__count) for sizing
+  }
+
   // Load annotations if specified
-  config.columnSortBy = [
-    { field: annotationName, order: 0 }
-  ]
+  // config.columnSortBy = [
+  //   { field: annotationName, order: 0 }
+  // ]
   config.columns = [
     { field: annotationName, display: 'text' }
   ]
@@ -319,7 +344,28 @@ export function renderDotPlot({
   config.columnColorModel = annotColorModel
 
 
-  config.colorScheme = dotPlotColorScheme
+  // Set color scheme (will be overridden for precomputed data below)
+  if (!isPrecomputed) {
+    config.colorScheme = dotPlotColorScheme
+  }
+
+  // For precomputed data, configure the sizer to use the __count series
+  if (isPrecomputed && processedDataset) {
+    // The color scheme should already have a sizer - we just need to configure it
+    config.sizeBy = {
+      seriesName: 'percent',
+      min: 0,
+      max: 75
+    }
+
+    // Use relative color scheme for raw expression values
+    // This will scale colors based on the actual data range across all genes and cell types
+    config.colorScheme = {
+      colors: ['#0000BB', '#CC0088', '#FF0000'],
+      values: [0, 0.5, 1],
+      scalingMode: 'relative'
+    }
+  }
 
   patchServiceWorkerCache()
 

app/javascript/lib/dot-plot-precompute-patch.js

@@ -0,0 +1,168 @@
+/**
+ * Monkeypatch for Morpheus to accept pre-computed dot plot data
+ * Data format: [mean_expression, percent_expressing]
+ * Mean expression values are normalized per-gene (per-row) to 0-1 range for proper color scaling
+ */
+
+(function() {
+  'use strict'
+
+  /**
+   * Apply the dot plot patch to Morpheus once it's loaded
+   * Waits for window.morpheus to be available before patching
+   */
+  function applyDotPlotPatch() {
+    if (typeof window.morpheus === 'undefined') {
+      // Morpheus not loaded yet, wait a bit and try again
+      setTimeout(applyDotPlotPatch, 100)
+      return
+    }
+
+    /**
+     * Convert your dot plot JSON format to a Morpheus dataset
+     */
+    window.morpheus.DotPlotConverter = {
+
+      createDataset(data) {
+        const cellTypes = data.values
+        const geneNames = Object.keys(data.genes)
+        const nRows = geneNames.length
+        const nCols = cellTypes.length
+
+        // Create dataset with Float32 data type
+        // The dataset name becomes the first series name by default
+        const dataset = new window.morpheus.Dataset({
+          name: 'Mean Expression',
+          rows: nRows,
+          columns: nCols,
+          dataType: 'Float32'
+        })
+
+        // Add second series for the size metric (percent expressing)
+        // Morpheus uses 'percent' for sizing in dot plots
+        dataset.addSeries({
+          name: 'percent',
+          dataType: 'Float32'
+        })
+
+        // Set up row metadata (genes)
+        const rowIds = dataset.getRowMetadata().add('id')
+        geneNames.forEach((gene, i) => {
+          rowIds.setValue(i, gene)
+        })
+
+        // Set up column metadata (cell types)
+        const colIds = dataset.getColumnMetadata().add('id')
+        const cellTypeMetadata = dataset.getColumnMetadata().add(data.annotation_name || 'Cell Type')
+        cellTypes.forEach((cellType, j) => {
+          colIds.setValue(j, cellType)
+          cellTypeMetadata.setValue(j, cellType)
+        })
+
+        // Fill in the data
+        // Series 0: mean expression (for color) - will be normalized per-gene (row)
+        // Series 1: percent expressing (for size) - will be scaled to 0-100
+        // Data format: values[0] = mean_expression, values[1] = percent_expressing
+        geneNames.forEach((gene, i) => {
+          const geneData = data.genes[gene]
+
+          geneData.forEach((values, j) => {
+            const meanExpression = values[0]
+            const percentExpressing = values[1]
+
+            // Use raw mean expression values, but convert zeros to NaN
+            // This excludes them from Morpheus color scaling while preserving actual values
+            const expressionValue = meanExpression === 0 ? NaN : meanExpression
+            dataset.setValue(i, j, expressionValue, 0) // Raw mean expression for color (zeros as NaN)
+            // Scale percent expressing to 0-100 range for better sizing
+            dataset.setValue(i, j, percentExpressing * 100, 1) // Percent expressing (0-100) for size
+          })
+        })
+
+        return dataset
+      },
+
+      /**
+       * Add custom properties to enable dot plot mode
+       */
+      configureDotPlot(dataset) {
+        // Add a property to indicate this is dot plot data
+        dataset._isDotPlot = true
+        dataset._dotPlotSizeSeries = 1 // Percent expressing
+        dataset._dotPlotColorSeries = 0 // Mean expression
+
+        return dataset
+      }
+    }
+
+    /**
+     * Register a custom JSON reader for dot plot format
+     */
+    const OriginalJsonReader = window.morpheus.JsonDatasetReader
+
+    window.morpheus.JsonDatasetReader = function() {
+      OriginalJsonReader.call(this)
+    }
+
+    window.morpheus.JsonDatasetReader.prototype = Object.create(OriginalJsonReader.prototype)
+
+    const originalRead = OriginalJsonReader.prototype.read
+    window.morpheus.JsonDatasetReader.prototype.read = function(fileOrUrl, callback) {
+      const self = this
+
+      // Check if it's our dot plot format
+      window.morpheus.Util.getText(fileOrUrl).then(text => {
+        try {
+          const data = JSON.parse(text)
+
+          // Check if it matches our dot plot format
+          if (data.annotation_name && data.values && data.genes) {
+            let dataset = window.morpheus.DotPlotConverter.createDataset(data)
+            dataset = window.morpheus.DotPlotConverter.configureDotPlot(dataset)
+            callback(null, dataset)
+          } else {
+            // Fall back to original reader
+            originalRead.call(self, fileOrUrl, callback)
+          }
+        } catch (err) {
+          callback(err)
+        }
+      }).catch(err => {
+        callback(err)
+      })
+    }
+
+    /**
+     * Helper to create dot plot directly from your data object
+     */
+    window.createMorpheusDotPlot = function(data) {
+      const dataset = window.morpheus.DotPlotConverter.createDataset(data)
+      return window.morpheus.DotPlotConverter.configureDotPlot(dataset)
+    }
+
+    /**
+     * Patch the HeatMap to properly handle dot plot sizing with __count series
+     */
+    const OriginalHeatMap = window.morpheus.HeatMap
+    window.morpheus.HeatMap = function(options) {
+      const heatmap = new OriginalHeatMap(options)
+
+      // Check if this is a precomputed dot plot dataset
+      if (options.dataset && options.dataset._isDotPlot) {
+        // Force the heatmap to use series 1 for sizing
+        if (heatmap.heatMapElementCanvas) {
+          heatmap.heatMapElementCanvas.sizeBySeriesIndex = 1
+        }
+      }
+
+      return heatmap
+    }
+
+    // Copy static properties
+    Object.setPrototypeOf(window.morpheus.HeatMap, OriginalHeatMap)
+    window.morpheus.HeatMap.prototype = OriginalHeatMap.prototype
+  }
+
+  // Start trying to apply the patch
+  applyDotPlotPatch()
+})()

app/javascript/lib/scp-api.jsx

@@ -750,7 +750,8 @@ export async function fetchMorpheusJson(
   annotationType,
   annotationScope,
   subsample,
-  mock=false
+  mock=false,
+  usePrecomputed=true
 ) {
   let geneString = genes
   if (Array.isArray(genes)) {
@@ -764,7 +765,8 @@ export async function fetchMorpheusJson(
     subsample,
     genes: geneString
   }
-  const apiUrl = `/studies/${studyAccession}/expression/morpheus${stringifyQuery(paramObj)}`
+  const endpoint = usePrecomputed ? 'dotplot' : 'morpheus'
+  const apiUrl = `/studies/${studyAccession}/expression/${endpoint}${stringifyQuery(paramObj)}`
   // don't camelcase the keys since those can be cluster names,
   // so send false for the 4th argument
   const [violin, perfTimes] = await scpApi(apiUrl, defaultInit(), mock, false)

db/migrate/20250609182851_make_all_facets_mongo_based.rb

@@ -1,7 +1,11 @@
 class MakeAllFacetsMongoBased < Mongoid::Migration
   def self.up
     SearchFacet.update_all(is_mongo_based: true)
-    CellMetadatum.where(name: 'organism_age').map(&:set_minmax_by_units!)
+    # Only process CellMetadatum records that have a valid study association
+    CellMetadatum.where(name: 'organism_age').each do |cell_metadatum|
+      next if cell_metadatum.study.nil?
+      cell_metadatum.set_minmax_by_units!
+    end
   end
 
   def self.down

db/migrate/20250616192718_create_dot_plot_gene_collection.rb

@@ -2,7 +2,8 @@ class CreateDotPlotGeneCollection < Mongoid::Migration
   def self.up
     # since these documents will be created by scp-ingest-pipeline, the collection needs to exist first to
     # prevent errors when the job tries to create them
-    DotPlotGene.collection.create
+    # Only create if it doesn't already exist
+    DotPlotGene.collection.create unless DotPlotGene.collection.database.collection_names.include?('dot_plot_genes')
   end
 
   def self.down

db/migrate/20251106100123_add_dot_plot_preprocessing_frontend_feature_flag.rb

@@ -0,0 +1,13 @@
+class AddDotPlotPreprocessingFrontendFeatureFlag < Mongoid::Migration
+  def self.up
+    FeatureFlag.find_or_create_by(name: 'dot_plot_preprocessing_frontend') do |flag|
+      flag.default_value = false
+      flag.description = 'Enable pre-computed dot plot data from backend preprocessing'
+    end
+  end
+
+  def self.down
+    flag = FeatureFlag.find_by(name: 'dot_plot_preprocessing_frontend')
+    flag.destroy if flag.present?
+  end
+end