Single-cell UMAP

import { Suspense, use, useEffect, useRef, useState } from 'react'

import { isFeature } from '@jbrowse/core/util'
import {
  JBrowseLinearGenomeView,
  loadPlugins,
  useCreateViewState,
} from '@jbrowse/react-linear-genome-view2'
import { observer } from 'mobx-react'

import type { MultiWiggleDisplayModel } from '@jbrowse/plugin-wiggle'
import type { ViewModel } from '@jbrowse/react-linear-genome-view2'

interface CellType {
  name: string
  group: string
  color: string
}

interface Cells {
  dataset: string
  cellTypes: CellType[]
  x: number[]
  y: number[]
  type: number[]
  genes: string[]
  geneLoc: string[]
  perCellGenes: string[]
  exprIndex: [number, number, number][]
  exprUrl: string
  perCellUrl: string
}

const WIDTH = 560
const HEIGHT = 400
const PAD = 14
const DOT = 2.6
const GREY = '#d8dbe0'

function geneValues(cells: Cells, expr: Uint8Array, gene: string) {
  const g = cells.genes.indexOf(gene)
  const out = new Uint8Array(cells.x.length)
  if (g !== -1) {
    const [offset, count] = cells.exprIndex[g]!
    for (let i = 0; i < count; i++) {
      const at = (offset + i) * 3
      out[expr[at]! | (expr[at + 1]! << 8)] = expr[at + 2]!
    }
  }
  return out
}

function ramp(v: number) {
  const t = v / 255
  const r = Math.round(232 + (8 - 232) * t)
  const g = Math.round(238 + (48 - 238) * t)
  const b = Math.round(247 + (107 - 247) * t)
  return `rgb(${r},${g},${b})`
}

function UmapScatter({
  cells,
  expr,
  gene,
  selected,
  onPick,
}: {
  cells: Cells
  expr: Uint8Array
  gene?: string
  selected: string[]
  onPick: (cellType: string) => void
}) {
  const ref = useRef<HTMLCanvasElement>(null)

  useEffect(() => {
    const ctx = ref.current?.getContext('2d')
    if (ctx) {
      const dpr = 2
      ctx.setTransform(dpr, 0, 0, dpr, 0, 0)
      ctx.clearRect(0, 0, WIDTH, HEIGHT)
      const px = (i: number) => PAD + cells.x[i]! * (WIDTH - 2 * PAD)
      const py = (i: number) => HEIGHT - PAD - cells.y[i]! * (HEIGHT - 2 * PAD)
      const values = gene ? geneValues(cells, expr, gene) : undefined
      const active = new Set(selected)

      const order = [...cells.x.keys()].sort((a, b) =>
        values
          ? values[a]! - values[b]!
          : Number(
              active.size > 0 &&
                active.has(cells.cellTypes[cells.type[a]!]!.name),
            ) -
            Number(
              active.size > 0 &&
                active.has(cells.cellTypes[cells.type[b]!]!.name),
            ),
      )
      for (const i of order) {
        const type = cells.cellTypes[cells.type[i]!]!
        ctx.fillStyle = values
          ? ramp(values[i]!)
          : active.size === 0 || active.has(type.name)
            ? type.color
            : GREY
        ctx.beginPath()
        ctx.arc(px(i), py(i), DOT, 0, 2 * Math.PI)
        ctx.fill()
      }
    }
  }, [cells, expr, gene, selected])

  return (
    <div>
      <canvas
        ref={ref}
        width={WIDTH * 2}
        height={HEIGHT * 2}
        style={{
          width: WIDTH,
          height: HEIGHT,
          maxWidth: '100%',
          cursor: 'pointer',
        }}
        onClick={event => {
          const box = event.currentTarget.getBoundingClientRect()
          const scale = WIDTH / box.width
          const mx = (event.clientX - box.left) * scale
          const my = (event.clientY - box.top) * scale
          let best = -1
          let bestDist = 64
          for (let i = 0; i < cells.x.length; i++) {
            const dx = PAD + cells.x[i]! * (WIDTH - 2 * PAD) - mx
            const dy = HEIGHT - PAD - cells.y[i]! * (HEIGHT - 2 * PAD) - my
            const d = dx * dx + dy * dy
            if (d < bestDist) {
              bestDist = d
              best = i
            }
          }
          if (best !== -1) {
            onPick(cells.cellTypes[cells.type[best]!]!.name)
          }
        }}
      />
      <div
        style={{ display: 'flex', flexWrap: 'wrap', gap: 8, maxWidth: WIDTH }}
      >
        {cells.cellTypes.map(type => (
          <button
            key={type.name}
            type="button"
            onClick={() => {
              onPick(type.name)
            }}
            style={{
              display: 'flex',
              alignItems: 'center',
              gap: 5,
              padding: '2px 7px',
              cursor: 'pointer',
              border: selected.includes(type.name)
                ? '2px solid CanvasText'
                : '2px solid transparent',
              background: 'none',
              font: 'inherit',
              fontSize: '0.85em',
            }}
          >
            <span
              style={{
                width: 10,
                height: 10,
                borderRadius: 5,
                background: type.color,
              }}
            />
            {type.name}
          </button>
        ))}
      </div>
    </div>
  )
}

const BASE = 'https://jbrowse.org/demos/scrna_pbmc5k'
const TRACK_ID = 'pbmc5k_scrna_pseudobulk'
const PER_CELL_TRACK_ID = 'pbmc5k_scrna_percell'
const ZARR_PLUGIN =
  'https://jbrowse.org/demos/zarr/jbrowse-plugin-zarr.umd.production.min.js'

const dataPromise = Promise.all([
  fetch(`${BASE}/cells.json`).then(res => res.json() as Promise<Cells>),
  fetch(`${BASE}/expr.bin`)
    .then(res => res.arrayBuffer())
    .then(buf => new Uint8Array(buf)),
  loadPlugins([{ name: 'Zarr', url: ZARR_PLUGIN }]).then(loaded =>
    loaded.map(p => p.plugin),
  ),
])

const assembly = {
  name: 'GRCh38',
  aliases: ['hg38'],
  uri: 'https://s3.amazonaws.com/jbrowse.org/genomes/GRCh38/fasta/GRCh38.fa.gz',
  refNameAliases: {
    uri: 'https://s3.amazonaws.com/jbrowse.org/genomes/GRCh38/hg38_aliases.txt',
  },
  geneticCodes: { MT: 2 },
}

function tracks(cells: Cells) {
  return [
    {
      type: 'FeatureTrack',
      trackId: 'hg38_refseq_curated',
      name: 'NCBI RefSeq genes',
      assemblyNames: ['GRCh38'],
      adapter: {
        type: 'Gff3TabixAdapter',
        uri: 'https://jbrowse.org/ucsc/hg38/ncbiRefSeqCurated.gff.gz',
        csi: true,
      },
      displayDefaults: { height: 90 },
    },
    {
      type: 'MultiQuantitativeTrack',
      trackId: TRACK_ID,
      name: 'scRNA pseudobulk by cell type',
      assemblyNames: ['GRCh38'],
      adapter: {
        type: 'MultiWiggleAdapter',
        subadapters: cells.cellTypes.map(type => ({
          type: 'BigWigAdapter',
          name: type.name,
          group: type.group,
          color: type.color,
          uri: `${BASE}/${type.name.replaceAll(' ', '_')}.bw`,
        })),
      },
      displayDefaults: { defaultRendering: 'multirowxy', height: 330 },
    },
    {
      type: 'MultiQuantitativeTrack',
      trackId: PER_CELL_TRACK_ID,
      name: 'Per-cell coverage (marker loci)',
      assemblyNames: ['GRCh38'],
      adapter: {
        type: 'MultiWiggleZarrAdapter',
        uri: `${BASE}/percell.zarr`,
      },
      displayDefaults: {
        defaultRendering: 'multirowdensity',
        height: 420,
        minScore: 0,
        maxScore: 4,
      },
    },
  ]
}

function Demo() {
  const [cells, expr, plugins] = use(dataPromise)
  const state = useCreateViewState({
    assembly,
    plugins,
    tracks: tracks(cells),
    defaultSession: {
      name: 'scRNA pseudobulk',
      view: {
        type: 'LinearGenomeView',
        assembly: 'GRCh38',
        loc: '11:60,453,846-60,472,752',
        tracks: ['hg38_refseq_curated', TRACK_ID, PER_CELL_TRACK_ID],
      },
    },
  })
  return state ? (
    <UmapAndGenomeView cells={cells} expr={expr} state={state} />
  ) : null
}

const UmapAndGenomeView = observer(function UmapAndGenomeView({
  cells,
  expr,
  state,
}: {
  cells: Cells
  expr: Uint8Array
  state: ViewModel
}) {
  const [picked, setPicked] = useState<string>()
  const [selected, setSelected] = useState<string[]>([])
  const { session } = state
  const { selection, view } = session
  const clicked = isFeature(selection)
    ? String(selection.get('name'))
    : undefined
  const gene = clicked && cells.genes.includes(clicked) ? clicked : picked

  function pickCellType(name: string) {
    const next = selected.includes(name)
      ? selected.filter(t => t !== name)
      : [...selected, name]
    setSelected(next)
    const display = view.getTrack(TRACK_ID)?.activeDisplay as
      | MultiWiggleDisplayModel
      | undefined
    display?.setSubtreeFilter(next.length ? next : undefined)
  }

  return (
    <div>
      <div style={{ display: 'flex', gap: 20, flexWrap: 'wrap' }}>
        <div>
          <UmapScatter
            cells={cells}
            expr={expr}
            gene={gene}
            selected={selected}
            onPick={name => {
              pickCellType(name)
            }}
          />
        </div>
        <div style={{ minWidth: 220 }}>
          <p style={{ marginTop: 0 }}>
            <b>{cells.dataset}</b>
          </p>
          <label>
            Color by gene{' '}
            <select
              value={gene ?? ''}
              onChange={event => {
                const name = event.target.value
                session.clearSelection()
                setPicked(name || undefined)
                const at = cells.genes.indexOf(name)
                if (at !== -1) {
                  void view.navToLocString(cells.geneLoc[at]!)
                }
              }}
            >
              <option value="">cell type</option>
              <optgroup label="With per-cell rows">
                {cells.perCellGenes.map(name => (
                  <option key={name} value={name}>
                    {name}
                  </option>
                ))}
              </optgroup>
              <optgroup label="Expression panel only">
                {cells.genes
                  .filter(name => !cells.perCellGenes.includes(name))
                  .map(name => (
                    <option key={name} value={name}>
                      {name}
                    </option>
                  ))}
              </optgroup>
            </select>
          </label>
          <p>
            {gene
              ? `Cells colored by ${gene} expression.`
              : 'Click a cluster or a legend swatch to filter the coverage rows to those cell types.'}
          </p>
          {clicked && !gene ? (
            <p>
              No expression in this demo's panel for <b>{clicked}</b>.
            </p>
          ) : null}
        </div>
      </div>
      <JBrowseLinearGenomeView viewState={state} />
    </div>
  )
})

export default function SingleCellUmap() {
  return (
    <Suspense fallback={<p>Loading cells...</p>}>
      <Demo />
    </Suspense>
  )
}