Merge pull request #2497 from GMOD/607_stats_estimation_v2

Add stats estimation to JB2

packages/core/data_adapters/BaseAdapter.ts

@@ -1,5 +1,5 @@
 import { Observable, merge } from 'rxjs'
-import { takeUntil } from 'rxjs/operators'
+import { filter, takeUntil, toArray } from 'rxjs/operators'
 import { isStateTreeNode, getSnapshot } from 'mobx-state-tree'
 import { ObservableCreate } from '../util/rxjs'
 import { checkAbortSignal, observeAbortSignal } from '../util'
@@ -10,7 +10,12 @@ import {
 } from '../configuration/configurationSchema'
 import { getSubAdapterType } from './dataAdapterCache'
 import { Region, NoAssemblyRegion } from '../util/types'
-import { blankStats, rectifyStats, scoresToStats } from '../util/stats'
+import {
+  blankStats,
+  rectifyStats,
+  scoresToStats,
+  BaseFeatureStats,
+} from '../util/stats'
 import BaseResult from '../TextSearch/BaseResults'
 import idMaker from '../util/idMaker'
 import PluginManager from '../PluginManager'
@@ -71,6 +76,7 @@ export abstract class BaseAdapter {
     this.config = config
     this.getSubAdapter = getSubAdapter
     this.pluginManager = pluginManager
+
     // note: we use switch on jest here for more simple feature IDs
     // in test environment
     if (typeof jest === 'undefined') {
@@ -94,6 +100,7 @@ export abstract class BaseAdapter {
  * subclasses must implement.
  */
 export abstract class BaseFeatureDataAdapter extends BaseAdapter {
+  private estimateStatsCache: BaseFeatureStats | undefined
   /**
    * Get all reference sequence names used in the data source
    *
@@ -258,6 +265,88 @@ export abstract class BaseFeatureDataAdapter extends BaseAdapter {
       scoreSum,
     })
   }
+
+  public async getGlobalStats(
+    regionToStart: Region,
+    opts?: BaseOptions,
+  ): Promise<BaseFeatureStats> {
+    // Estimates once, then cache stats for future calls
+    return this.estimateStatsCache
+      ? this.estimateStatsCache
+      : this.estimateGlobalStats(regionToStart, opts)
+  }
+
+  public async estimateGlobalStats(
+    region: Region,
+    opts?: BaseOptions,
+  ): Promise<BaseFeatureStats> {
+    const { statusCallback = () => {} } = opts || {}
+    const statsFromInterval = async (length: number, expansionTime: number) => {
+      const sampleCenter = region.start * 0.75 + region.end * 0.25
+      const start = Math.max(0, Math.round(sampleCenter - length / 2))
+      const end = Math.min(Math.round(sampleCenter + length / 2), region.end)
+
+      const feats = this.getFeatures(region, opts)
+      let features
+      try {
+        features = await feats
+          .pipe(
+            filter(
+              (f: Feature) =>
+                typeof f.get === 'function' &&
+                f.get('start') >= start &&
+                f.get('end') <= end,
+            ),
+            toArray(),
+          )
+          .toPromise()
+      } catch (e) {
+        if (`${e}`.match(/HTTP 404/)) {
+          throw new Error(`${e}`)
+        }
+        console.warn('Skipping feature density calculation: ', e)
+        return { featureDensity: Infinity }
+      }
+
+      // if no features in range or adapter has no f.get, cancel feature density calculation
+      if (features.length === 0) {
+        return { featureDensity: 0 }
+      }
+      const featureDensity = features.length / length
+      return maybeRecordStats(
+        length,
+        {
+          featureDensity: featureDensity,
+        },
+        features.length,
+        expansionTime,
+      )
+    }
+
+    const maybeRecordStats = async (
+      interval: number,
+      stats: BaseFeatureStats,
+      statsSampleFeatures: number,
+      expansionTime: number,
+    ): Promise<BaseFeatureStats> => {
+      const refLen = region.end - region.start
+      if (statsSampleFeatures >= 300 || interval * 2 > refLen) {
+        statusCallback('')
+        this.estimateStatsCache = stats
+      } else if (expansionTime <= 4) {
+        expansionTime++
+        return statsFromInterval(interval * 2, expansionTime)
+      } else {
+        statusCallback('')
+        console.error('Stats estimation reached timeout')
+        this.estimateStatsCache = { featureDensity: Infinity }
+      }
+      return this.estimateStatsCache as BaseFeatureStats
+    }
+
+    statusCallback('Calculating stats')
+    return statsFromInterval(100, 0)
+  }
 }
 
 export interface RegionsAdapter extends BaseAdapter {

packages/core/rpc/coreRpcMethods.ts

@@ -18,6 +18,7 @@ import {
 import { Region } from '../util/types'
 import { checkAbortSignal, renameRegionsIfNeeded } from '../util'
 import SimpleFeature, { SimpleFeatureSerialized } from '../util/simpleFeature'
+import { BaseFeatureStats } from '../util/stats'
 
 export class CoreGetRefNames extends RpcMethodType {
   name = 'CoreGetRefNames'
@@ -40,6 +41,7 @@ export class CoreGetRefNames extends RpcMethodType {
       sessionId,
       adapterConfig,
     )
+
     if (dataAdapter instanceof BaseFeatureDataAdapter) {
       return dataAdapter.getRefNames(deserializedArgs)
     }
@@ -188,6 +190,55 @@ export interface RenderArgsSerialized extends ServerSideRenderArgsSerialized {
   rendererType: string
 }
 
+export class CoreGetGlobalStats extends RpcMethodType {
+  name = 'CoreGetGlobalStats'
+
+  async serializeArguments(
+    args: RenderArgs & { signal?: AbortSignal; statusCallback?: Function },
+    rpcDriverClassName: string,
+  ) {
+    const assemblyManager =
+      this.pluginManager.rootModel?.session?.assemblyManager
+    if (!assemblyManager) {
+      return args
+    }
+    const renamedArgs = await renameRegionsIfNeeded(assemblyManager, {
+      ...args,
+      filters: args.filters && args.filters.toJSON().filters,
+    })
+
+    return super.serializeArguments(renamedArgs, rpcDriverClassName)
+  }
+
+  async execute(
+    args: {
+      adapterConfig: {}
+      regions: Region[]
+      signal?: RemoteAbortSignal
+      headers?: Record<string, string>
+      sessionId: string
+    },
+    rpcDriverClassName: string,
+  ): Promise<BaseFeatureStats> {
+    const deserializedArgs = await this.deserializeArguments(
+      args,
+      rpcDriverClassName,
+    )
+
+    const { adapterConfig, sessionId, regions } = deserializedArgs
+    const { dataAdapter } = await getAdapter(
+      this.pluginManager,
+      sessionId,
+      adapterConfig,
+    )
+
+    if (dataAdapter instanceof BaseFeatureDataAdapter) {
+      return dataAdapter.getGlobalStats(regions[0], deserializedArgs)
+    }
+    throw new Error('Data adapter not found')
+  }
+}
+
 /**
  * fetches features from an adapter and call a renderer with them
  */

packages/core/util/stats.ts

@@ -11,10 +11,14 @@ export interface UnrectifiedFeatureStats {
   featureCount: number
   basesCovered: number
 }
-export interface FeatureStats extends UnrectifiedFeatureStats {
+export interface BaseFeatureStats {
+  featureDensity: number
+}
+export interface FeatureStats
+  extends UnrectifiedFeatureStats,
+    BaseFeatureStats {
   scoreMean: number
   scoreStdDev: number
-  featureDensity: number
 }
 
 /*

plugins/alignments/src/BamAdapter/BamAdapter.ts

@@ -8,6 +8,7 @@ import { checkAbortSignal } from '@jbrowse/core/util'
 import { openLocation } from '@jbrowse/core/util/io'
 import { ObservableCreate } from '@jbrowse/core/util/rxjs'
 import { Feature } from '@jbrowse/core/util/simpleFeature'
+import { BaseFeatureStats } from '@jbrowse/core/util/stats'
 import { toArray } from 'rxjs/operators'
 import { readConfObject } from '@jbrowse/core/configuration'
 import BamSlightlyLazyFeature from './BamSlightlyLazyFeature'
@@ -188,6 +189,20 @@ export default class BamAdapter extends BaseFeatureDataAdapter {
     }, signal)
   }
 
+  async estimateGlobalStats(
+    region: Region,
+    opts?: BaseOptions,
+  ): Promise<BaseFeatureStats> {
+    const { bam } = await this.configure()
+    const featCount = await bam.lineCount(region.refName)
+    if (featCount < 0) {
+      return super.estimateGlobalStats(region, opts)
+    }
+
+    const featureDensity = featCount / (region.end - region.start)
+    return { featureDensity }
+  }
+
   freeResources(/* { region } */): void {}
 
   // depends on setup being called before the BAM constructor

plugins/alignments/src/LinearAlignmentsDisplay/models/model.tsx

@@ -32,6 +32,7 @@ const stateModelFactory = (
         height: 250,
         showCoverage: true,
         showPileup: true,
+        userFeatureScreenDensity: types.maybe(types.number),
       }),
     )
     .volatile(() => ({
@@ -126,9 +127,9 @@ const stateModelFactory = (
           height: self.snpCovHeight,
         }
       },
-      setUserBpPerPxLimit(limit: number) {
-        self.PileupDisplay.setUserBpPerPxLimit(limit)
-        self.SNPCoverageDisplay.setUserBpPerPxLimit(limit)
+      setUserFeatureScreenDensity(limit: number) {
+        self.PileupDisplay.setUserFeatureScreenDensity(limit)
+        self.SNPCoverageDisplay.setUserFeatureScreenDensity(limit)
       },
       setPileupDisplay(displayConfig: AnyConfigurationModel) {
         self.PileupDisplay = {

plugins/alignments/src/LinearPileupDisplay/configSchema.ts

@@ -23,10 +23,10 @@ function PileupConfigFactory(pluginManager: PluginManager) {
         SvgFeatureRenderer: SvgFeatureRendererConfigSchema,
       }),
       renderer: '',
-      maxDisplayedBpPerPx: {
+      maxFeatureScreenDensity: {
         type: 'number',
-        description: 'maximum bpPerPx that is displayed in the view',
-        defaultValue: 100,
+        description: 'maximum features per pixel that is displayed in the view',
+        defaultValue: 15,
       },
       colorScheme: {
         type: 'stringEnum',

plugins/alignments/src/LinearPileupDisplay/model.ts

@@ -416,6 +416,7 @@ const stateModelFactory = (
             colorBy,
             rpcDriverName,
           } = self
+
           return {
             ...superRenderProps(),
             notReady: !ready || (sortedBy && self.currBpPerPx !== view.bpPerPx),

plugins/alignments/src/LinearSNPCoverageDisplay/models/configSchema.ts

@@ -39,10 +39,10 @@ export default function SNPCoverageConfigFactory(pluginManager: PluginManager) {
         description: 'draw upside down',
         defaultValue: false,
       },
-      maxDisplayedBpPerPx: {
+      maxFeatureScreenDensity: {
         type: 'number',
-        description: 'maximum bpPerPx that is displayed in the view',
-        defaultValue: 100,
+        description: 'maximum features per pixel that is displayed in the view',
+        defaultValue: 15,
       },
       headroom: {
         type: 'number',

plugins/alignments/src/SNPCoverageAdapter/SNPCoverageAdapter.ts

@@ -16,6 +16,7 @@ import {
   getModificationPositions,
   Mismatch,
 } from '../BamAdapter/MismatchParser'
+import { BaseFeatureStats } from '@jbrowse/core/util/stats'
 
 interface SNPCoverageOptions extends BaseOptions {
   filters?: SerializableFilterChain
@@ -104,6 +105,14 @@ export default class SNPCoverageAdapter extends BaseFeatureDataAdapter {
     }, opts.signal)
   }
 
+  async estimateGlobalStats(
+    region: Region,
+    opts?: BaseOptions,
+  ): Promise<BaseFeatureStats> {
+    const { subadapter } = await this.configure()
+    return subadapter.estimateGlobalStats(region, opts)
+  }
+
   async getRefNames(opts: BaseOptions = {}) {
     const { subadapter } = await this.configure()
     return subadapter.getRefNames(opts)

plugins/bed/src/BedTabixAdapter/BedTabixAdapter.ts

@@ -15,6 +15,7 @@ import { ucscProcessedTranscript } from '../util'
 import MyConfigSchema from './configSchema'
 import PluginManager from '@jbrowse/core/PluginManager'
 import { getSubAdapterType } from '@jbrowse/core/data_adapters/dataAdapterCache'
+import { BaseFeatureStats } from '@jbrowse/core/util/stats'
 
 export default class BedTabixAdapter extends BaseFeatureDataAdapter {
   private parser: any
@@ -154,5 +155,17 @@ export default class BedTabixAdapter extends BaseFeatureDataAdapter {
     }, opts.signal)
   }
 
+  async estimateGlobalStats(
+    region: Region,
+    opts?: BaseOptions,
+  ): Promise<BaseFeatureStats> {
+    const featCount = await this.bed.lineCount(region.refName)
+    if (featCount === -1) {
+      return super.estimateGlobalStats(region, opts)
+    }
+    const featureDensity = featCount / (region.end - region.start)
+    return { featureDensity }
+  }
+
   public freeResources(): void {}
 }

plugins/config/src/ConfigurationEditorWidget/components/__snapshots__/ConfigurationEditor.test.js.snap

@@ -761,7 +761,7 @@ exports[`ConfigurationEditor widget renders with defaults of the PileupTrack sch
           class="MuiFormLabel-root MuiInputLabel-root MuiInputLabel-formControl MuiInputLabel-animated MuiInputLabel-shrink MuiFormLabel-filled"
           data-shrink="true"
         >
-          maxDisplayedBpPerPx
+          maxFeatureScreenDensity
         </label>
         <div
           class="MuiInputBase-root MuiInput-root MuiInput-underline MuiInputBase-formControl MuiInput-formControl"
@@ -770,13 +770,13 @@ exports[`ConfigurationEditor widget renders with defaults of the PileupTrack sch
             aria-invalid="false"
             class="MuiInputBase-input MuiInput-input"
             type="number"
-            value="1000"
+            value="0.5"
           />
         </div>
         <p
           class="MuiFormHelperText-root MuiFormHelperText-filled"
         >
-          maximum bpPerPx that is displayed in the view
+          maximum features per pixel that is displayed in the view
         </p>
       </div>
     </div>