Update STR page allele size histogram to show colors based on genotype quality or other values

.gitignore

@@ -33,3 +33,6 @@ hail-*.log
 # Playright test dirs
 /tests/playwright/
 /playwright/.cache/
+
+# Reads metadata databases
+reads/*.db

browser/package.json

@@ -23,6 +23,9 @@
     "@gnomad/ui": "2.0.0",
     "@hot-loader/react-dom": "^17.0.0",
     "@visx/axis": "^3.0.0",
+    "@visx/legend": "^3.12.0",
+    "@visx/scale": "^3.12.0",
+    "@visx/shape": "^3.12.0",
     "core-js": "3.5.0",
     "css-loader": "^6.7.3",
     "d3-array": "^1.2.4",

browser/src/ShortTandemRepeatPage/ShortTandemRepeatAdjacentRepeatSection.tsx

@@ -1,52 +1,69 @@
-import { max, min } from 'd3-array'
-import React, { useState } from 'react'
+import React, { SetStateAction, useState, Dispatch } from 'react'
 
 import { Modal, Select } from '@gnomad/ui'
 
 import ControlSection from '../VariantPage/ControlSection'
 
 import ShortTandemRepeatPopulationOptions from './ShortTandemRepeatPopulationOptions'
 import { ShortTandemRepeatAdjacentRepeat } from './ShortTandemRepeatPage'
-import ShortTandemRepeatAlleleSizeDistributionPlot from './ShortTandemRepeatAlleleSizeDistributionPlot'
-import ShortTandemRepeatGenotypeDistributionPlot from './ShortTandemRepeatGenotypeDistributionPlot'
+import ShortTandemRepeatAlleleSizeDistributionPlot, {
+  ScaleType,
+  Sex,
+  ColorBy,
+} from './ShortTandemRepeatAlleleSizeDistributionPlot'
+import ShortTandemRepeatGenotypeDistributionPlot, {
+  Bin as GenotypeBin,
+} from './ShortTandemRepeatGenotypeDistributionPlot'
 import ShortTandemRepeatGenotypeDistributionBinDetails from './ShortTandemRepeatGenotypeDistributionBinDetails'
 import ShortTandemRepeatGenotypeDistributionRepeatUnitsSelect from './ShortTandemRepeatGenotypeDistributionRepeatUnitsSelect'
 import ShortTandemRepeatAdjacentRepeatAttributes from './ShortTandemRepeatAdjacentRepeatAttributes'
 import {
   getSelectedAlleleSizeDistribution,
   getSelectedGenotypeDistribution,
   getGenotypeDistributionPlotAxisLabels,
+  genotypeRepunitPairs,
+  maxAlleleSizeDistributionRepeats,
+  maxGenotypeDistributionRepeats,
 } from './shortTandemRepeatHelpers'
+import { PopulationId } from '@gnomad/dataset-metadata/gnomadPopulations'
 
 type Props = {
   adjacentRepeat: ShortTandemRepeatAdjacentRepeat
-  populationIds: string[]
-  selectedPopulationId: string
-  onSelectPopulationId: (...args: any[]) => any
-  selectedScaleType: string
-  onSelectScaleType: (...args: any[]) => any
+  selectedScaleType: ScaleType
+  selectedPopulation: PopulationId | ''
+  selectedSex: Sex | ''
+  selectedColorBy: ColorBy | ''
+  populations: PopulationId[]
+  selectedGenotypeDistributionBin: GenotypeBin | null
+  setSelectedGenotypeDistributionBin: Dispatch<SetStateAction<GenotypeBin | null>>
+  setSelectedScaleType: Dispatch<SetStateAction<ScaleType>>
+  setSelectedPopulation: Dispatch<SetStateAction<PopulationId | ''>>
+  setSelectedSex: Dispatch<SetStateAction<Sex | ''>>
 }
 
 const ShortTandemRepeatAdjacentRepeatSection = ({
   adjacentRepeat,
-  populationIds,
-  selectedPopulationId,
-  onSelectPopulationId,
+  populations,
   selectedScaleType,
-  onSelectScaleType,
+  selectedPopulation,
+  selectedSex,
+  selectedColorBy,
+  setSelectedScaleType,
+  setSelectedPopulation,
+  setSelectedSex,
 }: Props) => {
   const [selectedRepeatUnit, setSelectedRepeatUnit] = useState(
     adjacentRepeat.repeat_units.length === 1 ? adjacentRepeat.repeat_units[0] : ''
   )
 
+  const genotypeDistributionPairs = genotypeRepunitPairs(adjacentRepeat)
+  const defaultGenotypeDistributionRepeatUnits =
+    genotypeDistributionPairs.length === 1 ? genotypeDistributionPairs[0] : ''
   const [selectedGenotypeDistributionRepeatUnits, setSelectedGenotypeDistributionRepeatUnits] =
-    useState(
-      adjacentRepeat.genotype_distribution.repeat_units.length === 1
-        ? adjacentRepeat.genotype_distribution.repeat_units[0].repeat_units.join(' / ')
-        : ''
-    )
+    useState<string[] | ''>(defaultGenotypeDistributionRepeatUnits)
 
-  const [selectedGenotypeDistributionBin, setSelectedGenotypeDistributionBin] = useState(null)
+  const [selectedGenotypeDistributionBin, setSelectedGenotypeDistributionBin] =
+    useState<GenotypeBin | null>(null)
 
   return (
     <section style={{ marginBottom: '3em' }}>
@@ -55,33 +72,33 @@ const ShortTandemRepeatAdjacentRepeatSection = ({
 
       <h4 style={{ marginBottom: '0.66em' }}>Allele Size Distribution</h4>
       <ShortTandemRepeatAlleleSizeDistributionPlot
-        // @ts-expect-error TS(2322) FIXME: Type '{ maxRepeats: number; alleleSizeDistribution... Remove this comment to see the full error message
-        maxRepeats={
-          adjacentRepeat.allele_size_distribution.distribution[
-            adjacentRepeat.allele_size_distribution.distribution.length - 1
-          ][0]
-        }
+        maxRepeats={maxAlleleSizeDistributionRepeats(adjacentRepeat)}
         alleleSizeDistribution={getSelectedAlleleSizeDistribution(adjacentRepeat, {
-          selectedPopulationId,
+          selectedPopulation,
+          selectedSex,
+          selectedColorBy,
           selectedRepeatUnit,
         })}
+        colorBy={selectedColorBy}
         repeatUnitLength={selectedRepeatUnit ? selectedRepeatUnit.length : null}
         scaleType={selectedScaleType}
       />
       <ControlSection>
         <ShortTandemRepeatPopulationOptions
           id={`${adjacentRepeat.id}-repeat-counts`}
-          populationIds={populationIds}
-          selectedPopulationId={selectedPopulationId}
-          onSelectPopulationId={onSelectPopulationId}
+          populations={populations}
+          selectedPopulation={selectedPopulation}
+          selectedSex={selectedSex}
+          setSelectedPopulation={setSelectedPopulation}
+          setSelectedSex={setSelectedSex}
         />
 
         <label htmlFor={`short-tandem-repeat-${adjacentRepeat.id}-repeat-unit`}>
           Repeat unit: {/* @ts-expect-error TS(2769) FIXME: No overload matches this call. */}
           <Select
             id={`short-tandem-repeat-${adjacentRepeat.id}-repeat-unit`}
             value={selectedRepeatUnit}
-            onChange={(e: any) => {
+            onChange={(e: { target: { value: string } }) => {
               setSelectedRepeatUnit(e.target.value)
             }}
           >
@@ -98,59 +115,63 @@ const ShortTandemRepeatAdjacentRepeatSection = ({
         </label>
 
         <label htmlFor={`short-tandem-repeat-${adjacentRepeat.id}-repeat-counts-scale`}>
-          Scale: {/* @ts-expect-error TS(2769) FIXME: No overload matches this call. */}
+          y-scale: {/* @ts-expect-error TS(2769) FIXME: No overload matches this call. */}
           <Select
             id={`short-tandem-repeat-${adjacentRepeat.id}-repeat-counts-scale`}
             value={selectedScaleType}
-            onChange={(e: any) => {
-              onSelectScaleType(e.target.value)
+            onChange={(e: { target: { value: ScaleType } }) => {
+              setSelectedScaleType(e.target.value)
             }}
           >
             <option value="linear">Linear</option>
             <option value="log">Log</option>
+            <option value="linear-truncated-50">Linear: Truncated at 50</option>
+            <option value="linear-truncated-200">Linear: Truncated at 200</option>
+            <option value="linear-truncated-1000">Linear: Truncated at 1000</option>
           </Select>
         </label>
       </ControlSection>
 
       <h4 style={{ marginBottom: '0.66em' }}>Genotype Distribution</h4>
       <ShortTandemRepeatGenotypeDistributionPlot
-        // @ts-expect-error TS(2322) FIXME: Type '{ axisLabels: any; maxRepeats: (string | und... Remove this comment to see the full error message
         axisLabels={getGenotypeDistributionPlotAxisLabels(adjacentRepeat, {
           selectedRepeatUnits: selectedGenotypeDistributionRepeatUnits,
         })}
-        maxRepeats={[
-          max(adjacentRepeat.genotype_distribution.distribution, (d: any) => max(d.slice(0, 2))),
-          max(adjacentRepeat.genotype_distribution.distribution, (d: any) => min(d.slice(0, 2))),
-        ]}
+        maxRepeats={maxGenotypeDistributionRepeats(adjacentRepeat)}
         genotypeDistribution={getSelectedGenotypeDistribution(adjacentRepeat, {
           selectedRepeatUnits: selectedGenotypeDistributionRepeatUnits,
-          selectedPopulationId,
+          selectedPopulation,
+          selectedSex,
         })}
-        onSelectBin={(bin: any) => {
+        onSelectBin={(bin: GenotypeBin) => {
           if (bin.xRange[0] !== bin.xRange[1] || bin.yRange[0] !== bin.yRange[1]) {
             setSelectedGenotypeDistributionBin(bin)
           }
         }}
+        xRanges={[]}
+        yRanges={[]}
       />
 
       <ControlSection>
         <ShortTandemRepeatPopulationOptions
           id={`${adjacentRepeat.id}-genotype-distribution`}
-          populationIds={populationIds}
-          selectedPopulationId={selectedPopulationId}
-          onSelectPopulationId={onSelectPopulationId}
+          populations={populations}
+          selectedPopulation={selectedPopulation}
+          selectedSex={selectedSex}
+          setSelectedPopulation={setSelectedPopulation}
+          setSelectedSex={setSelectedSex}
         />
 
         <ShortTandemRepeatGenotypeDistributionRepeatUnitsSelect
           shortTandemRepeatOrAdjacentRepeat={adjacentRepeat}
-          value={selectedGenotypeDistributionRepeatUnits}
-          onChange={setSelectedGenotypeDistributionRepeatUnits}
+          selectedRepeatUnits={selectedGenotypeDistributionRepeatUnits}
+          setSelectedRepeatUnits={setSelectedGenotypeDistributionRepeatUnits}
         />
       </ControlSection>
 
       {selectedGenotypeDistributionBin && (
         <Modal
-          title={(selectedGenotypeDistributionBin as any).label}
+          title={selectedGenotypeDistributionBin.label}
           size="large"
           // @ts-expect-error TS(2322) FIXME: Type '{ children: Element; title: any; size: "larg... Remove this comment to see the full error message
           initialFocusOnButton={false}
@@ -160,9 +181,11 @@ const ShortTandemRepeatAdjacentRepeatSection = ({
         >
           <ShortTandemRepeatGenotypeDistributionBinDetails
             shortTandemRepeatOrAdjacentRepeat={adjacentRepeat}
-            selectedPopulationId={selectedPopulationId}
             selectedRepeatUnits={selectedGenotypeDistributionRepeatUnits}
             bin={selectedGenotypeDistributionBin}
+            selectedPopulation={selectedPopulation}
+            selectedSex={selectedSex}
+            repeatUnitPairs={genotypeDistributionPairs}
           />
         </Modal>
       )}