gensbi.recipes.flux1joint#

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

Classes#

`Flux1JointDiffusionPipeline`	Model-agnostic joint pipeline parameterized by a `GenerativeMethod`.
`Flux1JointFlowPipeline`	Model-agnostic joint pipeline parameterized by a `GenerativeMethod`.
`Flux1JointSMPipeline`	Model-agnostic joint pipeline parameterized by a `GenerativeMethod`.

Functions#

`_flux1joint_config_from_path`(config_path, dim_joint)	Helper to parse common configuration for Flux1Joint pipelines.
`get_default_flux1joint_params`(dim_joint[, in_channels])	Return default parameters for the Flux1Joint model.
`parse_flux1joint_params`(config_path)	Parse a Flux1Joint configuration file.

Module Contents#

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

Bases: gensbi.recipes.joint_pipeline.JointPipeline

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(),
... )

_make_model(params)[source]#: Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)[source]#: Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)[source]#

Initialize the pipeline from a configuration file.

Parameters:

config_path (str) – Path to the configuration file.
**kwargs – Additional keyword arguments forwarded to the constructor.
dim_obs (int)
dim_cond (int)
checkpoint_dir (str)

ch_obs = 1[source]#

dim_joint[source]#

ema_model[source]#

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

Bases: gensbi.recipes.joint_pipeline.JointPipeline

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(),
... )

_make_model(params)[source]#: Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)[source]#: Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)[source]#

Initialize the pipeline from a configuration file.

Parameters:

config_path (str) – Path to the configuration file.
**kwargs – Additional keyword arguments forwarded to the constructor.
dim_obs (int)
dim_cond (int)
checkpoint_dir (str)

ch_obs = 1[source]#

dim_joint[source]#

ema_model[source]#

class gensbi.recipes.flux1joint.Flux1JointSMPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, sde_type='VP', params=None, training_config=None, condition_mask_kind='structured')[source]#

Bases: gensbi.recipes.joint_pipeline.JointPipeline

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.
sde_type (str)

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(),
... )

_make_model(params)[source]#: Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)[source]#: Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)[source]#

Initialize the pipeline from a configuration file.

Parameters:

config_path (str) – Path to the configuration file.
**kwargs – Additional keyword arguments forwarded to the constructor (e.g. sde_type="VE" for score matching).
dim_obs (int)
dim_cond (int)
checkpoint_dir (str)

ch_obs = 1[source]#

dim_joint[source]#

ema_model[source]#

gensbi.recipes.flux1joint._flux1joint_config_from_path(config_path, dim_joint)[source]#

Helper to parse common configuration for Flux1Joint pipelines.

Returns:

params (Flux1JointParams) – The parsed model parameters.
training_config (dict) – The parsed training configuration.
method (str) – The methodology (flow or diffusion) specified in the config.

Parameters:

config_path (str)
dim_joint (int)

gensbi.recipes.flux1joint.get_default_flux1joint_params(dim_joint, in_channels=1)[source]#

Return default parameters for the Flux1Joint model.

Parameters:

dim_joint (int)
in_channels (int)

gensbi.recipes.flux1joint.parse_flux1joint_params(config_path)[source]#

Parse a Flux1Joint configuration file.

Parameters:: config_path (str) – Path to the configuration file.
Returns:: config – Parsed configuration dictionary.
Return type:: dict