lnccbrown / lnccbrown/ssm-simulators

Add .to_bayesflow() method to Simulator class

Open
#275 1 comment 0 reactions 1 assignee Claimed by @stefanradev93 View on GitHub
connectors linear-ssm-simulators
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.