gensbi.recipes.flux1#

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

Classes#

Flux1DiffusionPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Flux1FlowPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Flux1SMPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Functions#

_flux1_config_from_path(config_path, dim_obs, dim_cond)

Helper to parse common configuration for Flux1 pipelines.

get_default_flux1_params(dim_obs, dim_cond[, ch_obs, ...])

Return default parameters for the Flux1 model.

parse_flux1_params(config_path)

Parse a Flux1 configuration file.

Module Contents#

class gensbi.recipes.flux1.Flux1DiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, params=None, training_config=None)[source]#

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

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(),
... )
_make_model(params)[source]#

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)[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)

ema_model[source]#
class gensbi.recipes.flux1.Flux1FlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, params=None, training_config=None)[source]#

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

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(),
... )
_make_model(params)[source]#

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)[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)

ema_model[source]#
class gensbi.recipes.flux1.Flux1SMPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, sde_type='VP', params=None, training_config=None)[source]#

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

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

  • sde_type (str)

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(),
... )
_make_model(params)[source]#

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)[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)

ema_model[source]#
gensbi.recipes.flux1._flux1_config_from_path(config_path, dim_obs, dim_cond)[source]#

Helper to parse common configuration for Flux1 pipelines.

Returns:

  • params (Flux1Params) – 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_obs (int)

  • dim_cond (int)

gensbi.recipes.flux1.get_default_flux1_params(dim_obs, dim_cond, ch_obs=1, ch_cond=1)[source]#

Return default parameters for the Flux1 model.

Parameters:

  • dim_obs (int)

  • dim_cond (int)

  • ch_obs (int)

  • ch_cond (int)

Return type:

gensbi.models.Flux1Params

gensbi.recipes.flux1.parse_flux1_params(config_path)[source]#

Parse a Flux1 configuration file.

Parameters:

config_path (str) – Path to the configuration file.

Returns:

config – Parsed configuration dictionary.

Return type:

dict