Pan-UKB GWAS

import { useEffect, useMemo, useState } from 'react'

import {
  JBrowseLinearGenomeView,
  useCreateViewState,
} from '@jbrowse/react-linear-genome-view2'

// Pan-UKBB per-phenotype flat files (tabix-indexed TSV, GRCh37/hg38-aliased)
const BASE = 'https://pan-ukb-us-east-1.s3.amazonaws.com/sumstats_flat_files'

// trimmed phenotype manifest (~7,200 phenotypes), generated by
// scripts/gen-panukb-manifest.mjs and hosted at jbrowse.org so the 2MB file
// doesn't live in this repo
const MANIFEST_URL = 'https://jbrowse.org/demos/panukbb/panukbPhenotypes.json'

interface Phenotype {
  id: string
  traitType: string
  description: string
  category: string
  n: number
  pops: string[]
  popsQc: string[]
}

const POP_LABELS: Record<string, string> = {
  EUR: 'European',
  AFR: 'African',
  AMR: 'Admixed American',
  CSA: 'Central/South Asian',
  EAS: 'East Asian',
  MID: 'Middle Eastern',
}

// hand-picked phenotypes with a known lead locus to zoom to on selection, so
// the demo opens on signal rather than an arbitrary window
const FEATURED: Record<string, string> = {
  'continuous-50-both_sexes-irnt': 'chr12:64,000,000..67,000,000', // height / HMGA2
  'continuous-21001-both_sexes-irnt': 'chr16:53,000,000..55,000,000', // BMI / FTO
  'continuous-4079-both_sexes-irnt': 'chr4:81,000,000..82,500,000', // DBP / FGF5
  'biomarkers-30690-both_sexes-irnt': 'chr1:54,500,000..56,000,000', // cholesterol / PCSK9
  'biomarkers-30780-both_sexes-irnt': 'chr19:11,000,000..11,500,000', // LDL / LDLR
  'icd10-E11-both_sexes': 'chr10:112,500,000..113,500,000', // T2D / TCF7L2
  'icd10-I25-both_sexes': 'chr9:21,500,000..22,500,000', // CAD / 9p21
  'icd10-J45-both_sexes': 'chr17:37,500,000..38,500,000', // asthma / ORMDL3
}

function popOptions(p: Phenotype) {
  return [
    ...(p.popsQc.length
      ? [{ col: 'neglog10_pval_meta_hq', label: 'Meta-analysis (QC pass)' }]
      : []),
    { col: 'neglog10_pval_meta', label: 'Meta-analysis (all)' },
    ...p.pops.map(pop => ({
      col: `neglog10_pval_${pop}`,
      label: POP_LABELS[pop] ?? pop,
    })),
  ]
}

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

const NCBI_REFSEQ_TRACK = {
  type: 'FeatureTrack',
  trackId: 'ncbi_refseq_hg38',
  name: 'NCBI RefSeq genes',
  assemblyNames: ['hg38'],
  adapter: {
    type: 'Gff3TabixAdapter',
    // `csi: true` makes the `uri` shorthand resolve a `.csi` index
    uri: 'https://jbrowse.org/ucsc/hg38/ncbiRefSeq.gff.gz',
    csi: true,
  },
  displayDefaults: {
    height: 200,
    labels: {
      name: "jexl:get(feature,'gene_id') || get(feature,'name') || get(feature,'id')",
    },
  },
}

function makeTrack(p: Phenotype, scoreColumn: string, label: string) {
  return {
    type: 'GWASTrack',
    trackId: 'panukb_gwas',
    name: `${p.description} — ${label}`,
    assemblyNames: ['hg38'],
    adapter: {
      type: 'GWASAdapter',
      // Pan-UKBB neglog10_pval_* columns are already -log10(p), so no transform
      scoreColumn,
      uri: `${BASE}/${p.id}.tsv.bgz`,
    },
    displayDefaults: { height: 250 },
  }
}

// Rebuilt on every phenotype/population change via the `key` below, so this is
// the case the managed <LinearGenomeView> is explicitly not for: that component
// owns its engine for the lifetime of the page and never tears it down, which
// is a whole orphaned RPC worker pool and autorun set per switch.
// useCreateViewState destroys the engine when the component unmounts, so the
// switch costs nothing.
function GenomeView({
  phenotype,
  scoreColumn,
  label,
}: {
  phenotype: Phenotype
  scoreColumn: string
  label: string
}) {
  const state = useCreateViewState({
    assembly,
    tracks: [makeTrack(phenotype, scoreColumn, label), NCBI_REFSEQ_TRACK],
    defaultSession: {
      name: 'Pan-UKB GWAS',
      view: {
        type: 'LinearGenomeView',
        init: {
          assembly: assembly.name,
          loc: FEATURED[phenotype.id] ?? 'chr1',
          tracks: ['panukb_gwas', 'ncbi_refseq_hg38'],
        },
      },
    },
  })
  return <JBrowseLinearGenomeView viewState={state} />
}

export default function PanUKBGWAS() {
  const [phenotypes, setPhenotypes] = useState<Phenotype[]>()
  const [query, setQuery] = useState('')
  const [selected, setSelected] = useState<Phenotype>()
  const [scoreColumn, setScoreColumn] = useState('neglog10_pval_meta_hq')

  useEffect(() => {
    const ac = new AbortController()
    fetch(MANIFEST_URL, { signal: ac.signal })
      .then(res => res.json() as Promise<Phenotype[]>)
      .then(data => {
        setPhenotypes(data)
        setSelected(
          data.find(p => p.id === 'continuous-50-both_sexes-irnt') ?? data[0],
        )
      })
      .catch((e: unknown) => {
        if (!ac.signal.aborted) {
          console.error(e)
        }
      })
    return () => {
      ac.abort()
    }
  }, [])

  const results = useMemo(() => {
    if (!phenotypes) {
      return []
    }
    const q = query.trim().toLowerCase()
    const matches = q
      ? phenotypes.filter(
          p =>
            p.description.toLowerCase().includes(q) ||
            p.category.toLowerCase().includes(q) ||
            p.id.toLowerCase().includes(q),
        )
      : phenotypes
    return matches.slice(0, 100)
  }, [phenotypes, query])

  const options = selected ? popOptions(selected) : []
  // keep the population valid as the selected phenotype changes
  const activeCol = options.some(o => o.col === scoreColumn)
    ? scoreColumn
    : (options[0]?.col ?? 'neglog10_pval_meta')
  const activeLabel = options.find(o => o.col === activeCol)?.label ?? activeCol

  return (
    <div>
      <div style={{ display: 'flex', gap: 16, marginBottom: 8 }}>
        <div style={{ flex: 1 }}>
          <input
            type="text"
            placeholder="Search ~7,200 phenotypes (e.g. height, diabetes, LDL)…"
            value={query}
            onChange={e => {
              setQuery(e.target.value)
            }}
            style={{ width: '100%', padding: 4, boxSizing: 'border-box' }}
          />
          <select
            size={6}
            value={selected?.id ?? ''}
            onChange={e => {
              const p = phenotypes?.find(x => x.id === e.target.value)
              if (p) {
                setSelected(p)
              }
            }}
            style={{ width: '100%', marginTop: 4 }}
          >
            {results.map(p => (
              <option key={p.id} value={p.id}>
                {p.description} [{p.traitType}, n={p.n.toLocaleString()}]
              </option>
            ))}
          </select>
          {/* GrayText, not a hex grey: a system color follows the host page's
              light/dark scheme, and #666 is unreadable on a dark one */}
          <div style={{ fontSize: '0.8em', color: 'GrayText' }}>
            {!phenotypes
              ? 'Loading phenotype catalog…'
              : query
                ? `${results.length}${results.length === 100 ? '+' : ''} match`
                : `${phenotypes.length.toLocaleString()} phenotypes`}
          </div>
        </div>
        <label>
          Population:{' '}
          <select
            value={activeCol}
            onChange={e => {
              setScoreColumn(e.target.value)
            }}
          >
            {options.map(o => (
              <option key={o.col} value={o.col}>
                {o.label}
              </option>
            ))}
          </select>
        </label>
      </div>
      {selected ? (
        <GenomeView
          key={`${selected.id}-${activeCol}`}
          phenotype={selected}
          scoreColumn={activeCol}
          label={activeLabel}
        />
      ) : null}
    </div>
  )
}