feat: drive Chart 2's box plot from real posterior draws, with fallback

2026-08-11 09:07:53 -04:00
parent 91e21feb74
commit d5bf18482a
3 changed files with 136 additions and 17 deletions
@@ -2,33 +2,40 @@ import ChartLegend from '../components/ChartLegend.jsx'
 import ApproxNote from '../components/ApproxNote.jsx'
 import { raceColor, OBSERVED_MARK_COLOR, SHORT_RACE_LABEL } from '../utils/colors.js'
 import { fitSkewedInterval } from '../utils/distributionApprox.js'
+import { quantile } from '../utils/kde.js'
 import { niceTicks } from '../utils/niceTicks.js'
+import { useDrawDistribution } from '../hooks/useDrawDistribution.js'

 /**
 * Arrest rate by student group, most recent year, disaggregated into two
 * panels (Female / Male), each a horizontal box-and-whisker across the 4
- * race categories. Whisker = the model's reported 90% interval, box = the
- * fitted approximation's 25th-75th percentile, white tick = median, dark
- * diamond = observed rate.
+ * race categories. Whisker = the model's reported 90% interval; box = the
+ * 25th-75th percentile of the group's 500 real posterior draws (or, if draws
+ * are unavailable, the fitSkewedInterval approximation); white tick =
+ * median; dark diamond = observed rate.
 */

 const RACE_ORDER = ['WH', 'BL', 'HI', 'AM']
 const SEX_PANELS = [{ sex: 'F', label: 'Female' }, { sex: 'M', label: 'Male' }]
 const ROW_HEIGHT = 34
+const MODEL = 'unified_m3_mod' // matches ChartPanel's WAVE_MODEL, which built `data`

-function buildBox(d) {
+function buildBox(d, draws) {
  const median = Math.max(d.modeledMedian || 0, 0)
  const lower = Math.max(Math.min(d.rateLower ?? median, median), 0)
  const upper = Math.max(d.rateUpper ?? median, median)
-  const fit = fitSkewedInterval({ median, lower, upper })
-  return {
-    lower, upper, median,
-    q1: Math.max(fit.quantile(0.25), 0),
-    q3: Math.max(fit.quantile(0.75), median),
+
+  if (draws && draws.length > 0) {
+    return { lower, upper, median, q1: Math.max(quantile(draws, 0.25), 0), q3: Math.max(quantile(draws, 0.75), median) }
  }
+  const fit = fitSkewedInterval({ median, lower, upper })
+  return { lower, upper, median, q1: Math.max(fit.quantile(0.25), 0), q3: Math.max(fit.quantile(0.75), median) }
 }

-export default function RateByGroupBar({ data }) {
+export default function RateByGroupBar({ data, leaid, state }) {
+  const groups = data.map((d) => ({ race: d.race, sex: d.sex, stuEnroll: d.enrollment }))
+  const { status, drawsByGroup } = useDrawDistribution({ leaid, state, model: MODEL, year: '21-22', groups })
+
  const maxRate = Math.max(
    ...data.map((d) => Math.max(d.observedRate, d.rateUpper ?? d.modeledMedian ?? 0)),
    0.5
@@ -40,11 +47,18 @@ export default function RateByGroupBar({ data }) {
      <h3 style={{ fontSize: '0.85rem', marginBottom: 'var(--space-1)', color: 'var(--cv-ink-2)' }}>
        Arrest rate by student group — 2021–22 (per 1,000)
      </h3>
-      <ApproxNote />
+      {status !== 'ready' && <ApproxNote />}

      <div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr', gap: 'var(--space-3)', marginTop: 'var(--space-2)' }}>
        {SEX_PANELS.map(({ sex, label }) => (
-          <SexPanel key={sex} label={label} rows={data.filter((d) => d.sex === sex)} ticks={ticks} niceMax={niceMax} />
+          <SexPanel
+            key={sex}
+            label={label}
+            rows={data.filter((d) => d.sex === sex)}
+            drawsByGroup={drawsByGroup}
+            ticks={ticks}
+            niceMax={niceMax}
+          />
        ))}
      </div>

@@ -56,14 +70,16 @@ export default function RateByGroupBar({ data }) {
      ]} />

      <p style={{ fontSize: '0.72rem', color: 'var(--cv-ink-3)', marginTop: 'var(--space-1)' }}>
-        Box = modeled 25th–75th percentile (fitted approximation); whisker = the model's reported
-        90% interval; white tick = median. The dark diamond is the observed rate.
+        {status === 'ready'
+          ? "Box = 25th–75th percentile of 500 real posterior draws; whisker = the model's reported 90% interval; white tick = median."
+          : "Box = modeled 25th–75th percentile (fitted approximation); whisker = the model's reported 90% interval; white tick = median."}
+        {' '}The dark diamond is the observed rate.
      </p>
    </div>
  )
 }

-function SexPanel({ label, rows, ticks, niceMax }) {
+function SexPanel({ label, rows, drawsByGroup, ticks, niceMax }) {
  const width = 300
  const margin = { top: 30, right: 16, bottom: 34, left: 66 }
  const innerWidth = width - margin.left - margin.right
@@ -99,7 +115,8 @@ function SexPanel({ label, rows, ticks, niceMax }) {
          const boxTop = midY - ROW_HEIGHT * 0.26
          const boxBottom = midY + ROW_HEIGHT * 0.26
          const color = raceColor(d.race)
-          const box = buildBox(d)
+          const draws = drawsByGroup?.[`${d.race}_${d.sex}`]
+          const box = buildBox(d, draws)
          const observedX = xScale(d.observedRate)

          return (
@@ -95,7 +95,7 @@ export default function ChartPanel({ district, state }) {
      }}>
        <ArrestsOverTime data={timeSeriesData} modelId={WAVE_MODEL} />

-        <RateByGroupBar data={rateByGroup} />
+        <RateByGroupBar data={rateByGroup} leaid={district.leaid} state={state} />

        <RateDensityRidgeline quadData={data.quadData} />
      </div>
@@ -0,0 +1,102 @@
+import { useEffect, useState } from 'react'
+import { getDb } from '../utils/duckdbClient.js'
+
+const HF_BASE = 'https://huggingface.co/datasets/civilytics/crdc-school-arrest-rates/resolve/main/parquet'
+
+// Module-level cache: one registered duckdb-wasm file buffer per
+// (model, year, state) shard, shared across every component instance and
+// district navigated to in this browser session. See Global Constraints —
+// in-memory only, no persistence across page loads.
+const shardCache = new Map()
+
+function shardKey(model, year, state) {
+  return `${model}__${year}__${state}`
+}
+
+function ensureShardRegistered(db, model, year, state) {
+  const key = shardKey(model, year, state)
+  if (!shardCache.has(key)) {
+    shardCache.set(
+      key,
+      (async () => {
+        const url = `${HF_BASE}/model_id=${model}/YEAR=${year}/LEA_STATE=${state}/data_0.parquet`
+        const res = await fetch(url)
+        if (!res.ok) throw new Error(`Failed to fetch draw shard: HTTP ${res.status}`)
+        const buffer = new Uint8Array(await res.arrayBuffer())
+        const fileName = `${key}.parquet`
+        await db.registerFileBuffer(fileName, buffer)
+        return fileName
+      })(),
+    )
+  }
+  return shardCache.get(key)
+}
+
+/**
+ * Fetches real posterior draws for one district/model/year from the Hugging
+ * Face parquet dataset via duckdb-wasm, converts predicted counts to
+ * rate-per-1,000 using each group's stu_enroll (not present in the draws
+ * table itself — joined here from data this app already has), and returns
+ * them keyed by "RACE_SEX".
+ *
+ * @param {{leaid: string, state: string, model: string, year: string,
+ *   groups: Array<{race: string, sex: string, stuEnroll: number}>}} params
+ * @returns {{status: 'loading'|'ready'|'error', drawsByGroup: Record<string, number[]> | null}}
+ */
+export function useDrawDistribution({ leaid, state, model, year, groups }) {
+  const [status, setStatus] = useState('loading')
+  const [drawsByGroup, setDrawsByGroup] = useState(null)
+
+  // groups is typically a fresh array literal every render; derive a stable
+  // primitive so the effect only re-runs when its actual content changes.
+  const groupsSignature = (groups || []).map((g) => `${g.race}:${g.sex}:${g.stuEnroll}`).join(',')
+
+  useEffect(() => {
+    if (!leaid || !state || !model || !year || !groups?.length) return
+    let cancelled = false
+    setStatus('loading')
+
+    async function run() {
+      let conn
+      try {
+        const db = await getDb()
+        const fileName = await ensureShardRegistered(db, model, year, state)
+        conn = await db.connect()
+        const stmt = await conn.prepare(`SELECT RACE, SEX, pred FROM read_parquet('${fileName}') WHERE LEAID = ?`)
+        const table = await stmt.query(leaid)
+        await stmt.close()
+        const rows = table.toArray().map((r) => r.toJSON())
+
+        const enrollByGroup = {}
+        for (const g of groups) enrollByGroup[`${g.race}_${g.sex}`] = g.stuEnroll || 0
+
+        const byGroup = {}
+        for (const row of rows) {
+          const key = `${row.RACE}_${row.SEX}`
+          const enroll = enrollByGroup[key]
+          if (!enroll) continue
+          const rate = (Number(row.pred) / enroll) * 1000
+          ;(byGroup[key] ??= []).push(rate)
+        }
+
+        if (!cancelled) {
+          setDrawsByGroup(byGroup)
+          setStatus('ready')
+        }
+      } catch (err) {
+        console.error('useDrawDistribution failed:', err)
+        if (!cancelled) setStatus('error')
+      } finally {
+        if (conn) await conn.close()
+      }
+    }
+
+    run()
+    return () => {
+      cancelled = true
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [leaid, state, model, year, groupsSignature])
+
+  return { status, drawsByGroup }
+}