setu.simulate_dataset#

setu.simulate_dataset(simulator_fn, prior_fn, n_samples, key, *, conditions_fn=None)[source]#

Generate a SimulationDataset from a simulator and prior.

Convenience function that calls the prior to sample theta, then vmaps the simulator over all theta samples to produce observations.

Parameters:

  • simulator_fn (Callable) – Function (theta, key) -> x that maps a single parameter vector to a single observation. Will be vmapped over all samples.

  • prior_fn (Callable) – Function (key) -> theta that draws a single parameter sample. Will be vmapped to draw n_samples.

  • n_samples (int) – Number of (theta, x) pairs to generate.

  • key (Array) – JAX random key.

  • conditions_fn (Optional[Callable]) – Optional function (key) -> conditions that draws a single condition vector. Will be vmapped to draw n_samples. If provided, simulator_fn should accept (theta, key, conditions) instead.

Return type:

SimulationDataset

Returns:

SimulationDataset with theta and x (and optionally conditions).

Example

>>> def prior(key):
...     return jax.random.normal(key, (3,))
>>> def simulator(theta, key):
...     return theta + 0.1 * jax.random.normal(key, theta.shape)
>>> dataset = simulate_dataset(simulator, prior, 10000, key)