react-msaview 6.5.0 → 8.0.0

package/src/components/tree/renderTreeCanvas.ts

@@ -1,31 +1,67 @@
-import {
-  calcDepthToLeaf,
-  forEachDescendant,
-  forEachLink,
-} from '../../hierarchy.ts'
+import { calcDepthToLeaf, nodeCoveringRows } from '../../hierarchy.ts'
 import { setFontSize } from '../../setFontSize.ts'
 import { getVisibleLeaves } from '../getVisibleLeaves.ts'
 import {
   drawHighlightLabel,
+  highlightLabelHeight,
   highlightRowFill,
 } from '../msa/renderHighlights.ts'
+import { bracketBarWidth, bracketGap } from './cladeBrackets.ts'
 
 import type { HierarchyNode } from '../../hierarchy.ts'
 import type { MsaViewModel } from '../../model.ts'
+import type { LabelBox } from '../msa/renderHighlights.ts'
 import type { RenderCtx } from '../renderCtx.ts'
 import type { ClickMapIndex } from './clickMap.ts'
 import type { Theme } from '@mui/material'
 
-export const padding = 600
-
-const extendBounds = 5
 const radius = 2.5
 const d = radius * 2
 
-// whether a node's row-center y falls within the block being drawn, padded by
-// extendBounds so bubbles/labels/triangles straddling the edge still render
-function inYBlock(y: number, offsetY: number, by: number) {
-  return y > offsetY - extendBounds && y < offsetY + by + extendBounds
+// How far outside a block a node's row-center can sit and still draw into it: a
+// collapsed triangle reaches 0.42 rows each way, a tip label an em, and a
+// highlight label its whole box. Five pixels clipped all three at block edges.
+function blockPad(model: MsaViewModel) {
+  return Math.max(model.rowHeight, highlightLabelHeight)
+}
+
+// whether a node's row-center y falls within the block being drawn, padded so
+// bubbles/labels/triangles straddling the edge still render
+function inYBlock(y: number, offsetY: number, by: number, pad: number) {
+  return y > offsetY - pad && y < offsetY + by + pad
+}
+
+/**
+ * Every node whose subtree reaches into the block, skipping the subtrees that
+ * do not. A subtree's leaves are one contiguous run of rows, recorded by
+ * clusterLayout as xMin/xMax, so a block of a 200k-tip tree visits a few dozen
+ * nodes.
+ *
+ * A link out of a visited node is drawn even when its child is pruned, so a
+ * branch crossing the block stays visible.
+ */
+function forEachNodeInBlock(
+  root: HierarchyNode,
+  offsetY: number,
+  by: number,
+  pad: number,
+  cb: (node: HierarchyNode) => void,
+) {
+  const top = offsetY - pad
+  const bottom = offsetY + by + pad
+  const stack = [root]
+  while (stack.length > 0) {
+    const node = stack.pop()!
+    if (node.xMax! < top || node.xMin! > bottom) {
+      continue
+    }
+    cb(node)
+    if (node.children) {
+      for (const child of node.children) {
+        stack.push(child)
+      }
+    }
+  }
 }
 
 // Calculate node x-coordinate for both phylogram (with branch lengths) and
@@ -66,31 +102,46 @@ function renderTree({
   maxDepthToLeaf: number
   blockSizeYOverride?: number
 }) {
-  const { hierarchy, showBranchLenEffective: showBranchLen, blockSize } = model
+  const {
+    hierarchy,
+    showBranchLenEffective: showBranchLen,
+    blockSize,
+    branchColors,
+  } = model
   const by = blockSizeYOverride ?? blockSize
-  ctx.strokeStyle = theme.palette.text.primary
-  forEachLink(hierarchy, (source, target) => {
+  const defaultStroke = theme.palette.text.primary
+  ctx.strokeStyle = defaultStroke
+  forEachNodeInBlock(hierarchy, offsetY, by, blockPad(model), source => {
     const sy = source.x!
-    const ty = target.x!
-    const tx = getNodeX(target, showBranchLen, tipX, maxDepthToLeaf)
     const sx = getNodeX(source, showBranchLen, tipX, maxDepthToLeaf)
-    if (tx === undefined || sx === undefined) {
+    if (sx === undefined || !source.children) {
       return
     }
-
-    const y1 = Math.min(sy, ty)
-    const y2 = Math.max(sy, ty)
-    // 1d line intersection to check if line crosses block at all, this is an
-    // optimization that allows us to skip drawing most tree links outside the
-    // block
-    if (offsetY + by >= y1 && y2 >= offsetY) {
-      ctx.beginPath()
-      ctx.moveTo(sx, sy)
-      ctx.lineTo(sx, ty)
-      ctx.lineTo(tx, ty)
-      ctx.stroke()
+    for (const target of source.children) {
+      const ty = target.x!
+      const tx = getNodeX(target, showBranchLen, tipX, maxDepthToLeaf)
+      if (tx === undefined) {
+        continue
+      }
+      const y1 = Math.min(sy, ty)
+      const y2 = Math.max(sy, ty)
+      // 1d line intersection to check if line crosses block at all, this is an
+      // optimization that allows us to skip drawing most tree links outside the
+      // block
+      if (offsetY + by >= y1 && y2 >= offsetY) {
+        // the edge belongs to the child, whose id the branch channel colors
+        if (branchColors) {
+          ctx.strokeStyle = branchColors.get(target.data.id) ?? defaultStroke
+        }
+        ctx.beginPath()
+        ctx.moveTo(sx, sy)
+        ctx.lineTo(sx, ty)
+        ctx.lineTo(tx, ty)
+        ctx.stroke()
+      }
     }
   })
+  ctx.strokeStyle = defaultStroke
 }
 
 function renderCollapsedTriangles({
@@ -121,6 +172,7 @@ function renderCollapsedTriangles({
     rowHeight,
     fontSize,
     marginLeft: ml,
+    branchColors,
   } = model
   // nothing collapsed is the common case, and the traversal below visits every
   // node in the tree on every block of every redraw
@@ -129,12 +181,12 @@ function renderCollapsedTriangles({
   }
   const by = blockSizeYOverride ?? blockSize
   const halfHeight = Math.max(2, rowHeight * 0.42)
-  forEachDescendant(hierarchy, node => {
+  forEachNodeInBlock(hierarchy, offsetY, by, blockPad(model), node => {
     const { id, name } = node.data
     if (collapsedSet.has(id)) {
       const apexX = getNodeX(node, showBranchLen, tipX, maxDepthToLeaf)
       const y = node.x!
-      const inBlock = inYBlock(y, offsetY, by)
+      const inBlock = inYBlock(y, offsetY, by, blockPad(model))
       // in cladogram mode the hidden tips align at the right edge, otherwise use
       // the branch-length extent recorded in the model's hierarchy getter
       const baseX = showBranchLen ? (node.collapsedTipXFar ?? tipX) : tipX
@@ -144,7 +196,7 @@ function renderCollapsedTriangles({
         ctx.lineTo(baseX, y - halfHeight)
         ctx.lineTo(baseX, y + halfHeight)
         ctx.closePath()
-        ctx.fillStyle = theme.palette.action.disabled
+        ctx.fillStyle = branchColors?.get(id) ?? theme.palette.action.disabled
         ctx.fill()
         ctx.strokeStyle = theme.palette.text.primary
         ctx.stroke()
@@ -179,6 +231,7 @@ function renderCollapsedTriangles({
 // always runs and `draw` only decides whether it also strokes the circles
 function renderNodeBubbles({
   ctx,
+  theme,
   clickMap,
   offsetY,
   model,
@@ -197,6 +250,7 @@ function renderNodeBubbles({
   blockSizeYOverride?: number
   collapsedSet: Set<string>
   draw: boolean
+  theme: Theme
 }) {
   const {
     hierarchy,
@@ -205,7 +259,7 @@ function renderNodeBubbles({
     marginLeft: ml,
   } = model
   const by = blockSizeYOverride ?? blockSize
-  forEachDescendant(hierarchy, node => {
+  forEachNodeInBlock(hierarchy, offsetY, by, blockPad(model), node => {
     const x = getNodeX(node, showBranchLen, tipX, maxDepthToLeaf)
     if (x === undefined) {
       return
@@ -213,11 +267,13 @@ function renderNodeBubbles({
     const { data } = node
     const y = node.x!
     const { id, name } = data
-    if (node.height >= 1 && inYBlock(y, offsetY, by)) {
+    if (node.height >= 1 && inYBlock(y, offsetY, by, blockPad(model))) {
       const isCollapsed = collapsedSet.has(id)
       if (draw) {
-        ctx.strokeStyle = 'black'
-        ctx.fillStyle = isCollapsed ? 'black' : 'white'
+        ctx.strokeStyle = theme.palette.text.primary
+        ctx.fillStyle = isCollapsed
+          ? theme.palette.text.primary
+          : theme.palette.background.default
         ctx.beginPath()
         ctx.arc(x, y, radius, 0, 2 * Math.PI)
         ctx.fill()
@@ -241,6 +297,53 @@ function renderNodeBubbles({
   })
 }
 
+// The label a newick file puts on an internal node, which is where a bootstrap
+// or posterior support value lands: ((A,B)95,(C,D)80);. withId falls back to the
+// path-derived id for an unnamed node, so a node whose name is its id has no
+// label of its own and an ungated pass prints node-0-1-1 across the tree.
+function renderNodeLabels({
+  ctx,
+  theme,
+  offsetY,
+  model,
+  tipX,
+  maxDepthToLeaf,
+  blockSizeYOverride,
+}: {
+  ctx: RenderCtx
+  theme: Theme
+  offsetY: number
+  model: MsaViewModel
+  tipX: number
+  maxDepthToLeaf: number
+  blockSizeYOverride?: number
+}) {
+  const {
+    hierarchy,
+    showBranchLenEffective: showBranchLen,
+    blockSize,
+    fontSize,
+  } = model
+  const by = blockSizeYOverride ?? blockSize
+  const pad = blockPad(model)
+  setFontSize(ctx, Math.max(6, fontSize - 3))
+  ctx.fillStyle = theme.palette.text.primary
+  ctx.textAlign = 'right'
+  forEachNodeInBlock(hierarchy, offsetY, by, pad, node => {
+    const { id, name } = node.data
+    if (node.height < 1 || name === id) {
+      return
+    }
+    const x = getNodeX(node, showBranchLen, tipX, maxDepthToLeaf)
+    const y = node.x!
+    if (x !== undefined && inYBlock(y, offsetY, by, pad)) {
+      ctx.fillText(name, x - radius - 1, y - radius - 1)
+    }
+  })
+  ctx.textAlign = 'start'
+  setFontSize(ctx, fontSize)
+}
+
 function renderTreeLabels({
   theme,
   model,
@@ -265,7 +368,7 @@ function renderTreeLabels({
   const {
     fontSize,
     showBranchLenEffective: showBranchLen,
-    treeMetadata,
+    tipLabelColors,
     blockSize,
     labelsAlignRight,
     drawTree,
@@ -279,6 +382,7 @@ function renderTreeLabels({
   // labels only exist for leaves, which are laid out top to bottom, so take the
   // same slice the MSA renderer uses instead of walking every tip per block
   const visibleLeaves = getVisibleLeaves({ model, offsetY, blockSizeY: by })
+  const pad = blockPad(model)
   const emHeight = ctx.measureText('M').width
   if (labelsAlignRight) {
     ctx.textAlign = 'right'
@@ -293,11 +397,11 @@ function renderTreeLabels({
     } = node
     const y = node.x!
 
-    const displayName = treeMetadata[name]?.genome || name
+    const displayName = model.rowDataOf(name)?.genome || name
     // a collapsed clade is drawn as a triangle + tip count; suppress its leaf
     // label when the "name" is just the auto-generated internal-node id
     const isAnonymousCollapsed = collapsedSet.has(id) && name === id
-    if (!isAnonymousCollapsed && inYBlock(y, offsetY, by)) {
+    if (!isAnonymousCollapsed && inYBlock(y, offsetY, by, pad)) {
       // note: +rowHeight/4 matches with -rowHeight/4 in msa
       const yp = y + fontSize / 4
       let xp = 0
@@ -311,7 +415,7 @@ function renderTreeLabels({
           ? ctx.measureText(displayName).width
           : measured * labelWidthScale
 
-      ctx.fillStyle = theme.palette.text.primary
+      ctx.fillStyle = tipLabelColors?.get(name) ?? theme.palette.text.primary
       if (labelsAlignRight) {
         const smallPadding = 2
         const offset = treeAreaWidthMinusMargin - smallPadding
@@ -353,8 +457,123 @@ function renderTreeLabels({
   ctx.setLineDash([])
 }
 
-// the row sets of `highlights`, tinted across the tree area under the labels
-// with the label in the gutter at the first row of each set
+// the first row of a set, as a loop: spreading into Math.min passes one
+// argument per row, and the ~125k argument limit throws RangeError on a set
+// covering the 230k-tip COVID tree
+function firstRow(rowIndices: number[]) {
+  let min = Infinity
+  for (const index of rowIndices) {
+    if (index < min) {
+      min = index
+    }
+  }
+  return min
+}
+
+// the highlight mark of the `clades` layer: a translucent rectangle from the
+// clade's common ancestor to the right edge of the tree area, under the
+// branches. A clade spanning hundreds of rows crosses several blocks, so each
+// block draws the whole rectangle and the canvas takes the part that is its.
+function renderCladeHighlights({
+  ctx,
+  model,
+  offsetY,
+  tipX,
+  maxDepthToLeaf,
+  blockSizeYOverride,
+}: {
+  ctx: RenderCtx
+  model: MsaViewModel
+  offsetY: number
+  tipX: number
+  maxDepthToLeaf: number
+  blockSizeYOverride?: number
+}) {
+  const {
+    resolvedClades,
+    hierarchy,
+    showBranchLenEffective: showBranchLen,
+    rowHeight,
+    treeAreaWidth,
+    marginLeft,
+    blockSize,
+  } = model
+  const by = blockSizeYOverride ?? blockSize
+  const right = treeAreaWidth - marginLeft
+  for (const { rows, color, mark } of resolvedClades) {
+    if (mark !== 'highlight') {
+      continue
+    }
+    const top = rows[0] * rowHeight
+    const bottom = (rows[1] + 1) * rowHeight
+    if (bottom < offsetY || top > offsetY + by) {
+      continue
+    }
+    const node = nodeCoveringRows(
+      hierarchy,
+      top + rowHeight / 2,
+      bottom - rowHeight / 2,
+    )
+    const x = getNodeX(node, showBranchLen, tipX, maxDepthToLeaf) ?? 0
+    ctx.fillStyle = color
+    ctx.fillRect(x, top, right - x, bottom - top)
+  }
+}
+
+// the bar of the bracket mark, in the gutter `cladeGutterWidth` reserves at the
+// right of the tree area. The label beside it is DOM text on screen and a
+// <text> in the export, since the canvas layer has no rotation.
+function renderCladeBrackets({
+  ctx,
+  model,
+  theme,
+  offsetY,
+  blockSizeYOverride,
+}: {
+  ctx: RenderCtx
+  model: MsaViewModel
+  theme: Theme
+  offsetY: number
+  blockSizeYOverride?: number
+}) {
+  const {
+    resolvedClades,
+    cladeGutterWidth,
+    rowHeight,
+    treeAreaWidth,
+    marginLeft,
+    blockSize,
+  } = model
+  if (cladeGutterWidth === 0) {
+    return
+  }
+  const by = blockSizeYOverride ?? blockSize
+  const x =
+    treeAreaWidth -
+    marginLeft -
+    cladeGutterWidth +
+    bracketGap +
+    bracketBarWidth / 2
+  ctx.lineWidth = bracketBarWidth
+  for (const { rows, mark, markColor } of resolvedClades) {
+    const top = rows[0] * rowHeight + 1
+    const bottom = (rows[1] + 1) * rowHeight - 1
+    if (mark !== 'bracket' || bottom < offsetY || top > offsetY + by) {
+      continue
+    }
+    ctx.strokeStyle = markColor ?? theme.palette.text.primary
+    ctx.beginPath()
+    ctx.moveTo(x, top)
+    ctx.lineTo(x, bottom)
+    ctx.stroke()
+  }
+  ctx.lineWidth = 1
+  ctx.strokeStyle = theme.palette.text.primary
+}
+
+// the `rowTint` encoding and the row sets of `highlights`, washed across the
+// tree area under the labels, with each set's label in the gutter at its first
+// row
 function renderRowHighlights({
   ctx,
   model,
@@ -370,25 +589,44 @@ function renderRowHighlights({
 }) {
   const {
     resolvedHighlights,
+    rowTints,
     rowHeight,
     treeAreaWidth,
     marginLeft,
     blockSize,
   } = model
   const by = blockSizeYOverride ?? blockSize
+  const pad = blockPad(model)
+  const placed: LabelBox[] = []
+  const fillRow = (index: number) => {
+    ctx.fillRect(-marginLeft, index * rowHeight, treeAreaWidth, rowHeight)
+  }
+  if (rowTints) {
+    const top = Math.max(0, Math.floor((offsetY - pad) / rowHeight))
+    const bottom = Math.min(
+      rowTints.length - 1,
+      Math.ceil((offsetY + by + pad) / rowHeight),
+    )
+    for (let index = top; index <= bottom; index++) {
+      const color = rowTints[index]
+      if (color) {
+        ctx.fillStyle = color
+        fillRow(index)
+      }
+    }
+  }
   for (const { rowIndices, label, color } of resolvedHighlights) {
     ctx.fillStyle = color ?? highlightRowFill
     for (const index of rowIndices) {
-      const y = index * rowHeight
-      if (inYBlock(y + rowHeight / 2, offsetY, by)) {
-        ctx.fillRect(-marginLeft, y, treeAreaWidth, rowHeight)
+      if (inYBlock(index * rowHeight + rowHeight / 2, offsetY, by, pad)) {
+        fillRow(index)
       }
     }
-    const first = Math.min(...rowIndices)
+    const first = firstRow(rowIndices)
     if (
       label &&
       rowIndices.length &&
-      inYBlock(first * rowHeight, offsetY, by)
+      inYBlock(first * rowHeight, offsetY, by, pad)
     ) {
       drawHighlightLabel({
         ctx,
@@ -397,6 +635,8 @@ function renderRowHighlights({
         x: -marginLeft + 2,
         y: first * rowHeight,
         spanWidth: treeAreaWidth,
+        bounds: { min: -marginLeft, max: treeAreaWidth - marginLeft },
+        placed,
       })
     }
   }
@@ -426,6 +666,7 @@ export function renderTreeCanvas({
     noTree,
     drawTree,
     drawNodeBubbles,
+    drawNodeLabels,
     highResScaleFactor,
     fontSize,
     showTreeText,
@@ -456,6 +697,14 @@ export function renderTreeCanvas({
   // every block of every redraw, and `collapsed` is a list
   const collapsedSet = new Set(model.collapsed)
 
+  renderCladeHighlights({
+    ctx,
+    model,
+    offsetY,
+    tipX,
+    maxDepthToLeaf,
+    blockSizeYOverride,
+  })
   renderRowHighlights({ ctx, model, theme, offsetY, blockSizeYOverride })
   setFontSize(ctx, fontSize)
 
@@ -484,6 +733,7 @@ export function renderTreeCanvas({
 
     renderNodeBubbles({
       ctx,
+      theme,
       offsetY,
       clickMap,
       model,
@@ -493,6 +743,18 @@ export function renderTreeCanvas({
       collapsedSet,
       draw: drawNodeBubbles,
     })
+
+    if (drawNodeLabels) {
+      renderNodeLabels({
+        ctx,
+        theme,
+        offsetY,
+        model,
+        tipX,
+        maxDepthToLeaf,
+        blockSizeYOverride,
+      })
+    }
   }
 
   if (showTreeText) {
@@ -509,6 +771,8 @@ export function renderTreeCanvas({
     })
   }
 
+  renderCladeBrackets({ ctx, model, theme, offsetY, blockSizeYOverride })
+
   // Finish the index so it's ready for queries
   clickMap?.finish()
 }