lnccbrown / lnccbrown/ssm-simulators

Roadmap: NRE-C support in the HSSM ecosystem

Open
#343 2 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.