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

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

Đang mở
#82 2 bình luận 0 reaction 1 người được giao Được @hanhou nhận Xem trên GitHub
enhancement
Ngôn ngữ chính
Python
Star
0
Fork
1
Merge trung bình
4 giờ 5 phút
Pull request đã merge (30 ngày)
24

Mô tả

## 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.

Hướng dẫn đóng góp

Chưa lập chỉ mục được hướng dẫn đóng góp cho kho mã nguồn này

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.