diff --git a/.gitignore b/.gitignore index 8f851f71..9082e158 100644 --- a/.gitignore +++ b/.gitignore @@ -122,6 +122,8 @@ alluvial_diagram_with_session.png .venv_raw .venv_new .venv_yellow +.venv-guessing +outputs/ EXTERNAL/Screenshots/Unsorted/BapunOpenField_Pf2D Refined.pdf EXTERNAL/Screenshots/Unsorted/BapunOpenField_Pf2D_Exports/BapunOpenField_Pf2D_pf2D.png SCRATCH/EXTERNAL/TESTING/Logging/debug_com.PhoHale.Spike3D.pipeline.log diff --git a/apps/__init__.py b/apps/__init__.py new file mode 100644 index 00000000..ad5b2f1e --- /dev/null +++ b/apps/__init__.py @@ -0,0 +1 @@ +"""Apps namespace package.""" diff --git a/apps/posterior_guessing/__init__.py b/apps/posterior_guessing/__init__.py new file mode 100644 index 00000000..8e340b7d --- /dev/null +++ b/apps/posterior_guessing/__init__.py @@ -0,0 +1 @@ +"""Posterior guessing FastAPI application package.""" diff --git a/apps/posterior_guessing/server.py b/apps/posterior_guessing/server.py new file mode 100644 index 00000000..29e5d74c --- /dev/null +++ b/apps/posterior_guessing/server.py @@ -0,0 +1,215 @@ +"""FastAPI server for the 2D posterior-guessing webapp.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np +from fastapi import FastAPI, HTTPException +from fastapi.responses import FileResponse +from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel, Field + + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from pho.posterior_guessing.bundle_io import load_bundle +from pho.posterior_guessing.persistence import save_prediction +from pho.posterior_guessing.scoring import renormalize, score_prediction + + +APP_DIR = Path(__file__).resolve().parent +STATIC_DIR = APP_DIR / 'static' +DEFAULT_BUNDLE_DIR = REPO_ROOT / 'data' / 'posterior_guessing' +DEFAULT_OUTPUT_ROOT = REPO_ROOT / 'outputs' / 'posterior_guessing' + +app = FastAPI(title='2D Posterior Guessing', version='0.1.0') +app.mount('/static', StaticFiles(directory=str(STATIC_DIR)), name='static') + +_BUNDLE_CACHE: Dict[str, Dict[str, Any]] = {} + + +class GuessRequest(BaseModel): + user_weights: List[List[float]] = Field(..., description='2D painted weights (n_x, n_y)') + save: bool = True + + +def _bundle_dir() -> Path: + return Path(DEFAULT_BUNDLE_DIR) + + +def _list_bundle_paths() -> List[Path]: + d = _bundle_dir() + if not d.exists(): + return [] + return sorted(d.glob('*_pf2d_guessing.npz')) + + +def _get_bundle(bundle_id: str) -> Dict[str, Any]: + if bundle_id in _BUNDLE_CACHE: + return _BUNDLE_CACHE[bundle_id] + + # Prefer exact stem match under data/posterior_guessing + candidates = [] + for p in _list_bundle_paths(): + data = None + # Match by filename stem or metadata session_id without loading all eagerly if name matches + if p.stem == bundle_id or p.stem.replace('_pf2d_guessing', '') == bundle_id: + candidates.append(p) + ## END for p in _list_bundle_paths()... + + if not candidates: + # Fallback: load each to check metadata session_id + for p in _list_bundle_paths(): + loaded = load_bundle(p, validate=True) + if loaded['bundle_id'] == bundle_id or p.stem == bundle_id: + _BUNDLE_CACHE[bundle_id] = loaded + _BUNDLE_CACHE[loaded['bundle_id']] = loaded + return loaded + ## END for p in _list_bundle_paths()... + raise HTTPException(status_code=404, detail=f'Bundle not found: {bundle_id}') + + loaded = load_bundle(candidates[0], validate=True) + _BUNDLE_CACHE[bundle_id] = loaded + _BUNDLE_CACHE[loaded['bundle_id']] = loaded + _BUNDLE_CACHE[Path(loaded['path']).stem] = loaded + return loaded + + +@app.get('/') +def index() -> FileResponse: + return FileResponse(STATIC_DIR / 'index.html') + + +@app.get('/api/health') +def health() -> Dict[str, Any]: + return {'ok': True, 'bundle_dir': str(_bundle_dir())} + + +@app.get('/api/bundles') +def list_bundles() -> Dict[str, Any]: + items = [] + for p in _list_bundle_paths(): + try: + data = load_bundle(p, validate=True) + items.append({ + 'bundle_id': data['bundle_id'], + 'path': data['path'], + 'n_neurons': int(np.asarray(data['neuron_ids']).shape[0]), + 'n_x': int(data['tuning_curves'].shape[1]), + 'n_y': int(data['tuning_curves'].shape[2]), + 'n_time': int(data['spike_counts'].shape[1]), + 'time_bin_size': float(np.asarray(data['time_bin_size']).reshape(())), + 'metadata': data.get('metadata', {}), + }) + _BUNDLE_CACHE[data['bundle_id']] = data + except Exception as exc: + items.append({'path': str(p), 'error': str(exc)}) + ## END for p in _list_bundle_paths()... + return {'bundles': items} + + +@app.get('/api/bundles/{bundle_id}') +def bundle_summary(bundle_id: str) -> Dict[str, Any]: + data = _get_bundle(bundle_id) + n_x, n_y, n_time = data['p_x_given_n'].shape + return { + 'bundle_id': data['bundle_id'], + 'path': data['path'], + 'n_neurons': int(data['neuron_ids'].shape[0]), + 'n_x': n_x, + 'n_y': n_y, + 'n_time': n_time, + 'time_bin_size': float(np.asarray(data['time_bin_size']).reshape(())), + 'time_bin_centers': np.asarray(data['time_bin_centers'], dtype=float).tolist(), + 'xbin': np.asarray(data['xbin'], dtype=float).tolist(), + 'ybin': np.asarray(data['ybin'], dtype=float).tolist(), + 'neuron_ids': np.asarray(data['neuron_ids']).tolist(), + 'occupancy': np.asarray(data['occupancy'], dtype=float).tolist(), + 'metadata': data.get('metadata', {}), + } + + +@app.get('/api/bundles/{bundle_id}/bins/{bin_index}') +def get_bin(bundle_id: str, bin_index: int) -> Dict[str, Any]: + """Return bin payload WITHOUT the true posterior (hidden until reveal).""" + data = _get_bundle(bundle_id) + n_time = int(data['spike_counts'].shape[1]) + if bin_index < 0 or bin_index >= n_time: + raise HTTPException(status_code=404, detail=f'bin_index out of range 0..{n_time - 1}') + + counts = np.asarray(data['spike_counts'])[:, bin_index] + active_mask = counts > 0 + active_idxs = np.where(active_mask)[0] + neuron_ids = np.asarray(data['neuron_ids']) + tuning = np.asarray(data['tuning_curves'], dtype=float) + + active_cells = [] + for i in active_idxs.tolist(): + active_cells.append({ + 'neuron_id': neuron_ids[i].item() if hasattr(neuron_ids[i], 'item') else neuron_ids[i], + 'neuron_index': int(i), + 'spike_count': int(counts[i]), + 'tuning_curve': tuning[i].tolist(), + }) + ## END for i in active_idxs.tolist()... + + n_x, n_y = tuning.shape[1], tuning.shape[2] + return { + 'bundle_id': data['bundle_id'], + 'bin_index': int(bin_index), + 'time_bin_center': float(np.asarray(data['time_bin_centers'])[bin_index]), + 'time_bin_size': float(np.asarray(data['time_bin_size']).reshape(())), + 'n_x': n_x, + 'n_y': n_y, + 'xbin': np.asarray(data['xbin'], dtype=float).tolist(), + 'ybin': np.asarray(data['ybin'], dtype=float).tolist(), + 'total_spikes': int(np.sum(counts)), + 'active_cells': active_cells, + # Intentionally omit p_x_given_n + } + + +@app.post('/api/bundles/{bundle_id}/bins/{bin_index}/reveal') +def reveal_bin(bundle_id: str, bin_index: int, body: GuessRequest) -> Dict[str, Any]: + data = _get_bundle(bundle_id) + n_x, n_y, n_time = data['p_x_given_n'].shape + if bin_index < 0 or bin_index >= n_time: + raise HTTPException(status_code=404, detail=f'bin_index out of range 0..{n_time - 1}') + + user = np.asarray(body.user_weights, dtype=np.float64) + if user.shape != (n_x, n_y): + raise HTTPException(status_code=400, detail=f'user_weights must have shape ({n_x}, {n_y}); got {user.shape}') + if float(np.sum(np.clip(user, 0.0, None))) <= 0: + raise HTTPException(status_code=400, detail='Paint a non-zero posterior before revealing') + + truth = np.asarray(data['p_x_given_n'][:, :, bin_index], dtype=np.float64) + scores = score_prediction(user, truth) + user_p = renormalize(user) + truth_p = renormalize(truth) + + saved = None + if body.save: + saved = save_prediction( + bundle_id=data['bundle_id'], + time_bin_index=bin_index, + user_weights=user, + true_posterior=truth, + scores=scores, + root=DEFAULT_OUTPUT_ROOT, + extra={'time_bin_center': float(np.asarray(data['time_bin_centers'])[bin_index])}, + save_npz=True, + ) + + return { + 'bundle_id': data['bundle_id'], + 'bin_index': int(bin_index), + 'scores': scores, + 'user_posterior': user_p.tolist(), + 'true_posterior': truth_p.tolist(), + 'saved': saved, + } diff --git a/apps/posterior_guessing/static/app.js b/apps/posterior_guessing/static/app.js new file mode 100644 index 00000000..50e13173 --- /dev/null +++ b/apps/posterior_guessing/static/app.js @@ -0,0 +1,344 @@ +/* 2D posterior guessing frontend */ + +const state = { + bundles: [], + bundleId: null, + summary: null, + binIndex: 0, + nX: 0, + nY: 0, + weights: null, + revealed: false, + painting: false, +}; + +const els = { + bundleSelect: document.getElementById('bundleSelect'), + binSlider: document.getElementById('binSlider'), + prevBin: document.getElementById('prevBin'), + nextBin: document.getElementById('nextBin'), + binLabel: document.getElementById('binLabel'), + spikeLabel: document.getElementById('spikeLabel'), + paintCanvas: document.getElementById('paintCanvas'), + truthCanvas: document.getElementById('truthCanvas'), + clearPaint: document.getElementById('clearPaint'), + revealBtn: document.getElementById('revealBtn'), + brushSize: document.getElementById('brushSize'), + brushStrength: document.getElementById('brushStrength'), + tuningGrid: document.getElementById('tuningGrid'), + emptyCells: document.getElementById('emptyCells'), + scoreBox: document.getElementById('scoreBox'), + scorePrimary: document.getElementById('scorePrimary'), + scoreSecondary: document.getElementById('scoreSecondary'), +}; + +const paintCtx = els.paintCanvas.getContext('2d'); +const truthCtx = els.truthCanvas.getContext('2d'); + +function ylOrRd(t) { + // Approximate YlOrRd + const stops = [ + [1.0, 1.0, 0.8], + [0.996, 0.878, 0.545], + [0.992, 0.682, 0.38], + [0.89, 0.29, 0.2], + [0.55, 0.02, 0.15], + ]; + return sampleStops(stops, t); +} + +function blues(t) { + const stops = [ + [0.97, 0.98, 1.0], + [0.73, 0.85, 0.92], + [0.42, 0.68, 0.84], + [0.19, 0.45, 0.69], + [0.03, 0.19, 0.42], + ]; + return sampleStops(stops, t); +} + +function sampleStops(stops, t) { + const x = Math.min(1, Math.max(0, t)); + const scaled = x * (stops.length - 1); + const i = Math.floor(scaled); + const f = scaled - i; + const a = stops[i]; + const b = stops[Math.min(i + 1, stops.length - 1)]; + return [ + a[0] + (b[0] - a[0]) * f, + a[1] + (b[1] - a[1]) * f, + a[2] + (b[2] - a[2]) * f, + ]; +} + +function renormalize(weights) { + let sum = 0; + for (let i = 0; i < weights.length; i++) { + for (let j = 0; j < weights[0].length; j++) { + if (weights[i][j] < 0) weights[i][j] = 0; + sum += weights[i][j]; + } + } + if (sum <= 0) return null; + const out = weights.map((row) => row.map((v) => v / sum)); + return out; +} + +function zeros2d(nX, nY) { + return Array.from({ length: nX }, () => Array(nY).fill(0)); +} + +function paintMass(weights) { + let sum = 0; + for (let i = 0; i < weights.length; i++) { + for (let j = 0; j < weights[0].length; j++) sum += Math.max(0, weights[i][j]); + } + return sum; +} + +function drawMap(ctx, canvas, map, cmap, { blankIfNull = true } = {}) { + const nX = state.nX; + const nY = state.nY; + if (!nX || !nY) return; + canvas.width = nX; + canvas.height = nY; + const img = ctx.createImageData(nX, nY); + if (!map) { + if (blankIfNull) { + for (let k = 0; k < img.data.length; k += 4) { + img.data[k] = 20; + img.data[k + 1] = 24; + img.data[k + 2] = 28; + img.data[k + 3] = 255; + } + ctx.putImageData(img, 0, 0); + } + return; + } + let maxV = 0; + for (let i = 0; i < nX; i++) { + for (let j = 0; j < nY; j++) maxV = Math.max(maxV, map[i][j]); + } + if (maxV <= 0) maxV = 1; + // ImageData is row-major with y downward; our arrays are [x][y] + for (let iy = 0; iy < nY; iy++) { + for (let ix = 0; ix < nX; ix++) { + const t = map[ix][iy] / maxV; + const [r, g, b] = cmap(t); + // Flip y for display so low y is at bottom + const dy = nY - 1 - iy; + const k = (dy * nX + ix) * 4; + img.data[k] = Math.round(r * 255); + img.data[k + 1] = Math.round(g * 255); + img.data[k + 2] = Math.round(b * 255); + img.data[k + 3] = map[ix][iy] > 0 ? 230 : 40; + } + } + ctx.putImageData(img, 0, 0); +} + +function redrawPaint() { + const norm = renormalize(state.weights.map((r) => r.slice())); + drawMap(paintCtx, els.paintCanvas, norm || null, ylOrRd); + els.revealBtn.disabled = paintMass(state.weights) <= 0 || state.revealed; +} + +function eventToBin(evt) { + const rect = els.paintCanvas.getBoundingClientRect(); + const px = (evt.clientX - rect.left) / rect.width; + const py = (evt.clientY - rect.top) / rect.height; + const ix = Math.min(state.nX - 1, Math.max(0, Math.floor(px * state.nX))); + // Flip y back from display + const iyDisp = Math.min(state.nY - 1, Math.max(0, Math.floor(py * state.nY))); + const iy = state.nY - 1 - iyDisp; + return { ix, iy }; +} + +function stampBrush(ix, iy) { + if (state.revealed) return; + const radius = Number(els.brushSize.value); + const strength = Number(els.brushStrength.value) / 10; + for (let dx = -radius; dx <= radius; dx++) { + for (let dy = -radius; dy <= radius; dy++) { + const x = ix + dx; + const y = iy + dy; + if (x < 0 || y < 0 || x >= state.nX || y >= state.nY) continue; + const dist = Math.sqrt(dx * dx + dy * dy); + if (dist > radius) continue; + const w = strength * Math.exp(-0.5 * (dist / Math.max(0.5, radius * 0.55)) ** 2); + state.weights[x][y] += w; + } + } + redrawPaint(); +} + +async function api(path, opts) { + const res = await fetch(path, opts); + if (!res.ok) { + let detail = res.statusText; + try { + const body = await res.json(); + detail = body.detail || JSON.stringify(body); + } catch (_) {} + throw new Error(detail); + } + return res.json(); +} + +function renderTuning(activeCells) { + els.tuningGrid.innerHTML = ''; + if (!activeCells.length) { + els.emptyCells.hidden = false; + return; + } + els.emptyCells.hidden = true; + activeCells.forEach((cell, idx) => { + const card = document.createElement('div'); + card.className = 'cell-card'; + const label = document.createElement('div'); + label.className = 'label'; + label.innerHTML = `Unit ${cell.neuron_id} · spikes ${cell.spike_count}`; + const plot = document.createElement('div'); + plot.id = `tune-${idx}`; + plot.style.height = '110px'; + card.appendChild(label); + card.appendChild(plot); + els.tuningGrid.appendChild(card); + + const z = cell.tuning_curve; + // Plotly heatmap expects z as [y][x]; our tuning is [x][y] + const zT = []; + for (let y = 0; y < state.nY; y++) { + const row = []; + for (let x = 0; x < state.nX; x++) row.push(z[x][y]); + zT.push(row); + } + Plotly.newPlot( + plot, + [{ + z: zT, + type: 'heatmap', + colorscale: 'Viridis', + showscale: false, + }], + { + margin: { l: 0, r: 0, t: 0, b: 0 }, + paper_bgcolor: 'rgba(0,0,0,0)', + plot_bgcolor: 'rgba(0,0,0,0)', + xaxis: { visible: false }, + yaxis: { visible: false }, + }, + { displayModeBar: false, staticPlot: true, responsive: true }, + ); + }); +} + +async function loadBin(binIndex) { + state.binIndex = binIndex; + state.revealed = false; + els.truthCanvas.hidden = true; + els.scoreBox.hidden = true; + const data = await api(`/api/bundles/${encodeURIComponent(state.bundleId)}/bins/${binIndex}`); + state.nX = data.n_x; + state.nY = data.n_y; + state.weights = zeros2d(state.nX, state.nY); + els.binSlider.value = String(binIndex); + els.binLabel.textContent = `Bin ${binIndex} / ${state.summary.n_time - 1} · t=${data.time_bin_center.toFixed(3)}s`; + els.spikeLabel.textContent = `Spikes ${data.total_spikes} · ${data.active_cells.length} cells`; + redrawPaint(); + renderTuning(data.active_cells); +} + +async function selectBundle(bundleId) { + state.bundleId = bundleId; + state.summary = await api(`/api/bundles/${encodeURIComponent(bundleId)}`); + els.binSlider.min = '0'; + els.binSlider.max = String(Math.max(0, state.summary.n_time - 1)); + await loadBin(0); +} + +async function reveal() { + if (state.revealed || paintMass(state.weights) <= 0) return; + const body = { + user_weights: state.weights, + save: true, + }; + const result = await api(`/api/bundles/${encodeURIComponent(state.bundleId)}/bins/${state.binIndex}/reveal`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); + state.revealed = true; + els.revealBtn.disabled = true; + drawMap(paintCtx, els.paintCanvas, result.user_posterior, ylOrRd); + drawMap(truthCtx, els.truthCanvas, result.true_posterior, blues); + els.truthCanvas.hidden = false; + els.scoreBox.hidden = false; + els.scorePrimary.textContent = `Hellinger affinity ${result.scores.hellinger_affinity.toFixed(3)}`; + els.scoreSecondary.textContent = `Cosine ${result.scores.cosine_similarity.toFixed(3)}`; +} + +async function init() { + const listing = await api('/api/bundles'); + state.bundles = listing.bundles.filter((b) => !b.error); + els.bundleSelect.innerHTML = ''; + if (!state.bundles.length) { + els.bundleSelect.innerHTML = ''; + els.binLabel.textContent = 'Run: python scripts/export_posterior_guessing_bundle.py --synthetic'; + return; + } + state.bundles.forEach((b) => { + const opt = document.createElement('option'); + opt.value = b.bundle_id; + opt.textContent = `${b.bundle_id} (${b.n_time} bins)`; + els.bundleSelect.appendChild(opt); + }); + els.bundleSelect.addEventListener('change', () => selectBundle(els.bundleSelect.value)); + els.binSlider.addEventListener('input', () => loadBin(Number(els.binSlider.value))); + els.prevBin.addEventListener('click', () => { + const next = Math.max(0, state.binIndex - 1); + loadBin(next); + }); + els.nextBin.addEventListener('click', () => { + const next = Math.min(state.summary.n_time - 1, state.binIndex + 1); + loadBin(next); + }); + els.clearPaint.addEventListener('click', () => { + if (state.revealed) return; + state.weights = zeros2d(state.nX, state.nY); + redrawPaint(); + }); + els.revealBtn.addEventListener('click', () => reveal().catch((e) => alert(e.message))); + + els.paintCanvas.addEventListener('mousedown', (evt) => { + state.painting = true; + const { ix, iy } = eventToBin(evt); + stampBrush(ix, iy); + }); + els.paintCanvas.addEventListener('mousemove', (evt) => { + if (!state.painting) return; + const { ix, iy } = eventToBin(evt); + stampBrush(ix, iy); + }); + window.addEventListener('mouseup', () => { state.painting = false; }); + els.paintCanvas.addEventListener('mouseleave', () => { state.painting = false; }); + + window.addEventListener('keydown', (evt) => { + if (evt.key === 'Enter') { + evt.preventDefault(); + reveal().catch((e) => alert(e.message)); + } else if (evt.key === 'ArrowLeft') { + els.prevBin.click(); + } else if (evt.key === 'ArrowRight') { + els.nextBin.click(); + } + }); + + await selectBundle(state.bundles[0].bundle_id); +} + +init().catch((err) => { + console.error(err); + els.binLabel.textContent = `Failed to start: ${err.message}`; +}); diff --git a/apps/posterior_guessing/static/index.html b/apps/posterior_guessing/static/index.html new file mode 100644 index 00000000..577f23a2 --- /dev/null +++ b/apps/posterior_guessing/static/index.html @@ -0,0 +1,64 @@ + + + + + + 2D Posterior Guessing + + + + +
+
Posterior Guessing
+ + +
+ +
+
+
+ + + +
+
+ Bin — + Spikes — + Paint on the blank canvas, then press Enter to reveal +
+
+ +
+
Your posterior estimate
+
+ + +
+
+ + + + +
+
+ User (YlOrRd) + True (Blues) — shown after Enter +
+
+ +
+
Cells that fired in this bin
+
+ +
+
+ + + + diff --git a/apps/posterior_guessing/static/styles.css b/apps/posterior_guessing/static/styles.css new file mode 100644 index 00000000..d5819f11 --- /dev/null +++ b/apps/posterior_guessing/static/styles.css @@ -0,0 +1,242 @@ +:root { + --bg: #1a1f24; + --panel: #242b33; + --text: #e8eef4; + --muted: #9aa7b5; + --accent: #e2a35a; + --line: #3a4552; + --user: #f0a060; + --truth: #6aa6d8; +} + +* { box-sizing: border-box; } + +html, body { + margin: 0; + min-height: 100%; + background: + radial-gradient(1200px 600px at 10% -10%, #2c3642 0%, transparent 55%), + radial-gradient(900px 500px at 100% 0%, #3a2f28 0%, transparent 50%), + var(--bg); + color: var(--text); + font-family: "IBM Plex Sans", "Segoe UI", sans-serif; +} + +.topbar { + display: flex; + align-items: center; + gap: 1.25rem; + padding: 0.9rem 1.25rem; + border-bottom: 1px solid var(--line); + background: rgba(20, 24, 28, 0.72); + backdrop-filter: blur(8px); + position: sticky; + top: 0; + z-index: 10; +} + +.brand { + font-family: "IBM Plex Serif", Georgia, serif; + font-size: 1.35rem; + letter-spacing: 0.02em; +} + +.bundle-picker { + display: flex; + align-items: center; + gap: 0.5rem; + color: var(--muted); + font-size: 0.9rem; +} + +select, button, input[type="range"] { + accent-color: var(--accent); +} + +select, button { + background: #313a44; + color: var(--text); + border: 1px solid var(--line); + border-radius: 6px; + padding: 0.35rem 0.65rem; +} + +button:disabled { + opacity: 0.45; + cursor: not-allowed; +} + +button:not(:disabled):hover { + border-color: var(--accent); +} + +.score-box { + margin-left: auto; + display: flex; + gap: 1rem; + font-variant-numeric: tabular-nums; +} + +.score-box span { + background: #2d3844; + border: 1px solid var(--line); + border-radius: 6px; + padding: 0.35rem 0.7rem; +} + +.layout { + display: grid; + grid-template-columns: minmax(320px, 1.1fr) minmax(280px, 0.9fr); + grid-template-areas: + "scrub scrub" + "paint cells"; + gap: 1rem; + padding: 1rem 1.25rem 2rem; +} + +.scrub-panel { grid-area: scrub; } +.paint-panel { grid-area: paint; } +.cells-panel { grid-area: cells; } + +.panel { + background: rgba(36, 43, 51, 0.92); + border: 1px solid var(--line); + border-radius: 10px; + padding: 0.9rem 1rem 1rem; +} + +.panel-title { + font-size: 0.95rem; + margin-bottom: 0.65rem; + color: var(--muted); + text-transform: uppercase; + letter-spacing: 0.06em; +} + +.scrub-row { + display: flex; + align-items: center; + gap: 0.75rem; +} + +.scrub-row input[type="range"] { + flex: 1; +} + +.meta-row { + display: flex; + flex-wrap: wrap; + gap: 1rem; + margin-top: 0.55rem; + color: var(--muted); + font-size: 0.92rem; +} + +.hint kbd { + background: #3a4552; + border-radius: 4px; + padding: 0.05rem 0.35rem; + color: var(--text); +} + +.canvas-wrap { + position: relative; + width: 100%; + max-width: 560px; + aspect-ratio: 5 / 4; + background: + linear-gradient(45deg, #1e242b 25%, transparent 25%), + linear-gradient(-45deg, #1e242b 25%, transparent 25%), + linear-gradient(45deg, transparent 75%, #1e242b 75%), + linear-gradient(-45deg, transparent 75%, #1e242b 75%); + background-size: 24px 24px; + background-position: 0 0, 0 12px, 12px -12px, -12px 0; + background-color: #151a1f; + border: 1px solid var(--line); + border-radius: 8px; + overflow: hidden; + cursor: crosshair; +} + +#paintCanvas, #truthCanvas { + position: absolute; + inset: 0; + width: 100%; + height: 100%; + image-rendering: pixelated; +} + +#truthCanvas.overlay { + opacity: 0.72; + mix-blend-mode: screen; + pointer-events: none; +} + +.paint-controls { + display: flex; + flex-wrap: wrap; + gap: 0.75rem 1rem; + align-items: center; + margin-top: 0.75rem; + color: var(--muted); + font-size: 0.9rem; +} + +.legend { + margin-top: 0.55rem; + color: var(--muted); + font-size: 0.85rem; + display: flex; + gap: 0.75rem; + align-items: center; +} + +.swatch { + width: 0.85rem; + height: 0.85rem; + border-radius: 2px; + display: inline-block; +} + +.swatch.user { background: linear-gradient(90deg, #2b1a0f, #f0a060); } +.swatch.truth { background: linear-gradient(90deg, #0f1a24, #6aa6d8); } + +.tuning-grid { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(140px, 1fr)); + gap: 0.65rem; + max-height: 70vh; + overflow: auto; +} + +.cell-card { + background: #1b2229; + border: 1px solid var(--line); + border-radius: 8px; + padding: 0.4rem; +} + +.cell-card .label { + font-size: 0.8rem; + color: var(--muted); + margin-bottom: 0.25rem; +} + +.cell-card .label strong { + color: var(--text); +} + +.empty { + color: var(--muted); + padding: 1rem 0; +} + +@media (max-width: 900px) { + .layout { + grid-template-columns: 1fr; + grid-template-areas: + "scrub" + "paint" + "cells"; + } +} diff --git a/data/posterior_guessing/synthetic_demo_pf2d_guessing.npz b/data/posterior_guessing/synthetic_demo_pf2d_guessing.npz new file mode 100644 index 00000000..0544fe54 Binary files /dev/null and b/data/posterior_guessing/synthetic_demo_pf2d_guessing.npz differ diff --git a/pho/__init__.py b/pho/__init__.py new file mode 100644 index 00000000..5752231f --- /dev/null +++ b/pho/__init__.py @@ -0,0 +1 @@ +"""Spike3D local helpers package.""" diff --git a/pho/posterior_guessing/__init__.py b/pho/posterior_guessing/__init__.py new file mode 100644 index 00000000..e307697d --- /dev/null +++ b/pho/posterior_guessing/__init__.py @@ -0,0 +1,23 @@ +"""Portable 2D placefield / decoder bundles and posterior-guessing helpers.""" + +from pho.posterior_guessing.bundle_io import BUNDLE_REQUIRED_KEYS, load_bundle, save_bundle, validate_bundle +from pho.posterior_guessing.export import export_from_arrays, export_from_decoder +from pho.posterior_guessing.persistence import save_prediction +from pho.posterior_guessing.scoring import cosine_similarity, hellinger_affinity, renormalize, score_prediction +from pho.posterior_guessing.synthetic import make_synthetic_bundle + + +__all__ = [ + 'BUNDLE_REQUIRED_KEYS', + 'cosine_similarity', + 'export_from_arrays', + 'export_from_decoder', + 'hellinger_affinity', + 'load_bundle', + 'make_synthetic_bundle', + 'renormalize', + 'save_bundle', + 'save_prediction', + 'score_prediction', + 'validate_bundle', +] diff --git a/pho/posterior_guessing/bundle_io.py b/pho/posterior_guessing/bundle_io.py new file mode 100644 index 00000000..d5b5feab --- /dev/null +++ b/pho/posterior_guessing/bundle_io.py @@ -0,0 +1,176 @@ +"""NPZ schema load/save/validate for 2D posterior-guessing bundles.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Mapping, MutableMapping, Optional, Union + +import numpy as np + + +BUNDLE_REQUIRED_KEYS = ( + 'neuron_ids', + 'tuning_curves', + 'occupancy', + 'P_x', + 'spike_counts', + 'p_x_given_n', + 'time_bin_centers', + 'xbin', + 'ybin', + 'time_bin_size', + 'metadata_json', +) + + +PathLike = Union[str, Path] + + +def _as_path(path: PathLike) -> Path: + return Path(path).expanduser().resolve() + + +def validate_bundle(data: Mapping[str, Any], *, require_positive_mass: bool = True) -> Dict[str, Any]: + """Validate array shapes/dtypes for a guessing bundle; return normalized metadata dict. + + Raises + ------ + KeyError + Missing required keys. + ValueError + Inconsistent shapes or invalid probability mass. + """ + missing = [k for k in BUNDLE_REQUIRED_KEYS if k not in data] + if missing: + raise KeyError(f'Missing required bundle keys: {missing}') + + neuron_ids = np.asarray(data['neuron_ids']) + tuning_curves = np.asarray(data['tuning_curves'], dtype=np.float64) + occupancy = np.asarray(data['occupancy'], dtype=np.float64) + P_x = np.asarray(data['P_x'], dtype=np.float64) + spike_counts = np.asarray(data['spike_counts']) + p_x_given_n = np.asarray(data['p_x_given_n'], dtype=np.float64) + time_bin_centers = np.asarray(data['time_bin_centers'], dtype=np.float64) + xbin = np.asarray(data['xbin'], dtype=np.float64) + ybin = np.asarray(data['ybin'], dtype=np.float64) + time_bin_size = float(np.asarray(data['time_bin_size']).reshape(())) + + if tuning_curves.ndim != 3: + raise ValueError(f'tuning_curves must be (n_neurons, n_x, n_y); got shape {tuning_curves.shape}') + n_neurons, n_x, n_y = tuning_curves.shape + if neuron_ids.shape != (n_neurons,): + raise ValueError(f'neuron_ids shape {neuron_ids.shape} != ({n_neurons},)') + if occupancy.shape != (n_x, n_y): + raise ValueError(f'occupancy shape {occupancy.shape} != ({n_x}, {n_y})') + if P_x.shape != (n_x, n_y): + raise ValueError(f'P_x shape {P_x.shape} != ({n_x}, {n_y})') + if spike_counts.ndim != 2 or spike_counts.shape[0] != n_neurons: + raise ValueError(f'spike_counts must be (n_neurons, n_time); got {spike_counts.shape}') + n_time = spike_counts.shape[1] + if p_x_given_n.shape != (n_x, n_y, n_time): + raise ValueError(f'p_x_given_n shape {p_x_given_n.shape} != ({n_x}, {n_y}, {n_time})') + if time_bin_centers.shape != (n_time,): + raise ValueError(f'time_bin_centers shape {time_bin_centers.shape} != ({n_time},)') + if xbin.ndim != 1 or xbin.size != n_x + 1: + raise ValueError(f'xbin must be length n_x+1={n_x + 1}; got shape {xbin.shape}') + if ybin.ndim != 1 or ybin.size != n_y + 1: + raise ValueError(f'ybin must be length n_y+1={n_y + 1}; got shape {ybin.shape}') + if time_bin_size <= 0: + raise ValueError(f'time_bin_size must be > 0; got {time_bin_size}') + + if require_positive_mass: + prior_sum = float(np.sum(P_x)) + if not np.isfinite(prior_sum) or prior_sum <= 0: + raise ValueError('P_x must have positive finite mass') + # Check a sample of posterior columns for finite mass + col_sums = np.sum(p_x_given_n.reshape(n_x * n_y, n_time), axis=0) + if not np.all(np.isfinite(col_sums)): + raise ValueError('p_x_given_n contains non-finite values') + if np.any(col_sums <= 0): + raise ValueError('each p_x_given_n[..., t] must have positive mass') + + metadata = parse_metadata(data['metadata_json']) + return metadata + + +def parse_metadata(metadata_json: Any) -> Dict[str, Any]: + """Parse metadata_json which may be str, bytes, or already a dict.""" + if isinstance(metadata_json, dict): + return dict(metadata_json) + if isinstance(metadata_json, (bytes, bytearray, np.bytes_)): + metadata_json = bytes(metadata_json).decode('utf-8') + if isinstance(metadata_json, np.ndarray): + if metadata_json.shape == (): + metadata_json = metadata_json.item() + else: + metadata_json = str(metadata_json) + if not isinstance(metadata_json, str): + metadata_json = str(metadata_json) + metadata_json = metadata_json.strip() + if not metadata_json: + return {} + return json.loads(metadata_json) + + +def metadata_to_json(metadata: Optional[Mapping[str, Any]] = None) -> str: + return json.dumps(dict(metadata or {}), sort_keys=True) + + +def save_bundle(path: PathLike, data: Mapping[str, Any], *, validate: bool = True) -> Path: + """Write a compressed NPZ bundle. Returns the resolved output path.""" + out_path = _as_path(path) + out_path.parent.mkdir(parents=True, exist_ok=True) + + payload: MutableMapping[str, Any] = {k: data[k] for k in BUNDLE_REQUIRED_KEYS if k in data} + # Allow callers to pass metadata as dict + if 'metadata_json' in payload and not isinstance(payload['metadata_json'], (str, bytes, np.ndarray)): + payload['metadata_json'] = metadata_to_json(payload['metadata_json']) + elif 'metadata' in data and 'metadata_json' not in payload: + payload['metadata_json'] = metadata_to_json(data['metadata']) + + if validate: + validate_bundle(payload) + + # Ensure metadata is stored as a unicode string for NPZ + payload['metadata_json'] = metadata_to_json(parse_metadata(payload['metadata_json'])) + np.savez_compressed(out_path, **payload) + return out_path + + +def load_bundle(path: PathLike, *, validate: bool = True) -> Dict[str, Any]: + """Load an NPZ guessing bundle into a plain dict of arrays (+ parsed metadata).""" + in_path = _as_path(path) + with np.load(in_path, allow_pickle=False) as npz: + data = {k: npz[k] for k in npz.files} + + metadata = parse_metadata(data.get('metadata_json', '{}')) + if validate: + validate_bundle(data) + data['metadata'] = metadata + data['metadata_json'] = metadata_to_json(metadata) + data['path'] = str(in_path) + data['bundle_id'] = metadata.get('session_id') or in_path.stem + return data + + +def filter_active_time_bins(data: Mapping[str, Any], *, min_total_spikes: int = 1) -> Dict[str, Any]: + """Return a shallow-copied bundle keeping only time bins with enough spikes.""" + spike_counts = np.asarray(data['spike_counts']) + keep = np.sum(spike_counts, axis=0) >= min_total_spikes + if not np.any(keep): + raise ValueError('No time bins meet the spike filter') + + out = dict(data) + out['spike_counts'] = spike_counts[:, keep] + out['p_x_given_n'] = np.asarray(data['p_x_given_n'])[:, :, keep] + out['time_bin_centers'] = np.asarray(data['time_bin_centers'])[keep] + metadata = parse_metadata(data.get('metadata_json', data.get('metadata', {}))) + metadata = dict(metadata) + metadata['filtered_active_bins'] = True + metadata['min_total_spikes'] = int(min_total_spikes) + metadata['n_time_original'] = int(spike_counts.shape[1]) + metadata['n_time_kept'] = int(np.sum(keep)) + out['metadata'] = metadata + out['metadata_json'] = metadata_to_json(metadata) + return out diff --git a/pho/posterior_guessing/export.py b/pho/posterior_guessing/export.py new file mode 100644 index 00000000..85b167bf --- /dev/null +++ b/pho/posterior_guessing/export.py @@ -0,0 +1,123 @@ +"""Export helpers for 2D posterior-guessing NPZ bundles.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Union + +import numpy as np + +from pho.posterior_guessing.bundle_io import filter_active_time_bins, metadata_to_json, save_bundle + + +PathLike = Union[str, Path] + + +def export_from_arrays(*, neuron_ids, tuning_curves, occupancy, P_x, spike_counts, p_x_given_n, time_bin_centers, xbin, ybin, time_bin_size: float, out_path: PathLike, metadata: Optional[Mapping[str, Any]] = None, filter_empty_bins: bool = True, min_total_spikes: int = 1) -> Path: + """Build and save a guessing bundle from explicit arrays. Returns output path.""" + data: Dict[str, Any] = { + 'neuron_ids': np.asarray(neuron_ids), + 'tuning_curves': np.asarray(tuning_curves, dtype=np.float64), + 'occupancy': np.asarray(occupancy, dtype=np.float64), + 'P_x': np.asarray(P_x, dtype=np.float64), + 'spike_counts': np.asarray(spike_counts), + 'p_x_given_n': np.asarray(p_x_given_n, dtype=np.float64), + 'time_bin_centers': np.asarray(time_bin_centers, dtype=np.float64), + 'xbin': np.asarray(xbin, dtype=np.float64), + 'ybin': np.asarray(ybin, dtype=np.float64), + 'time_bin_size': np.asarray(float(time_bin_size), dtype=np.float64), + 'metadata_json': metadata_to_json(metadata), + } + if filter_empty_bins: + data = filter_active_time_bins(data, min_total_spikes=min_total_spikes) + return save_bundle(out_path, data, validate=True) + + +def _reshape_flat_spatial(flat: np.ndarray, n_x: int, n_y: int) -> np.ndarray: + """Reshape flat position axis of length n_x*n_y into (n_x, n_y, ...).""" + flat = np.asarray(flat) + n_flat = n_x * n_y + if flat.shape[0] != n_flat: + raise ValueError(f'Expected leading axis length {n_flat} (n_x*n_y); got {flat.shape}') + trailing = flat.shape[1:] + return flat.reshape((n_x, n_y) + trailing) + + +def export_from_decoder(decoder: Any, out_path: PathLike, *, metadata: Optional[Mapping[str, Any]] = None, filter_empty_bins: bool = True, min_total_spikes: int = 1, session_id: Optional[str] = None) -> Path: + """Export from a BayesianPlacemapPositionDecoder-like object when sibling packages are available. + + Expects 2D placefields: `decoder.pf.ratemap.tuning_curves` shaped (n_neurons, n_x, n_y) + and `decoder.p_x_given_n` shaped (n_x, n_y, n_time) or flat (n_x*n_y, n_time). + """ + pf = getattr(decoder, 'pf', None) + if pf is None: + raise ValueError('decoder has no .pf placefield attribute') + ratemap = getattr(pf, 'ratemap', None) or getattr(pf, '_ratemap', None) + if ratemap is None: + raise ValueError('decoder.pf has no ratemap') + + tuning_curves = np.asarray(ratemap.tuning_curves, dtype=np.float64) + if tuning_curves.ndim != 3: + raise ValueError(f'Expected 2D tuning_curves (n_neurons, n_x, n_y); got {tuning_curves.shape}') + n_neurons, n_x, n_y = tuning_curves.shape + + occupancy = np.asarray(ratemap.occupancy, dtype=np.float64) + xbin = np.asarray(getattr(ratemap, 'xbin', None) if getattr(ratemap, 'xbin', None) is not None else pf.xbin, dtype=np.float64) + ybin = np.asarray(getattr(ratemap, 'ybin', None) if getattr(ratemap, 'ybin', None) is not None else pf.ybin, dtype=np.float64) + if ybin is None or np.asarray(ybin).size == 0: + raise ValueError('2D export requires ybin edges') + + neuron_ids = np.asarray(getattr(decoder, 'neuron_IDs', getattr(ratemap, 'neuron_ids', np.arange(n_neurons)))) + if neuron_ids.shape != (n_neurons,): + neuron_ids = np.asarray(list(neuron_ids))[:n_neurons] + + P_x = np.asarray(decoder.P_x, dtype=np.float64) + if P_x.ndim == 2 and P_x.shape[1] == 1: + P_x = _reshape_flat_spatial(P_x[:, 0], n_x, n_y) + elif P_x.ndim == 1: + P_x = _reshape_flat_spatial(P_x, n_x, n_y) + elif P_x.shape != (n_x, n_y): + raise ValueError(f'Could not reshape P_x with shape {P_x.shape} to ({n_x}, {n_y})') + + spike_counts = np.asarray(decoder.unit_specific_time_binned_spike_counts) + p_x_given_n = np.asarray(decoder.p_x_given_n, dtype=np.float64) + if p_x_given_n.ndim == 2: + p_x_given_n = _reshape_flat_spatial(p_x_given_n, n_x, n_y) + elif p_x_given_n.shape[:2] != (n_x, n_y): + raise ValueError(f'p_x_given_n shape {p_x_given_n.shape} incompatible with ({n_x}, {n_y}, n_time)') + + # Align time axes if spike counts and posterior lengths differ (common in docs) + n_time_spikes = spike_counts.shape[1] + n_time_post = p_x_given_n.shape[2] + n_time = min(n_time_spikes, n_time_post) + spike_counts = spike_counts[:, :n_time] + p_x_given_n = p_x_given_n[:, :, :n_time] + + tbc = getattr(getattr(decoder, 'time_binning_container', None), 'centers', None) + if tbc is None: + time_bin_centers = np.arange(n_time, dtype=np.float64) * float(decoder.time_bin_size) + else: + time_bin_centers = np.asarray(tbc, dtype=np.float64)[:n_time] + + meta = dict(metadata or {}) + if session_id is not None: + meta['session_id'] = session_id + meta.setdefault('source', 'export_from_decoder') + meta.setdefault('ndim', 2) + + return export_from_arrays( + neuron_ids=neuron_ids, + tuning_curves=tuning_curves, + occupancy=occupancy, + P_x=P_x, + spike_counts=spike_counts, + p_x_given_n=p_x_given_n, + time_bin_centers=time_bin_centers, + xbin=xbin, + ybin=ybin, + time_bin_size=float(decoder.time_bin_size), + out_path=out_path, + metadata=meta, + filter_empty_bins=filter_empty_bins, + min_total_spikes=min_total_spikes, + ) diff --git a/pho/posterior_guessing/persistence.py b/pho/posterior_guessing/persistence.py new file mode 100644 index 00000000..c86a0bbf --- /dev/null +++ b/pho/posterior_guessing/persistence.py @@ -0,0 +1,72 @@ +"""Persist user posterior guesses for later analysis.""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Union + +import numpy as np + +from pho.posterior_guessing.scoring import renormalize, score_prediction + + +PathLike = Union[str, Path] + + +def default_predictions_dir(bundle_id: str, root: PathLike = 'outputs/posterior_guessing') -> Path: + return Path(root).expanduser().resolve() / str(bundle_id) + + +def save_prediction(*, bundle_id: str, time_bin_index: int, user_weights: np.ndarray, true_posterior: np.ndarray, scores: Optional[Mapping[str, float]] = None, root: PathLike = 'outputs/posterior_guessing', extra: Optional[Mapping[str, Any]] = None, save_npz: bool = True) -> Dict[str, Any]: + """Score (if needed), append JSONL record, and optionally save per-bin NPZ. + + Returns the record dict that was appended. + """ + user = np.asarray(user_weights, dtype=np.float64) + truth = np.asarray(true_posterior, dtype=np.float64) + if user.shape != truth.shape: + raise ValueError(f'shape mismatch: user {user.shape} vs true {truth.shape}') + + user_p = renormalize(user) + if scores is None: + scores = score_prediction(user_p, truth) + + out_dir = default_predictions_dir(bundle_id, root=root) + out_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now(timezone.utc).isoformat() + record: Dict[str, Any] = { + 'bundle_id': bundle_id, + 'time_bin_index': int(time_bin_index), + 'timestamp': timestamp, + 'hellinger_affinity': float(scores['hellinger_affinity']), + 'hellinger_distance': float(scores.get('hellinger_distance', 1.0 - float(scores['hellinger_affinity']))), + 'cosine_similarity': float(scores['cosine_similarity']), + 'user_mass': float(np.sum(np.clip(user, 0.0, None))), + 'shape': list(user.shape), + } + if extra: + record.update(dict(extra)) + + jsonl_path = out_dir / 'predictions.jsonl' + with open(jsonl_path, 'a', encoding='utf-8') as f: + f.write(json.dumps(record) + '\n') + + if save_npz: + npz_path = out_dir / f'bin_{int(time_bin_index):05d}_{timestamp.replace(":", "").replace("-", "")}.npz' + np.savez_compressed( + npz_path, + user_weights=user, + user_posterior=user_p, + true_posterior=renormalize(truth), + time_bin_index=np.asarray(int(time_bin_index)), + hellinger_affinity=np.asarray(float(scores['hellinger_affinity'])), + cosine_similarity=np.asarray(float(scores['cosine_similarity'])), + metadata_json=np.asarray(json.dumps(record)), + ) + record['npz_path'] = str(npz_path) + + record['jsonl_path'] = str(jsonl_path) + return record diff --git a/pho/posterior_guessing/scoring.py b/pho/posterior_guessing/scoring.py new file mode 100644 index 00000000..ff459305 --- /dev/null +++ b/pho/posterior_guessing/scoring.py @@ -0,0 +1,66 @@ +"""Scoring helpers for painted vs true 2D posteriors.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +import numpy as np + + +def renormalize(weights: np.ndarray, *, eps: float = 0.0) -> np.ndarray: + """Return a non-negative map renormalized to sum to 1. + + Parameters + ---------- + weights: + Raw painted weights (any shape). Negative values are clipped to 0. + eps: + If total mass <= eps, returns a uniform distribution over the same shape. + """ + w = np.asarray(weights, dtype=np.float64) + w = np.clip(w, 0.0, None) + total = float(np.sum(w)) + if not np.isfinite(total) or total <= eps: + return np.full(w.shape, 1.0 / w.size, dtype=np.float64) + return w / total + + +def hellinger_distance(p: np.ndarray, q: np.ndarray) -> float: + """Hellinger distance H(p, q) in [0, 1] for discrete distributions.""" + p_n = renormalize(p) + q_n = renormalize(q) + return float(np.sqrt(0.5 * np.sum((np.sqrt(p_n) - np.sqrt(q_n)) ** 2))) + + +def hellinger_affinity(p: np.ndarray, q: np.ndarray) -> float: + """Primary score: 1 - H(p, q), so identical maps score 1.""" + return float(1.0 - hellinger_distance(p, q)) + + +def cosine_similarity(p: np.ndarray, q: np.ndarray) -> float: + """Secondary score: cosine similarity of flattened renormalized maps.""" + p_n = renormalize(p).ravel() + q_n = renormalize(q).ravel() + denom = float(np.linalg.norm(p_n) * np.linalg.norm(q_n)) + if denom <= 0: + return 0.0 + return float(np.dot(p_n, q_n) / denom) + + +def score_prediction(user_weights: np.ndarray, true_posterior: np.ndarray) -> Dict[str, float]: + """Return hellinger_affinity + cosine_similarity for a user paint vs truth.""" + user_p = renormalize(user_weights) + true_p = renormalize(true_posterior) + return { + 'hellinger_affinity': hellinger_affinity(user_p, true_p), + 'hellinger_distance': hellinger_distance(user_p, true_p), + 'cosine_similarity': cosine_similarity(user_p, true_p), + } + + +def assert_compatible_maps(user_weights: np.ndarray, true_posterior: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + user = np.asarray(user_weights, dtype=np.float64) + truth = np.asarray(true_posterior, dtype=np.float64) + if user.shape != truth.shape: + raise ValueError(f'shape mismatch: user {user.shape} vs true {truth.shape}') + return user, truth diff --git a/pho/posterior_guessing/synthetic.py b/pho/posterior_guessing/synthetic.py new file mode 100644 index 00000000..323520f5 --- /dev/null +++ b/pho/posterior_guessing/synthetic.py @@ -0,0 +1,145 @@ +"""Synthetic 2D placefield / decoder bundle for demos and tests.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Optional, Union + +import numpy as np +from scipy.special import factorial + +from pho.posterior_guessing.export import export_from_arrays + + +PathLike = Union[str, Path] + + +def _gaussian_2d(xx: np.ndarray, yy: np.ndarray, mu_x: float, mu_y: float, sigma: float, peak_rate: float) -> np.ndarray: + return peak_rate * np.exp(-0.5 * (((xx - mu_x) / sigma) ** 2 + ((yy - mu_y) / sigma) ** 2)) + + +def zhang_posterior_2d(tau: float, P_x: np.ndarray, tuning_curves: np.ndarray, spike_counts_t: np.ndarray) -> np.ndarray: + """Compute P(x|n) for one time bin from 2D tuning curves. + + Parameters + ---------- + tau: + Time bin size. + P_x: + Prior over (n_x, n_y), positive mass. + tuning_curves: + (n_neurons, n_x, n_y) firing rates. + spike_counts_t: + (n_neurons,) spike counts in the bin. + """ + n_neurons, n_x, n_y = tuning_curves.shape + assert spike_counts_t.shape == (n_neurons,) + + # Work in flat spatial space for the product + F_flat = tuning_curves.reshape(n_neurons, n_x * n_y).T # (n_pos, n_neurons) + prior = np.asarray(P_x, dtype=np.float64).reshape(n_x * n_y) + prior = prior / np.sum(prior) + + cell_prob = np.ones(n_x * n_y, dtype=np.float64) + for cell in range(n_neurons): + n_i = float(spike_counts_t[cell]) + f_i = F_flat[:, cell] + coeff = 1.0 / float(factorial(n_i)) + cell_prob *= ((tau * f_i) ** n_i) * coeff * np.exp(-tau * f_i) + ## END for cell in range(n_neurons)... + + posterior = prior * cell_prob + total = float(np.sum(posterior)) + if total <= 0 or not np.isfinite(total): + posterior = prior.copy() + else: + posterior = posterior / total + return posterior.reshape(n_x, n_y) + + +def make_synthetic_bundle_arrays(*, n_neurons: int = 8, n_x: int = 20, n_y: int = 16, n_time: int = 40, time_bin_size: float = 0.25, seed: int = 0, arena_xy: tuple = ((0.0, 100.0), (0.0, 80.0))) -> Dict[str, Any]: + """Create in-memory synthetic 2D decoding arrays (no disk write).""" + rng = np.random.default_rng(seed) + (x_min, x_max), (y_min, y_max) = arena_xy + xbin = np.linspace(x_min, x_max, n_x + 1) + ybin = np.linspace(y_min, y_max, n_y + 1) + xcent = 0.5 * (xbin[:-1] + xbin[1:]) + ycent = 0.5 * (ybin[:-1] + ybin[1:]) + xx, yy = np.meshgrid(xcent, ycent, indexing='ij') + + neuron_ids = np.arange(1, n_neurons + 1) + tuning_curves = np.zeros((n_neurons, n_x, n_y), dtype=np.float64) + for i in range(n_neurons): + mu_x = rng.uniform(x_min + 10, x_max - 10) + mu_y = rng.uniform(y_min + 10, y_max - 10) + sigma = rng.uniform(8.0, 18.0) + peak = rng.uniform(4.0, 18.0) + tuning_curves[i] = _gaussian_2d(xx, yy, mu_x, mu_y, sigma, peak) + 0.05 + ## END for i in range(n_neurons)... + + # Soft occupancy biased toward center + occupancy = np.exp(-0.5 * (((xx - xx.mean()) / (0.35 * (x_max - x_min))) ** 2 + ((yy - yy.mean()) / (0.35 * (y_max - y_min))) ** 2)) + occupancy = occupancy * 10.0 + 0.5 + P_x = occupancy / np.sum(occupancy) + + # Simulate animal trajectory as random walk on grid, sample Poisson spikes + ix = n_x // 2 + iy = n_y // 2 + spike_counts = np.zeros((n_neurons, n_time), dtype=np.int64) + p_x_given_n = np.zeros((n_x, n_y, n_time), dtype=np.float64) + for t in range(n_time): + ix = int(np.clip(ix + rng.integers(-2, 3), 0, n_x - 1)) + iy = int(np.clip(iy + rng.integers(-2, 3), 0, n_y - 1)) + local_rates = tuning_curves[:, ix, iy] + counts = rng.poisson(local_rates * time_bin_size) + # Ensure a few bins are empty and most have activity + if t % 11 == 0: + counts = np.zeros_like(counts) + spike_counts[:, t] = counts + p_x_given_n[:, :, t] = zhang_posterior_2d(time_bin_size, P_x, tuning_curves, counts.astype(np.float64)) + ## END for t in range(n_time)... + + time_bin_centers = (np.arange(n_time, dtype=np.float64) + 0.5) * time_bin_size + return { + 'neuron_ids': neuron_ids, + 'tuning_curves': tuning_curves, + 'occupancy': occupancy, + 'P_x': P_x, + 'spike_counts': spike_counts, + 'p_x_given_n': p_x_given_n, + 'time_bin_centers': time_bin_centers, + 'xbin': xbin, + 'ybin': ybin, + 'time_bin_size': float(time_bin_size), + 'metadata': { + 'session_id': 'synthetic_demo', + 'source': 'make_synthetic_bundle', + 'seed': int(seed), + 'ndim': 2, + }, + } + + +def make_synthetic_bundle(out_path: Optional[PathLike] = None, *, filter_empty_bins: bool = True, **kwargs) -> Path: + """Create and optionally save a synthetic guessing bundle. + + Default out_path: data/posterior_guessing/synthetic_demo_pf2d_guessing.npz + """ + arrays = make_synthetic_bundle_arrays(**kwargs) + if out_path is None: + out_path = Path('data/posterior_guessing/synthetic_demo_pf2d_guessing.npz') + return export_from_arrays( + neuron_ids=arrays['neuron_ids'], + tuning_curves=arrays['tuning_curves'], + occupancy=arrays['occupancy'], + P_x=arrays['P_x'], + spike_counts=arrays['spike_counts'], + p_x_given_n=arrays['p_x_given_n'], + time_bin_centers=arrays['time_bin_centers'], + xbin=arrays['xbin'], + ybin=arrays['ybin'], + time_bin_size=arrays['time_bin_size'], + out_path=out_path, + metadata=arrays['metadata'], + filter_empty_bins=filter_empty_bins, + ) diff --git a/pyproject.toml b/pyproject.toml index 60346129..bafd296a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,9 @@ dependencies = [ "pip~=24.0", "streamlit>=1.27.1,<2", "streamlit-extras>=0.3.2,<0.4", + "fastapi>=0.115.0,<1", + "uvicorn[standard]>=0.30.0,<1", + "httpx>=0.27.0,<1", "regex>=2023.10.3,<2024", "panel>=1.4,<2", "params>=0.9.0,<0.10", diff --git a/scripts/export_posterior_guessing_bundle.py b/scripts/export_posterior_guessing_bundle.py new file mode 100644 index 00000000..07580123 --- /dev/null +++ b/scripts/export_posterior_guessing_bundle.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Export a portable 2D posterior-guessing NPZ bundle. + +Examples +-------- + uv run python scripts/export_posterior_guessing_bundle.py --synthetic + uv run python scripts/export_posterior_guessing_bundle.py --from-pkl /path/to/decoder_or_pipeline.pkl --out data/posterior_guessing/session_pf2d_guessing.npz +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +def _load_object_from_pickle(pkl_path: Path): + try: + import dill as pickle # type: ignore + except ImportError: + import pickle + + with open(pkl_path, 'rb') as f: + return pickle.load(f) + + +def _resolve_decoder(obj): + """Best-effort extraction of a 2D decoder from a pipeline or decoder pickle.""" + # Direct decoder-like + if hasattr(obj, 'p_x_given_n') and hasattr(obj, 'pf'): + return obj + + # Common pipeline attribute paths + candidates = [] + if hasattr(obj, 'computation_results'): + try: + for _ctx, result in dict(obj.computation_results).items(): + computed = getattr(result, 'computed_data', None) + if computed is None: + continue + for key in ('pf2D_Decoder', 'pf2D_decoder', 'decoder', 'pf2D'): + if hasattr(computed, key): + candidates.append(getattr(computed, key)) + elif isinstance(computed, dict) and key in computed: + candidates.append(computed[key]) + ## END for key in (...)... + ## END for _ctx, result in dict(obj.computation_results).items()... + except Exception: + pass + + for cand in candidates: + if hasattr(cand, 'p_x_given_n') and hasattr(cand, 'pf'): + return cand + ## END for cand in candidates... + + raise ValueError('Could not locate a 2D decoder with .pf and .p_x_given_n in the pickle') + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description='Export 2D posterior-guessing NPZ bundle') + parser.add_argument('--synthetic', action='store_true', help='Write the synthetic demo bundle') + parser.add_argument('--from-pkl', type=Path, default=None, help='Path to decoder or pipeline pickle') + parser.add_argument('--out', type=Path, default=None, help='Output NPZ path') + parser.add_argument('--session-id', type=str, default=None, help='Session id stored in metadata') + parser.add_argument('--no-filter-empty', action='store_true', help='Keep time bins with zero spikes') + parser.add_argument('--seed', type=int, default=0, help='RNG seed for --synthetic') + parser.add_argument('--n-time', type=int, default=40, help='Number of synthetic time bins') + args = parser.parse_args(argv) + + if not args.synthetic and args.from_pkl is None: + parser.error('Specify --synthetic and/or --from-pkl') + + filter_empty_bins = not args.no_filter_empty + + if args.synthetic: + from pho.posterior_guessing.synthetic import make_synthetic_bundle + + out = args.out or (REPO_ROOT / 'data' / 'posterior_guessing' / 'synthetic_demo_pf2d_guessing.npz') + path = make_synthetic_bundle(out_path=out, filter_empty_bins=filter_empty_bins, seed=args.seed, n_time=args.n_time) + print(f'Wrote synthetic bundle: {path}') + return 0 + + from pho.posterior_guessing.export import export_from_decoder + + pkl_path = Path(args.from_pkl).expanduser().resolve() + obj = _load_object_from_pickle(pkl_path) + decoder = _resolve_decoder(obj) + out = args.out or (REPO_ROOT / 'data' / 'posterior_guessing' / f'{(args.session_id or pkl_path.stem)}_pf2d_guessing.npz') + path = export_from_decoder( + decoder, + out, + metadata={'source_pkl': str(pkl_path)}, + filter_empty_bins=filter_empty_bins, + session_id=args.session_id or pkl_path.stem, + ) + print(f'Wrote bundle from pickle: {path}') + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/test_posterior_guessing.py b/tests/test_posterior_guessing.py new file mode 100644 index 00000000..bc015405 --- /dev/null +++ b/tests/test_posterior_guessing.py @@ -0,0 +1,154 @@ +"""Tests for posterior guessing bundle IO, scoring, and API.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import pytest +from fastapi.testclient import TestClient + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from pho.posterior_guessing.bundle_io import load_bundle, save_bundle, validate_bundle +from pho.posterior_guessing.export import export_from_arrays +from pho.posterior_guessing.persistence import save_prediction +from pho.posterior_guessing.scoring import cosine_similarity, hellinger_affinity, renormalize, score_prediction +from pho.posterior_guessing.synthetic import make_synthetic_bundle, make_synthetic_bundle_arrays + + +@pytest.fixture() +def synthetic_arrays(): + return make_synthetic_bundle_arrays(n_neurons=5, n_x=12, n_y=10, n_time=15, seed=7) + + +def test_validate_and_roundtrip(tmp_path, synthetic_arrays): + out = tmp_path / 'demo_pf2d_guessing.npz' + path = export_from_arrays( + neuron_ids=synthetic_arrays['neuron_ids'], + tuning_curves=synthetic_arrays['tuning_curves'], + occupancy=synthetic_arrays['occupancy'], + P_x=synthetic_arrays['P_x'], + spike_counts=synthetic_arrays['spike_counts'], + p_x_given_n=synthetic_arrays['p_x_given_n'], + time_bin_centers=synthetic_arrays['time_bin_centers'], + xbin=synthetic_arrays['xbin'], + ybin=synthetic_arrays['ybin'], + time_bin_size=synthetic_arrays['time_bin_size'], + out_path=out, + metadata=synthetic_arrays['metadata'], + filter_empty_bins=True, + ) + loaded = load_bundle(path) + validate_bundle(loaded) + assert loaded['tuning_curves'].ndim == 3 + assert loaded['p_x_given_n'].shape[2] == loaded['spike_counts'].shape[1] + assert loaded['bundle_id'] == 'synthetic_demo' + + +def test_make_synthetic_bundle_writes(tmp_path): + path = make_synthetic_bundle(out_path=tmp_path / 'synthetic_demo_pf2d_guessing.npz', n_time=20, seed=1) + assert path.exists() + data = load_bundle(path) + assert data['spike_counts'].shape[1] == data['p_x_given_n'].shape[2] + + +def test_scoring_identical_and_disjoint(): + p = renormalize(np.array([[0.0, 1.0], [0.0, 0.0]])) + assert hellinger_affinity(p, p) == pytest.approx(1.0, abs=1e-9) + assert cosine_similarity(p, p) == pytest.approx(1.0, abs=1e-9) + + q = renormalize(np.array([[0.0, 0.0], [1.0, 0.0]])) + assert hellinger_affinity(p, q) == pytest.approx(0.0, abs=1e-9) + scores = score_prediction(p, q) + assert scores['hellinger_affinity'] == pytest.approx(0.0, abs=1e-9) + + +def test_save_prediction(tmp_path, synthetic_arrays): + truth = synthetic_arrays['p_x_given_n'][:, :, 1] + user = truth.copy() + record = save_prediction( + bundle_id='unit_test', + time_bin_index=1, + user_weights=user, + true_posterior=truth, + root=tmp_path, + save_npz=True, + ) + assert record['hellinger_affinity'] == pytest.approx(1.0, abs=1e-6) + jsonl = Path(record['jsonl_path']) + assert jsonl.exists() + line = jsonl.read_text(encoding='utf-8').strip().splitlines()[-1] + parsed = json.loads(line) + assert parsed['bundle_id'] == 'unit_test' + assert Path(record['npz_path']).exists() + + +@pytest.fixture() +def api_client(tmp_path, monkeypatch, synthetic_arrays): + bundle_dir = tmp_path / 'bundles' + bundle_dir.mkdir() + out_root = tmp_path / 'outputs' + path = export_from_arrays( + neuron_ids=synthetic_arrays['neuron_ids'], + tuning_curves=synthetic_arrays['tuning_curves'], + occupancy=synthetic_arrays['occupancy'], + P_x=synthetic_arrays['P_x'], + spike_counts=synthetic_arrays['spike_counts'], + p_x_given_n=synthetic_arrays['p_x_given_n'], + time_bin_centers=synthetic_arrays['time_bin_centers'], + xbin=synthetic_arrays['xbin'], + ybin=synthetic_arrays['ybin'], + time_bin_size=synthetic_arrays['time_bin_size'], + out_path=bundle_dir / 'synthetic_demo_pf2d_guessing.npz', + metadata=synthetic_arrays['metadata'], + filter_empty_bins=True, + ) + assert path.exists() + + import apps.posterior_guessing.server as server + + monkeypatch.setattr(server, 'DEFAULT_BUNDLE_DIR', bundle_dir) + monkeypatch.setattr(server, 'DEFAULT_OUTPUT_ROOT', out_root) + server._BUNDLE_CACHE.clear() + return TestClient(server.app), server + + +def test_api_hides_posterior_until_reveal(api_client): + client, server = api_client + listing = client.get('/api/bundles') + assert listing.status_code == 200 + bundles = listing.json()['bundles'] + assert len(bundles) == 1 + bundle_id = bundles[0]['bundle_id'] + + bin_payload = client.get(f'/api/bundles/{bundle_id}/bins/0') + assert bin_payload.status_code == 200 + body = bin_payload.json() + assert 'p_x_given_n' not in body + assert 'true_posterior' not in body + assert 'active_cells' in body + n_x, n_y = body['n_x'], body['n_y'] + + # Empty paint rejected + bad = client.post(f'/api/bundles/{bundle_id}/bins/0/reveal', json={'user_weights': np.zeros((n_x, n_y)).tolist(), 'save': False}) + assert bad.status_code == 400 + + data = load_bundle(bundles[0]['path']) + truth = np.asarray(data['p_x_given_n'][:, :, 0]) + # Use a noisy version of truth as the paint + rng = np.random.default_rng(0) + paint = np.clip(truth + 0.01 * rng.random(truth.shape), 0, None) + + revealed = client.post(f'/api/bundles/{bundle_id}/bins/0/reveal', json={'user_weights': paint.tolist(), 'save': True}) + assert revealed.status_code == 200 + result = revealed.json() + assert 'true_posterior' in result + assert 'scores' in result + assert result['scores']['hellinger_affinity'] > 0.5 + assert result['saved'] is not None