Render start/stop codons on LASFD (#731)

* toggles

* rendering of high-contrast, feature-label-aware, start/stop codons on LASFD

* tie loading of refSeq to zoomThreshold level rather than hard-coded "3"

* make use of util functions

* Ensure codon rendering occurs after feature rendering, to avoid codons getting covered by features

Author: haessar

packages/jbrowse-plugin-apollo/src/LinearApolloReferenceSequenceDisplay/drawSequenceTrack.ts

@@ -3,29 +3,7 @@ import { type BlockSet } from '@jbrowse/core/util/blockTypes'
 import { type Theme } from '@mui/material'
 
 import { type ApolloSessionModel } from '../session'
-
-function colorCode(letter: string, theme: Theme) {
-  const letterUpper = letter.toUpperCase()
-  if (
-    letterUpper === 'A' ||
-    letterUpper === 'C' ||
-    letterUpper === 'G' ||
-    letterUpper === 'T'
-  ) {
-    return theme.palette.bases[letterUpper].main.toString()
-  }
-  return 'lightgray'
-}
-
-function codonColorCode(letter: string, theme: Theme, highContrast?: boolean) {
-  if (letter === 'M') {
-    return theme.palette.startCodon
-  }
-  if (letter === '*') {
-    return highContrast ? theme.palette.text.primary : theme.palette.stopCodon
-  }
-  return
-}
+import { codonColorCode, colorCode } from '../util/displayUtils'
 
 function drawLetter(
   seqTrackctx: CanvasRenderingContext2D,

packages/jbrowse-plugin-apollo/src/LinearApolloSixFrameDisplay/stateModel/base.ts

@@ -39,6 +39,8 @@ export function baseModelFactory(
       graphical: true,
       table: false,
       showFeatureLabels: true,
+      showStartCodons: false,
+      showStopCodons: true,
       showCheckResults: true,
       zoomThreshold: 200,
       heightPreConfig: types.maybe(
@@ -179,6 +181,12 @@ export function baseModelFactory(
       toggleShowFeatureLabels() {
         self.showFeatureLabels = !self.showFeatureLabels
       },
+      toggleShowStartCodons() {
+        self.showStartCodons = !self.showStartCodons
+      },
+      toggleShowStopCodons() {
+        self.showStopCodons = !self.showStopCodons
+      },
       toggleShowCheckResults() {
         self.showCheckResults = !self.showCheckResults
       },
@@ -193,7 +201,14 @@ export function baseModelFactory(
       const { filteredFeatureTypes, trackMenuItems: superTrackMenuItems } = self
       return {
         trackMenuItems() {
-          const { graphical, table, showFeatureLabels, showCheckResults } = self
+          const {
+            graphical,
+            table,
+            showFeatureLabels,
+            showStartCodons,
+            showStopCodons,
+            showCheckResults,
+          } = self
           return [
             ...superTrackMenuItems(),
             {
@@ -232,6 +247,22 @@ export function baseModelFactory(
                     self.toggleShowFeatureLabels()
                   },
                 },
+                {
+                  label: 'Show start codons',
+                  type: 'checkbox',
+                  checked: showStartCodons,
+                  onClick: () => {
+                    self.toggleShowStartCodons()
+                  },
+                },
+                {
+                  label: 'Show stop codons',
+                  type: 'checkbox',
+                  checked: showStopCodons,
+                  onClick: () => {
+                    self.toggleShowStopCodons()
+                  },
+                },
                 {
                   label: 'Check Results',
                   type: 'checkbox',
@@ -326,7 +357,7 @@ export function baseModelFactory(
               void (
                 self.session as unknown as ApolloSessionModel
               ).apolloDataStore.loadFeatures(self.regions)
-              if (self.lgv.bpPerPx <= 3) {
+              if (self.lgv.bpPerPx <= self.zoomThreshold) {
                 void (
                   self.session as unknown as ApolloSessionModel
                 ).apolloDataStore.loadRefSeq(self.regions)

packages/jbrowse-plugin-apollo/src/LinearApolloSixFrameDisplay/stateModel/rendering.ts

@@ -1,15 +1,62 @@
 /* eslint-disable @typescript-eslint/no-unnecessary-condition */
 import type PluginManager from '@jbrowse/core/PluginManager'
 import { type AnyConfigurationSchemaType } from '@jbrowse/core/configuration/configurationSchema'
-import { doesIntersect2 } from '@jbrowse/core/util'
+import {
+  defaultCodonTable,
+  doesIntersect2,
+  getFrame,
+  revcom,
+} from '@jbrowse/core/util'
 import { type Theme, createTheme } from '@mui/material'
 import { autorun } from 'mobx'
 import { type Instance, addDisposer, types } from 'mobx-state-tree'
 
 import { type ApolloSessionModel } from '../../session'
+import { codonColorCode } from '../../util/displayUtils'
 
 import { layoutsModelFactory } from './layouts'
 
+function drawCodon(
+  ctx: CanvasRenderingContext2D,
+  codon: string,
+  leftPx: number,
+  index: number,
+  theme: Theme,
+  highContrast: boolean,
+  bpPerPx: number,
+  bp: number,
+  rowHeight: number,
+  showFeatureLabels: boolean,
+  showStartCodons: boolean,
+  showStopCodons: boolean,
+) {
+  const frameOffsets = (
+    showFeatureLabels ? [0, 4, 2, 0, 14, 12, 10] : [0, 2, 1, 0, 7, 6, 5]
+  ).map((b) => b * rowHeight)
+  const strands = [-1, 1] as const
+  for (const strand of strands) {
+    const frame = getFrame(bp, bp + 3, strand, 0)
+    const top = frameOffsets.at(frame)
+    if (top === undefined) {
+      continue
+    }
+    const left = Math.round(leftPx + index / bpPerPx)
+    const width = Math.round(3 / bpPerPx) === 0 ? 1 : Math.round(3 / bpPerPx)
+    const codonCode = strand === 1 ? codon : revcom(codon)
+    const aminoAcidCode =
+      defaultCodonTable[codonCode as keyof typeof defaultCodonTable]
+    const fillColor = codonColorCode(aminoAcidCode, theme, highContrast)
+    if (
+      fillColor &&
+      ((showStopCodons && aminoAcidCode == '*') ||
+        (showStartCodons && aminoAcidCode != '*'))
+    ) {
+      ctx.fillStyle = fillColor
+      ctx.fillRect(left, top, width, rowHeight)
+    }
+  }
+}
+
 export function renderingModelFactory(
   pluginManager: PluginManager,
   configSchema: AnyConfigurationSchemaType,
@@ -135,17 +182,29 @@ export function renderingModelFactory(
           self,
           autorun(
             () => {
-              const { canvas, featureLayouts, featuresHeight, lgv } = self
+              const {
+                apolloRowHeight,
+                canvas,
+                featureLayouts,
+                featuresHeight,
+                lgv,
+                session,
+                theme,
+                showFeatureLabels,
+                showStartCodons,
+                showStopCodons,
+              } = self
               if (!lgv.initialized || self.regionCannotBeRendered()) {
                 return
               }
-              const { displayedRegions, dynamicBlocks } = lgv
+              const { bpPerPx, offsetPx, displayedRegions, dynamicBlocks } = lgv
 
               const ctx = canvas?.getContext('2d')
               if (!ctx) {
                 return
               }
               ctx.clearRect(0, 0, dynamicBlocks.totalWidthPx, featuresHeight)
+
               for (const [idx, featureLayout] of featureLayouts.entries()) {
                 const displayedRegion = displayedRegions[idx]
                 for (const [row, featureLayoutRow] of featureLayout.entries()) {
@@ -171,6 +230,45 @@ export function renderingModelFactory(
                   }
                 }
               }
+
+              if (showStartCodons || showStopCodons) {
+                const { apolloDataStore } = session
+                for (const block of dynamicBlocks.contentBlocks) {
+                  const assembly = apolloDataStore.assemblies.get(
+                    block.assemblyName,
+                  )
+                  const ref = assembly?.getByRefName(block.refName)
+                  const roundedStart = Math.floor(block.start)
+                  const roundedEnd = Math.ceil(block.end)
+                  let seq = ref?.getSequence(roundedStart, roundedEnd)
+                  if (!seq) {
+                    break
+                  }
+                  seq = seq.toUpperCase()
+                  const baseOffsetPx = (block.start - roundedStart) / bpPerPx
+                  const seqLeftPx = Math.round(
+                    block.offsetPx - offsetPx - baseOffsetPx,
+                  )
+                  for (let i = 0; i < seq.length; i++) {
+                    const bp = roundedStart + i
+                    const codon = seq.slice(i, i + 3)
+                    drawCodon(
+                      ctx,
+                      codon,
+                      seqLeftPx,
+                      i,
+                      theme,
+                      true,
+                      bpPerPx,
+                      bp,
+                      apolloRowHeight,
+                      showFeatureLabels,
+                      showStartCodons,
+                      showStopCodons,
+                    )
+                  }
+                }
+              }
             },
             { name: 'LinearApolloSixFrameDisplayRenderFeatures' },
           ),

packages/jbrowse-plugin-apollo/src/util/displayUtils.ts

@@ -1,4 +1,5 @@
 import { type CheckResultIdsType } from '@apollo-annotation/mst'
+import { type Theme } from '@mui/material'
 import { makeStyles } from 'tss-react/mui'
 
 export { default as EditZoomThresholdDialog } from '../components/EditZoomThresholdDialog'
@@ -147,3 +148,30 @@ export function clusterResultByMessage<
   )
   return clusters
 }
+
+export function codonColorCode(
+  letter: string,
+  theme: Theme,
+  highContrast?: boolean,
+) {
+  if (letter === 'M') {
+    return theme.palette.startCodon
+  }
+  if (letter === '*') {
+    return highContrast ? theme.palette.text.primary : theme.palette.stopCodon
+  }
+  return
+}
+
+export function colorCode(letter: string, theme: Theme) {
+  const letterUpper = letter.toUpperCase()
+  if (
+    letterUpper === 'A' ||
+    letterUpper === 'C' ||
+    letterUpper === 'G' ||
+    letterUpper === 'T'
+  ) {
+    return theme.palette.bases[letterUpper].main.toString()
+  }
+  return 'lightgray'
+}