gensbi.recipes

Cookie cutter modules for creating and training SBI models.

Submodules

• gensbi.recipes.conditional_pipeline
• gensbi.recipes.flux1
• gensbi.recipes.flux1joint
• gensbi.recipes.joint_pipeline
• gensbi.recipes.pipeline
• gensbi.recipes.simformer
• gensbi.recipes.unconditional_pipeline
• gensbi.recipes.utils

Classes

ConditionalDiffusionPipeline	Diffusion pipeline for training and using a Conditional model for simulation-based inference.
ConditionalFlowPipeline	Flow pipeline for training and using a Conditional model for simulation-based inference.
Flux1DiffusionPipeline	Diffusion pipeline for training and using a Conditional model for simulation-based inference.
Flux1FlowPipeline	Flow pipeline for training and using a Conditional model for simulation-based inference.
Flux1JointDiffusionPipeline	Diffusion pipeline for training and using a Joint model for simulation-based inference.
Flux1JointFlowPipeline	Flow pipeline for training and using a Joint model for simulation-based inference.
JointDiffusionPipeline	Diffusion pipeline for training and using a Joint model for simulation-based inference.
JointFlowPipeline	Flow pipeline for training and using a Joint model for simulation-based inference.
SimformerDiffusionPipeline	Diffusion pipeline for training and using a Joint model for simulation-based inference.
SimformerFlowPipeline	Flow pipeline for training and using a Joint model for simulation-based inference.
UnconditionalDiffusionPipeline	Diffusion pipeline for training and using an Unconditional model for simulation-based inference.
UnconditionalFlowPipeline	Flow pipeline for training and using an Unconditional model for simulation-based inference.

Package Contents

class gensbi.recipes.ConditionalDiffusionPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, id_embedding_strategy=('absolute', 'absolute'), params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Diffusion pipeline for training and using a Conditional model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int or tuple of int) – Dimension of the parameter space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• dim_cond (int or tuple of int) – Dimension of the observation space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• params (ConditionalParams, optional) – Parameters for the Conditional model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the ConditionalDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import ConditionalDiffusionPipeline
from gensbi.models import Flux1, Flux1Params

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be split into observations and conditions.
# We also apply normalization at this stage.
def split_obs_cond(data):
    data = normalize(data, means, stds)
    return (
        data[:, :dim_obs],
        data[:, dim_obs:],
    )  # assuming first dim_obs are obs, last dim_cond are cond


# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the split_obs_cond function to prepare the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Flux1, we need to specify dimensions, embedding strategies, and other architecture details.
params = Flux1Params(
    in_channels=1,
    vec_in_dim=None,
    context_in_dim=1,
    mlp_ratio=3,
    num_heads=2,
    depth=4,
    depth_single_blocks=8,
    axes_dim=[
        10,
    ],
    qkv_bias=True,
    dim_obs=dim_obs,
    dim_cond=dim_cond,
    id_embedding_strategy=("absolute", "absolute"),
    theta=10 * dim_joint,
    rngs=nnx.Rngs(default=42),
    param_dtype=jnp.float32,
)

model = Flux1(params)

# %% Instantiate the pipeline
# The ConditionalDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
training_config = ConditionalDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = ConditionalDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)

plt.savefig(
    "conditional_diffusion_pipeline_marginals.png", dpi=100, bbox_inches="tight"
)
plt.show()

# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

Note

Sampling in the latent space (latent diffusion/flow) is not currently supported.

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

classmethod get_default_training_config()

Return a dictionary of default training configuration parameters.

Returns:

training_config – Default training configuration.

Return type:

dict

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

get_train_step_fn(loss_fn)

Return the training step function, which performs a single optimization step.

Returns:

train_step – JIT-compiled training step function.

Return type:

Callable

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, x_o, nsamples=10000, nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

cond_ids

loss_fn

obs_ids

path

class gensbi.recipes.ConditionalFlowPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, id_embedding_strategy=('absolute', 'absolute'), params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Flow pipeline for training and using a Conditional model for simulation-based inference.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int or tuple of int) – Dimension of the parameter space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• dim_cond (int or tuple of int) – Dimension of the observation space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• ch_obs (int, optional) – Number of channels per token in the observation data. Default is 1.

• ch_cond (int, optional) – Number of channels per token in the conditional data. Default is 1.

• params (ConditionalParams, optional) – Parameters for the Conditional model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the ConditionalFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import ConditionalFlowPipeline
from gensbi.models import Flux1, Flux1Params

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)


# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be split into observations and conditions.
# We also apply normalization at this stage.
def split_obs_cond(data):
    data = normalize(data, means, stds)
    return (
        data[:, :dim_obs],
        data[:, dim_obs:],
    )  # assuming first dim_obs are obs, last dim_cond are cond


# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the split_obs_cond function to prepare the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Flux1, we need to specify dimensions, embedding strategies, and other architecture details.
params = Flux1Params(
    in_channels=1,
    vec_in_dim=None,
    context_in_dim=1,
    mlp_ratio=3,
    num_heads=2,
    depth=4,
    depth_single_blocks=8,
    axes_dim=[
        10,
    ],
    qkv_bias=True,
    dim_obs=dim_obs,
    dim_cond=dim_cond,
    theta=10 * dim_joint,
    id_embedding_strategy=("absolute", "absolute"),
    rngs=nnx.Rngs(default=42),
    param_dtype=jnp.float32,
)

model = Flux1(params)

# %% Instantiate the pipeline
# The ConditionalFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions, and training configuration.
training_config = ConditionalFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = ConditionalFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])
# %% Plot the samples
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("conditional_flow_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

Note

Sampling in the latent space (latent diffusion/flow) is not currently supported.

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

get_train_step_fn(loss_fn)

Return the training step function, which performs a single optimization step.

Returns:

train_step – JIT-compiled training step function.

Return type:

Callable

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, x_o, nsamples=10000, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

cond_ids

loss_fn

obs_ids

p0_obs

path

class gensbi.recipes.Flux1DiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, params=None, training_config=None)

Bases: gensbi.recipes.conditional_pipeline.ConditionalDiffusionPipeline

Diffusion pipeline for training and using a Conditional model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int or tuple of int) – Dimension of the parameter space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• dim_cond (int or tuple of int) – Dimension of the observation space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• params (ConditionalParams, optional) – Parameters for the Conditional model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the ConditionalDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import ConditionalDiffusionPipeline
from gensbi.models import Flux1, Flux1Params

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be split into observations and conditions.
# We also apply normalization at this stage.
def split_obs_cond(data):
    data = normalize(data, means, stds)
    return (
        data[:, :dim_obs],
        data[:, dim_obs:],
    )  # assuming first dim_obs are obs, last dim_cond are cond


# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the split_obs_cond function to prepare the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Flux1, we need to specify dimensions, embedding strategies, and other architecture details.
params = Flux1Params(
    in_channels=1,
    vec_in_dim=None,
    context_in_dim=1,
    mlp_ratio=3,
    num_heads=2,
    depth=4,
    depth_single_blocks=8,
    axes_dim=[
        10,
    ],
    qkv_bias=True,
    dim_obs=dim_obs,
    dim_cond=dim_cond,
    id_embedding_strategy=("absolute", "absolute"),
    theta=10 * dim_joint,
    rngs=nnx.Rngs(default=42),
    param_dtype=jnp.float32,
)

model = Flux1(params)

# %% Instantiate the pipeline
# The ConditionalDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
training_config = ConditionalDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = ConditionalDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)

plt.savefig(
    "conditional_diffusion_pipeline_marginals.png", dpi=100, bbox_inches="tight"
)
plt.show()

# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

Note

Sampling in the latent space (latent diffusion/flow) is not currently supported.

_get_default_params()

Return default parameters for the Flux1 model.

_make_model(params)

Create and return the Flux1 model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_cond = 1

ch_obs = 1

dim_cond

dim_obs

ema_model

class gensbi.recipes.Flux1FlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, params=None, training_config=None)

Bases: gensbi.recipes.conditional_pipeline.ConditionalFlowPipeline

Flow pipeline for training and using a Conditional model for simulation-based inference.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int or tuple of int) – Dimension of the parameter space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• dim_cond (int or tuple of int) – Dimension of the observation space (number of tokens). Can represent unstructured data, time-series, or patchified 2D images. For images, provide a tuple (height, width).

• ch_obs (int, optional) – Number of channels per token in the observation data. Default is 1.

• ch_cond (int, optional) – Number of channels per token in the conditional data. Default is 1.

• params (ConditionalParams, optional) – Parameters for the Conditional model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the ConditionalFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import ConditionalFlowPipeline
from gensbi.models import Flux1, Flux1Params

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)


# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be split into observations and conditions.
# We also apply normalization at this stage.
def split_obs_cond(data):
    data = normalize(data, means, stds)
    return (
        data[:, :dim_obs],
        data[:, dim_obs:],
    )  # assuming first dim_obs are obs, last dim_cond are cond


# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the split_obs_cond function to prepare the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(split_obs_cond)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Flux1, we need to specify dimensions, embedding strategies, and other architecture details.
params = Flux1Params(
    in_channels=1,
    vec_in_dim=None,
    context_in_dim=1,
    mlp_ratio=3,
    num_heads=2,
    depth=4,
    depth_single_blocks=8,
    axes_dim=[
        10,
    ],
    qkv_bias=True,
    dim_obs=dim_obs,
    dim_cond=dim_cond,
    theta=10 * dim_joint,
    id_embedding_strategy=("absolute", "absolute"),
    rngs=nnx.Rngs(default=42),
    param_dtype=jnp.float32,
)

model = Flux1(params)

# %% Instantiate the pipeline
# The ConditionalFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions, and training configuration.
training_config = ConditionalFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = ConditionalFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])
# %% Plot the samples
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("conditional_flow_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

Note

Sampling in the latent space (latent diffusion/flow) is not currently supported.

_get_default_params()

Return default parameters for the Flux1 model.

_make_model(params)

Create and return the Flux1 model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_cond = 1

ch_obs = 1

dim_cond

dim_obs

ema_model

class gensbi.recipes.Flux1JointDiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointDiffusionPipeline

Diffusion pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from gensbi.recipes import JointDiffusionPipeline
from gensbi.utils.plotting import plot_marginals

from gensbi.models import Simformer, SimformerParams
import matplotlib.pyplot as plt

from numpyro import distributions as dist


from flax import nnx

# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_diffusion_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

_get_default_params()

Return default parameters for the Simformer model.

_make_model(params)

Create and return the Simformer model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_obs = 1

dim_joint

ema_model

class gensbi.recipes.Flux1JointFlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointFlowPipeline

Flow pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (JointParams, optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import JointFlowPipeline
from gensbi.models import Simformer, SimformerParams

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).
def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
# We also map the process_data function to prepare (normalize) the data for the model.
# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_flow_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

_get_default_params()

Return default parameters for the Simformer model.

_make_model(params)

Create and return the Simformer model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_obs = 1

dim_joint

ema_model

class gensbi.recipes.JointDiffusionPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.pipeline.AbstractPipeline

Diffusion pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from gensbi.recipes import JointDiffusionPipeline
from gensbi.utils.plotting import plot_marginals

from gensbi.models import Simformer, SimformerParams
import matplotlib.pyplot as plt

from numpyro import distributions as dist


from flax import nnx

# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_diffusion_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

classmethod get_default_training_config()

Return a dictionary of default training configuration parameters.

Returns:

training_config – Default training configuration.

Return type:

dict

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, x_o, nsamples=10000, nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

condition_mask_kind = 'structured'

loss_fn

path

class gensbi.recipes.JointFlowPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.pipeline.AbstractPipeline

Flow pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (JointParams, optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import JointFlowPipeline
from gensbi.models import Simformer, SimformerParams

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).
def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
# We also map the process_data function to prepare (normalize) the data for the model.
# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_flow_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, x_o, nsamples=10000, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

condition_mask_kind = 'structured'

dim_joint

loss_fn

p0_joint

p0_obs

path

class gensbi.recipes.SimformerDiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, edge_mask=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointDiffusionPipeline

Diffusion pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from gensbi.recipes import JointDiffusionPipeline
from gensbi.utils.plotting import plot_marginals

from gensbi.models import Simformer, SimformerParams
import matplotlib.pyplot as plt

from numpyro import distributions as dist


from flax import nnx

# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_diffusion_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

_get_default_params()

Return default parameters for the Simformer model.

_make_model(params)

Create and return the Simformer model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

sample(key, x_o, nsamples=10000, nsteps=18, use_ema=True, return_intermediates=False)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

ch_obs = 1

dim_joint

edge_mask = None

ema_model

class gensbi.recipes.SimformerFlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, edge_mask=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointFlowPipeline

Flow pipeline for training and using a Joint model for simulation-based inference.

Parameters:

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• dim_cond (int) – Dimension of the observation space.

• ch_obs (int, optional) – Number of channels for the observation space. Default is 1.

• params (JointParams, optional) – Parameters for the Joint model. If None, default parameters are used.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

• condition_mask_kind (str, optional) – Kind of condition mask to use. One of [“structured”, “posterior”].

Examples

Minimal example on how to instantiate and use the JointFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from numpyro import distributions as dist
from flax import nnx

from gensbi.recipes import JointFlowPipeline
from gensbi.models import Simformer, SimformerParams

from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


# %%

theta_prior = dist.Uniform(
    low=jnp.array([-2.0, -2.0, -2.0]), high=jnp.array([2.0, 2.0, 2.0])
)

dim_obs = 3
dim_cond = 3
dim_joint = dim_obs + dim_cond


# %%
def simulator(key, nsamples):
    theta_key, sample_key = jax.random.split(key, 2)
    thetas = theta_prior.sample(theta_key, (nsamples,))

    xs = thetas + 1 + jax.random.normal(sample_key, thetas.shape) * 0.1

    thetas = thetas[..., None]
    xs = xs[..., None]

    # when making a dataset for the joint pipeline, thetas need to come first
    data = jnp.concatenate([thetas, xs], axis=1)

    return data


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator is a simple function that generates parameters (theta) and data (x).
# In this example, we use a simple Gaussian simulator.
train_data = simulator(jax.random.PRNGKey(0), 100_000)
val_data = simulator(jax.random.PRNGKey(1), 2000)
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).


# %% Prepare the data for the pipeline
# The pipeline expects the data to be normalized but not split (for joint pipelines).
def process_data(data):
    return normalize(data, means, stds)


# %%
train_data.shape

# %%

# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
# We also map the process_data function to prepare (normalize) the data for the model.
# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data for the model.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

# %% Define your model
# specific model parameters are defined here.
# For Simformer, we need to specify dimensions, embedding strategies, and other architecture details.
params = SimformerParams(
    rngs=nnx.Rngs(0),
    in_channels=1,
    dim_value=20,
    dim_id=10,
    dim_condition=10,
    dim_joint=dim_joint,
    fourier_features=128,
    num_heads=4,
    num_layers=6,
    widening_factor=3,
    qkv_features=40,
    num_hidden_layers=1,
)

model = Simformer(params)

# %% Instantiate the pipeline
# The JointFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
# We also specify the condition_mask_kind, which determines how conditioning is handled during training.
training_config = JointFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = JointFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    dim_cond,
    condition_mask_kind="posterior",
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# To generate samples, we first need an observation (and its corresponding condition).
# We generate a new sample from the simulator, normalize it, and extract the condition x_o.

new_sample = simulator(jax.random.PRNGKey(20), 1)
true_theta = new_sample[:, :dim_obs, :]  # extract observation from the joint sample

new_sample = normalize(new_sample, means, stds)
x_o = new_sample[:, dim_obs:, :]  # extract condition from the joint sample

# Then we invoke the pipeline's sample method.
samples = pipeline.sample(rngs.sample(), x_o, nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means[:dim_obs], stds[:dim_obs])

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples
# against the true parameters.
plot_marginals(
    np.array(samples[..., 0]),
    gridsize=30,
    true_param=np.array(true_theta[0, :, 0]),
    range=[(1, 3), (1, 3), (-0.6, 0.5)],
)
plt.savefig("joint_flow_pipeline_marginals.png", dpi=100, bbox_inches="tight")
plt.show()

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

_get_default_params()

Return default parameters for the Simformer model.

_make_model(params)

Create and return the Simformer model to be trained.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• config_path (str) – Path to the configuration file.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

sample(key, x_o, nsamples=10000, step_size=0.01, use_ema=True, time_grid=None)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

ch_obs = 1

dim_joint

edge_mask = None

ema_model

class gensbi.recipes.UnconditionalDiffusionPipeline(model, train_dataset, val_dataset, dim_obs, ch_obs=1, params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Diffusion pipeline for training and using an Unconditional model for simulation-based inference.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• ch_obs (int) – Number of channels in the observation space.

• params (optional) – Parameters for the model. Serves no use if a custom model is provided.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the UnconditionalDiffusionPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from gensbi.recipes import UnconditionalDiffusionPipeline
from gensbi.utils.model_wrapping import _expand_dims, _expand_time
from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt
from gensbi.models import Simformer, SimformerParams


from flax import nnx


# %% define a simulator
def simulator(key, nsamples):
    return 3 + jax.random.normal(key, (nsamples, 2)) * jnp.array([0.5, 1]).reshape(
        1, 2
    )  # a simple 2D gaussian


# %%


# %% Define your training and validation datasets.
train_data = simulator(jax.random.PRNGKey(0), 100_000).reshape(-1, 2, 1)
val_data = simulator(jax.random.PRNGKey(1), 2000).reshape(-1, 2, 1)
# %%
# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means

    return normalize(data, means, stds)


def process_data(data):
    return normalize(data, means, stds)


# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)
# %% Define your model
# Here we define a MLP velocity field model,
# this model only works for inputs of shape (batch, dim, 1).
# For more complex models, please refer to the transformer-based models in gensbi.models.


class MLP(nnx.Module):
    def __init__(self, input_dim: int = 2, hidden_dim: int = 512, *, rngs: nnx.Rngs):

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim

        din = input_dim + 1

        self.linear1 = nnx.Linear(din, self.hidden_dim, rngs=rngs)
        self.linear2 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear3 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear4 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear5 = nnx.Linear(self.hidden_dim, self.input_dim, rngs=rngs)

    def __call__(self, t: jax.Array, obs: jax.Array, node_ids, *args, **kwargs):
        obs = _expand_dims(obs)[
            ..., 0
        ]  # for this specific model, we use samples of shape (batch, dim), while for transformer models we use (batch, dim, c)
        t = _expand_time(t)
        if t.ndim == 3:
            t = t.reshape(t.shape[0], t.shape[1])
        t = jnp.broadcast_to(t, (obs.shape[0], 1))

        h = jnp.concatenate([obs, t], axis=-1)

        x = self.linear1(h)
        x = jax.nn.gelu(x)

        x = self.linear2(x)
        x = jax.nn.gelu(x)

        x = self.linear3(x)
        x = jax.nn.gelu(x)

        x = self.linear4(x)
        x = jax.nn.gelu(x)

        x = self.linear5(x)

        return x[..., None]  # return shape (batch, dim, 1)


model = MLP(
    rngs=nnx.Rngs(42)
)  # your nnx.Module model here, e.g., a simple MLP, or the Simformer model
# if you define a custom model, it should take as input the following arguments:
#    t: Array,
#    obs: Array,
#    node_ids: Array (optional, if your model is a transformer-based model)
#    *args
#    **kwargs

# the obs input should have shape (batch_size, dim_joint, c), and the output will be of the same shape
# %% Instantiate the pipeline
dim_obs = 2  # Dimension of the parameter space
ch_obs = 1  # Number of channels of the parameter space

# The UnconditionalDiffusionPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
training_config = UnconditionalDiffusionPipeline.get_default_training_config()
training_config["nsteps"] = 10000

pipeline = UnconditionalDiffusionPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    ch_obs,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# We generate new samples using the trained model.
samples = pipeline.sample(rngs.sample(), nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means, stds)

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples.
samples.mean(axis=0), samples.std(axis=0)
# %%

plot_marginals(
    np.array(samples[..., 0]), true_param=[3, 3], gridsize=20, range=[(-2, 8), (-2, 8)]
)
plt.savefig("unconditional_diffusion_samples.png", dpi=300, bbox_inches="tight")
plt.show()

# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

classmethod get_default_training_config(sde='EDM')

Return a dictionary of default training configuration parameters.

Returns:

training_config – Default training configuration.

Return type:

dict

get_loss_fn()

Return the loss function for training/validation.

get_sampler(nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, nsamples=10000, nsteps=18, use_ema=True, return_intermediates=False, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

abstractmethod sample_batched(*args, **kwargs)

Generate samples from the trained model in batches.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int) – Number of samples to generate.

• chunk_size (int, optional) – Size of each batch for sampling. Default is 50.

• show_progress_bars (bool, optional) – Whether to display progress bars during sampling. Default is True.

• args (tuple) – Additional positional arguments for the sampler.

• kwargs (dict) – Additional keyword arguments for the sampler.

Returns:

samples – Generated samples of shape (nsamples, batch_size_cond, dim_obs, ch_obs).

Return type:

array-like

loss_fn

obs_ids

path

class gensbi.recipes.UnconditionalFlowPipeline(model, train_dataset, val_dataset, dim_obs, ch_obs=1, params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Flow pipeline for training and using an Unconditional model for simulation-based inference.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (grain dataset or iterator over batches) – Training dataset.

• val_dataset (grain dataset or iterator over batches) – Validation dataset.

• dim_obs (int) – Dimension of the parameter space.

• ch_obs (int) – Number of channels in the observation space.

• params (optional) – Parameters for the model. Serves no use if a custom model is provided.

• training_config (dict, optional) – Configuration for training. If None, default configuration is used.

Examples

Minimal example on how to instantiate and use the UnconditionalFlowPipeline:

# %% Imports
import os

# Set JAX backend (use 'cuda' for GPU, 'cpu' otherwise)
# os.environ["JAX_PLATFORMS"] = "cuda"

import grain
import numpy as np
import jax
from jax import numpy as jnp
from gensbi.recipes import UnconditionalFlowPipeline
from gensbi.utils.model_wrapping import _expand_dims, _expand_time
from gensbi.utils.plotting import plot_marginals
import matplotlib.pyplot as plt


from flax import nnx


# %% define a simulator
def simulator(key, nsamples):
    return 3 + jax.random.normal(key, (nsamples, 2)) * jnp.array([0.5, 1]).reshape(
        1, 2
    )  # a simple 2D gaussian


# %% Define your training and validation datasets.
# We generate a training dataset and a validation dataset using the simulator.
# The simulator generates samples from a 2D Gaussian distribution.
train_data = simulator(jax.random.PRNGKey(0), 100_000).reshape(-1, 2, 1)
val_data = simulator(jax.random.PRNGKey(1), 2000).reshape(-1, 2, 1)

# %% Normalize the dataset
# It is important to normalize the data to have zero mean and unit variance.
# This helps the model training process.
means = jnp.mean(train_data, axis=0)
stds = jnp.std(train_data, axis=0)


def normalize(data, means, stds):
    return (data - means) / stds


def unnormalize(data, means, stds):
    return data * stds + means


def process_data(data):
    return normalize(data, means, stds)


# %% Create the input pipeline using Grain
# We use Grain to create an efficient input pipeline.
# This involves shuffling, repeating for multiple epochs, and batching the data.
# We also map the process_data function to prepare (normalize) the data.
batch_size = 256

train_dataset_grain = (
    grain.MapDataset.source(np.array(train_data))
    .shuffle(42)
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)

val_dataset_grain = (
    grain.MapDataset.source(np.array(val_data))
    .shuffle(
        42
    )  # Use a different seed/strategy for validation if needed, but shuffling is fine
    .repeat()
    .to_iter_dataset()
    .batch(batch_size)
    .map(process_data)
    # .mp_prefetch() # Uncomment if you want to use multiprocessing prefetching
)


# %% Define your model
# Here we define a MLP velocity field model,
# this model only works for inputs of shape (batch, dim, 1).
# For more complex models, please refer to the transformer-based models in gensbi.models.
class MLP(nnx.Module):
    def __init__(self, input_dim: int = 2, hidden_dim: int = 128, *, rngs: nnx.Rngs):

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim

        din = input_dim + 1

        self.linear1 = nnx.Linear(din, self.hidden_dim, rngs=rngs)
        self.linear2 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear3 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear4 = nnx.Linear(self.hidden_dim, self.hidden_dim, rngs=rngs)
        self.linear5 = nnx.Linear(self.hidden_dim, self.input_dim, rngs=rngs)

    def __call__(self, t: jax.Array, obs: jax.Array, node_ids, *args, **kwargs):
        obs = _expand_dims(obs)[
            ..., 0
        ]  # for this specific model, we use samples of shape (batch, dim), while for transformer models we use (batch, dim, c)
        t = _expand_time(t)
        t = jnp.broadcast_to(t, (obs.shape[0], 1))

        h = jnp.concatenate([obs, t], axis=-1)

        x = self.linear1(h)
        x = jax.nn.gelu(x)

        x = self.linear2(x)
        x = jax.nn.gelu(x)

        x = self.linear3(x)
        x = jax.nn.gelu(x)

        x = self.linear4(x)
        x = jax.nn.gelu(x)

        x = self.linear5(x)

        return x[..., None]  # return shape (batch, dim, 1)


model = MLP(
    rngs=nnx.Rngs(42)
)  # your nnx.Module model here, e.g., a simple MLP, or the Simformer model
# if you define a custom model, it should take as input the following arguments:
#    t: Array,
#    obs: Array,
#    node_ids: Array (optional, if your model is a transformer-based model)
#    *args
#    **kwargs

# the obs input should have shape (batch_size, dim_joint, c), and the output will be of the same shape

# %% Instantiate the pipeline
# The UnconditionalFlowPipeline handles the training loop and sampling.
# We configure it with the model, datasets, dimensions using a default training configuration.
training_config = UnconditionalFlowPipeline.get_default_training_config()
training_config["nsteps"] = 10000

dim_obs = 2  # Dimension of the parameter space
ch_obs = 1  # Number of channels of the parameter space

pipeline = UnconditionalFlowPipeline(
    model,
    train_dataset_grain,
    val_dataset_grain,
    dim_obs,
    ch_obs,
    training_config=training_config,
)

# %% Train the model
# We create a random key for training and start the training process.
rngs = nnx.Rngs(42)
pipeline.train(
    rngs, save_model=False
)  # if you want to save the model, set save_model=True

# %% Sample from the posterior
# We generate new samples using the trained model.
samples = pipeline.sample(rngs.sample(), nsamples=100_000)
# Finally, we unnormalize the samples to get them back to the original scale.
samples = unnormalize(samples, means, stds)

# %% Plot the samples
# We verify the model's performance by plotting the marginal distributions of the generated samples.
plot_marginals(
    np.array(samples[..., 0]), true_param=[3, 3], gridsize=30, range=[(-2, 8), (-2, 8)]
)
plt.savefig("unconditional_flow_samples.png", dpi=300, bbox_inches="tight")
plt.show()
# %%

Note

If you plan on using multiprocessing prefetching, ensure that your script is wrapped in a if __name__ == "__main__": guard. See https://docs.python.org/3/library/multiprocessing.html

abstractmethod _get_default_params()

Return a dictionary of default model parameters.

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

get_loss_fn()

Return the loss function for training/validation.

get_sampler(step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

classmethod init_pipeline_from_config()

Abstractmethod:

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

sample(key, nsamples=10000, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Generate samples from the trained model.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

abstractmethod sample_batched(*args, **kwargs)

Generate samples from the trained model in batches.

Parameters:

• key (jax.random.PRNGKey) – Random number generator key.

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int) – Number of samples to generate.

• chunk_size (int, optional) – Size of each batch for sampling. Default is 50.

• show_progress_bars (bool, optional) – Whether to display progress bars during sampling. Default is True.

• args (tuple) – Additional positional arguments for the sampler.

• kwargs (dict) – Additional keyword arguments for the sampler.

Returns:

samples – Generated samples of shape (nsamples, batch_size_cond, dim_obs, ch_obs).

Return type:

array-like

loss_fn

obs_ids

p0_obs

path