AllenNeuralDynamics / AllenNeuralDynamics/aind-dynamic-foraging-bfm-wrapper

HB runs log no figures; plotting.py is orphaned and session sites are never persisted

Open
#82 2 comments 0 reactions 1 assignee Claimed by @hanhou View on GitHub
enhancement
Dominant language
Python
Stars
0
Forks
1
Avg merge
4h 5m
Merged PRs (30d)
24

Description

## Context

Gate 2 of dispatcher #107 asked for posterior/diagnostic figures from a real HB fit. There
are none, and there is no code path that could produce them from a run.

## Findings

| # | Verified | Where |
|---|---|---|
| 1 | `hb_trainer.py` contains **no** plotting code at all — no `matplotlib`, `savefig`, `plot`, or figure logging. The only media a run emits is the per-subject table. | `code/model_trainers/hb_trainer.py` |
| 2 | `hierarchical_bayes/plotting.py` (`plot_population_recovery`, `plot_conditioning_curve`, `plot_shrinkage`, `to_bounded`) is imported by **nothing** in either repo — the only grep match is a stale `__pycache__` `.pyc`. | models `hierarchical_bayes/plotting.py` |
| 3 | Session-level sites are never persisted: `save_fit(..., include_session_sites=False)` is the default and `HBTrainer` never overrides it, nor does `hb_hattori.yaml` expose a knob. The D10 artifact's `sites_saved` confirms population + subject sites only. | models `hierarchical_bayes/artifacts.py:88`, `:115-117` |

Consequence of finding 3, from `save_fit`'s own docstring: session-level parameters and
per-session log likelihoods "are what WAIC, PSIS-LOO and per-session comparisons need". None
of those are possible from a saved fit today.

The three committed figures in `studies/08-hb-vs-gru-heldout/figures/` were therefore
produced outside the run path, on synthetic data with known ground truth. They validate the
plotting code and must not be read as results.

## Trajectory figure — feasible, and without touching the likelihood

`hattori2019_choice_prob` computes Q inside a `jax.lax.scan` but returns only
`choice_prob.T`; the Q trajectory is the scan carry and is discarded
(`likelihood.py:74-95`). So:

- **Per-trial choice probability with a credible band**: needs no likelihood change, only
the session-level parameters from finding 3.
- **Per-trial Q trajectory**: recoverable by replaying the deterministic update rule from
`(choices, rewards, learn_rate_rew, learn_rate_unrew, forget_rate_unchosen)` in a plotting
helper. Emitting Q from the scan directly would be a one-line change to the hot likelihood
every fit runs, so the replay route is preferred.

This is the figure worth having: a credible band on a latent trajectory is exactly what
per-session MLE structurally cannot produce.

## Done when

- [ ] A run logs sampler-trustworthiness figures: trace + rank per population parameter, and an energy/divergence summary
- [ ] A run logs per-parameter population posterior densities on the interpretable (post-`Phi`) scale, via `to_bounded`
- [ ] A run logs the shrinkage figure (pooled vs unpooled subject estimates)
- [ ] The zero-shot -> few-shot(k) conditioning curve is logged with GRU / `baseline_rl` comparators on shared axes
- ~~ `include_session_sites` is reachable from config, so per-session diagnostics and LOO/WAIC become possible — **moved to dispatcher #115** (latent DV inference stack)~~ — **moved to dispatcher #115** (done there for the config knob; the Q replay is still open).
- ~~ Per-session choice-probability trace with a credible band, and Q trajectory by replay — Q-trajectory replay **moved to dispatcher #115**; the choice-probability trace with a credible band stays here~~ — **moved to dispatcher #115** (done there for the config knob; the Q replay is still open).
- [ ] The committed synthetic figures are relabelled so they cannot be mistaken for results

## Notes

Verified at wrapper `d7309dd` / models `26ce8eb`. Blocked in part by the artifact
round-trip bug filed alongside this — figures generated from a reloaded fit need
`az.from_netcdf` to work.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.