lnccbrown / lnccbrown/ssm-simulators
Roadmap: NRE-C support in the HSSM ecosystem
- Dominant language
- Jupyter Notebook
- Stars
- 24
- Forks
- 18
- Avg merge
- 1d 14h
- Merged PRs (30d)
- 9
Description
Productionize the NRE-C ratio-estimation work (`ratio_estimation.ipynb` + `nre_lib/` in HSSMSpine), reusing the existing `ssm-simulators` `generate` CLI for data generation.
```mermaid
flowchart LR
GEN["ssm-simulators\ngenerate CLI"] -->|artifacts| LOAD["artifact loader\n(PR 1)"]
LOAD --> TRAIN["LANfactory\ntrain RatioApproximator"]
TRAIN -->|".keras"| PUB["publish to HF hub\n(PR 5)"]
TRAIN -->|".keras"| INF["inference\nNUTS via NeuralDistribution"]
PUB -.->|"pretrained, by name"| INF
INF -->|".nc"| OUT["analysis / tutorial"]
```
- [ ] 1. `bayesflow`-compatible data artifacts
Define the abstraction everything else builds against.
- Artifact schema for NRE training data: parameter draws + one `(rt, choice)` pair each; metadata (model, bounds, seed, ssms version). Prefer npz/parquet over pickle.
- Loader: artifacts → bayesflow-ready dataset (dict / `OfflineDataset`), validated against the ssms model config.
- Note: ssms 0.13.2 has a `ratio_estimator` config stub but no implementation and no bayesflow tooling — this PR writes the missing contract.
**Done when:** schema documented; loader + round-trip test with a tiny committed fixture.
- [ ] 2. Data generation (ssm-simulators)
- Implement the `ratio_estimator` approach in the `generate` CLI: emit schema-conformant `(theta, rt, choice)` artifacts, skipping the KDE-label step (not needed for NRE).
**Done when:** `generate --config ...` produces artifacts PR 1's loader accepts.
- [ ] 3. Training + inference (probably LANfactory)
Move the `nre_lib` cores out of HSSMSpine:
- Train a `RatioApproximator` from artifacts (offline; keep online simulation as fallback); save `.keras`.
- Sample posteriors via `NeuralDistribution` + NUTS; save `.nc`. Uniform-over-bounds default priors with per-parameter overrides.
- Preserve `nre_lib`'s fixes: per-call simulator seeding, `workers=1` default, `initvals` support.
- bayesflow/keras stay optional deps.
**Done when:** train + sample work for `ddm` and `angle` with tests; no per-model hardcoding.
- [ ] 4. Tutorial
- Rewrite `ratio_estimation.ipynb` as a docs tutorial on the new APIs; retire `nre_lib`.
- Reconcile with the existing NRE/LRE tutorials in HSSM (`sbi_nre_integration.ipynb`, `bayesflow_lre_integration.ipynb`) — one coherent story, no overlap.
**Done when:** runs end to end; notebook CI strategy decided (tiny budgets + committed artifacts, or skip list).
- [ ] 5. Publish trained NREs to Hugging Face
Adapt the publish workflow from `LAN_pipeline_minimal/publish/publish_network.py` (resolve run → stage artifacts → validate → upload → record) for NRE `.keras` artifacts:
- Validation gate before upload (e.g. posterior-recovery smoke check); a skipped gate is not a passed gate.
- Upload layout consumable by HSSM the same way pretrained LANs are (hub repo + canonical filename); record provenance (source run, commit, gate scores).
- Token stays laptop-side; never runs on the cluster.
**Done when:** a trained ratio approximator can be published to the hub and loaded from HSSM by name, with the gate and provenance recorded.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with ratio_estimation.ipynb and the nre_lib/ directory in HSSMSpine, then inspect the existing ssm-simulators generate CLI and the LAN_pipeline_minimal/publish/publish_network.py workflow. The roadmap spans artifact schemas, generation, training, inference, tutorials, and publishing across multiple projects; done requires the listed end-to-end tests, tutorials, and Hugging Face loading workflow.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- huggingface, jupyter-notebook, keras, python
- Domain
- cli, data, documentation, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100