lnccbrown / lnccbrown/ssm-simulators
Add .to_bayesflow() method to Simulator class
- Dominant language
- Jupyter Notebook
- Stars
- 24
- Forks
- 18
- Avg merge
- 1d 14h
- Merged PRs (30d)
- 9
Description
## Summary
Add a `.to_bayesflow()` method to the `Simulator` class that returns a BayesFlow-compatible simulator object, enabling users to train BayesFlow amortized inference networks directly from `ssm-simulators` models.
## Motivation
[BayesFlow](https://github.com/bayesflow-org/bayesflow) provides amortized Bayesian inference via neural networks. Currently, users who want to use BayesFlow with SSM models need to manually write wrapper functions that bridge the two libraries. A `.to_bayesflow()` method would make this seamless.
## BayesFlow Simulator Interface
BayesFlow expects simulators to implement a `sample(batch_shape, **kwargs) -> dict[str, np.ndarray]` method. The returned dictionary maps variable names to arrays. BayesFlow composes priors and likelihoods via `SequentialSimulator`, where outputs from earlier stages are passed as kwargs to later stages based on matching key names.
A minimal BayesFlow simulator can also be a `LambdaSimulator` wrapping any callable `f(batch_shape) -> dict`.
## Proposed API
```python
sim = Simulator("ddm")
# Returns a BayesFlow-compatible simulator (prior + likelihood)
bf_simulator = sim.to_bayesflow(
prior=some_prior_distribution, # optional: prior over model params
)
# The returned object should work directly with BayesFlow's training pipeline
# e.g., bayesflow.make_simulator([bf_simulator]) or used as-is
```
## Implementation Considerations
- The method should return an object compatible with `bayesflow.simulators.Simulator` (i.e., has a `sample(batch_shape) -> dict` method)
- The returned dict should contain at minimum `"obs_data"` (rts/choices) and `"params"` (theta values)
- Should handle the parameter bounds from `self.config["param_bounds"]` for defining a default uniform prior if none is provided
- Consider whether BayesFlow should be an optional dependency (imported only when `.to_bayesflow()` is called)
- Key mapping between ssm-simulators output (`rts`, `choices`) and BayesFlow's expected dict format needs to be defined
## Example Usage Vision
```python
from ssms.basic_simulators import Simulator
import bayesflow as bf
# Setup simulator
sim = Simulator("angle")
# Convert to BayesFlow simulator
bf_sim = sim.to_bayesflow()
# Use directly in BayesFlow training pipeline
workflow = bf.BasicWorkflow(simulator=bf_sim)
workflow.fit_online(...)
```
## Related
- BayesFlow repo: https://github.com/bayesflow-org/bayesflow
- Current `Simulator` class: `ssms/basic_simulators/simulator_class.py`
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.