gensbi.recipes.joint_pipeline#

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

Classes#

JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Functions#

sample_condition_mask(key, num_samples, theta_dim, x_dim)

sample_structured_conditional_mask(key, num_samples, ...)

Sample structured conditional masks for the Joint model.

Module Contents#

class gensbi.recipes.joint_pipeline.JointPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, method, ch_obs=1, condition_mask_kind='structured', params=None, training_config=None)[source]#

Bases: gensbi.recipes.pipeline.AbstractPipeline

Model-agnostic joint 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 JointWrapper interface.

Parameters:

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

  • train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

  • val_dataset (iterable) – Validation dataset.

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

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

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

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

  • condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

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

  • training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )
abstractmethod _make_model()[source]#

Create and return the model to be trained.

_wrap_model()[source]#

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

classmethod get_default_params(*args, **kwargs)[source]#
Abstractmethod:

get_log_prob_fn(x_o, use_ema=True, prior=None, model_extras=None, **kwargs)[source]#

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.

  • prior (numpyro.distributions.Distribution, optional) –

    Obs-space prior for log-probability evaluation. The method’s prior lives on the full joint space (dim_joint, ch) and cannot be automatically marginalized for arbitrary priors.

    • Default Gaussian: auto-constructed — no need to provide.

    • Custom prior: must supply the correct obs-space marginal.

  • 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

Raises:

ValueError – If the joint prior is non-Gaussian and no prior is provided.

get_loss_fn()[source]#

Return the loss function for training/validation.

get_sampler(x_o, use_ema=True, model_extras=None, **sampler_kwargs)[source]#

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.

Returns:

sampler(key, nsamples) -> samples

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)[source]#
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, prior=None, *, key=None, **kwargs)[source]#

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.

  • prior (numpyro.distributions.Distribution, optional) – Obs-space prior distribution. See get_log_prob_fn() for details.

  • 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)[source]#

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

condition_mask_kind = 'structured'[source]#
dim_joint[source]#
loss_obj[source]#
method[source]#
path[source]#
gensbi.recipes.joint_pipeline.sample_condition_mask(key, num_samples, theta_dim, x_dim, kind='structured')[source]#
gensbi.recipes.joint_pipeline.sample_structured_conditional_mask(key, num_samples, theta_dim, x_dim, p_joint=0.2, p_posterior=0.2, p_likelihood=0.2, p_rnd1=0.2, p_rnd2=0.2, rnd1_prob=0.3, rnd2_prob=0.7)[source]#

Sample structured conditional masks for the Joint model.

Parameters:

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

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

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

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

  • p_joint (float) – Probability of selecting the joint mask.

  • p_posterior (float) – Probability of selecting the posterior mask.

  • p_likelihood (float) – Probability of selecting the likelihood mask.

  • p_rnd1 (float) – Probability of selecting the first random mask.

  • p_rnd2 (float) – Probability of selecting the second random mask.

  • rnd1_prob (float) – Probability of a True value in the first random mask.

  • rnd2_prob (float) – Probability of a True value in the second random mask.

Returns:

condition_mask – Array of shape (num_samples, theta_dim + x_dim) with boolean masks.

Return type:

jnp.ndarray