gensbi.recipes.conditional_pipeline

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

Classes

ConditionalPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Module Contents

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

Bases: gensbi.recipes.pipeline.AbstractPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the ConditionalWrapper interface.

Parameters:

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

• train_dataset (iterable) – Training dataset yielding (obs, cond) batches.

• val_dataset (iterable) – Validation dataset yielding (obs, cond) batches.

• dim_obs (int or tuple of int) – Dimension of the observation/parameter space.

• dim_cond (int or tuple of int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

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

• ch_cond (int, optional) – Number of channels per conditioning token. Default is 1.

• id_embedding_strategy (tuple of str, optional) – Embedding strategy for observation and conditioning IDs. Default is ("absolute", "absolute").

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration. If None, uses defaults augmented by method.get_extra_training_config().

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = ConditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=5, dim_cond=3,
...     method=FlowMatchingMethod(),
... )

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_params(*args, **kwargs)

Abstractmethod:

get_log_prob_fn(x_o, use_ema=True, model_extras=None, **kwargs)

Get a log-probability function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• model_extras (dict, optional) – Additional model extras. Cannot override protected keys.

• **kwargs – Forwarded to method.build_log_prob_fn.

Returns:

log_prob_fn(x_1) -> log_prob

Return type:

Callable

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, use_ema=True, model_extras=None, **sampler_kwargs)

Get a sampler function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• model_extras (dict, optional) – Additional keyword arguments passed to the model during sampling (e.g. {"edge_mask": mask}). Cannot override the protected keys cond, obs_ids, cond_ids.

• **sampler_kwargs – Forwarded to method.build_sampler_fn (e.g. step_size, nsteps, solver, time_grid).

Returns:

sampler(key, nsamples) -> samples

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)

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

log_prob(x_1, x_o, use_ema=True, *, key=None, **kwargs)

Compute log-probability of x_1 given x_o.

Parameters:

• x_1 (array-like) – Data samples to evaluate.

• x_o (array-like) – Conditioning variable.

• use_ema (bool, optional) – Use the EMA model. Default is True.

• key (jax.random.PRNGKey, optional) – Required when exact_divergence=False (Hutchinson).

• **kwargs – Forwarded to get_log_prob_fn().

Returns:

Log-probabilities.

Return type:

Array

sample(key, x_o, nsamples=10000, use_ema=True, **sampler_kwargs)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

• nsamples (int, optional) – Number of samples. Default is 10 000.

• use_ema (bool, optional) – Use the EMA model. Default is True.

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

loss_obj

method

path