SciGenML

ConditionalFlowMatching

A conditional flow matching generative model.

The conditional flow matching generative model is a generative model that uses a flow matching approach to generate data.

It is a special case of the conditional stochastic interpolant generative model where there is no noise in the interpolant.

Abstract type for all conditional generative models.

DeterministicInterpolantCoefs

A struct that contains the alpha and beta function coefficients and their derivatives.

The interpolant is defined as:
x(t) = α(t) * x0 + β(t) * x1

The derivatives are:
dx/dt = α'(t) * x0 + β'(t) * x1

FlowMatching

A flow matching generative model.

The flow matching generative model is a generative model that uses a flow matching approach to generate data.

It is a special case of the stochastic interpolant generative model where there is no noise in the interpolant.

FollmerStochasticInterpolant

A follmer stochastic interpolant generative model.

Abstract type for all generative models.

ScoreBasedDiffusionModel

A score-based diffusion model generative model.

The score-based diffusion model generative model is a generative model that uses a score-based diffusion model approach to generate data.

It is a special case of the stochastic interpolant generative model where there is no noise in the interpolant.

StochasticInterpolant

A stochastic interpolant generative model.

StochasticInterpolantCoefs

A struct that contains the alpha and beta function coefficients and their derivatives.

The interpolant is defined as:
x(t) = α(t) * x0 + β(t) * x1 + γ(t) * Z

The derivatives are:
dx/dt = α'(t) * x0 + β'(t) * x1 + γ'(t) * Z

compute_interpolant(x0, x1, z, t, interpolant_coefs::StochasticInterpolantCoefs)

Computes the stochastic interpolant at a given time.

Args:
    x0: The starting point.
    x1: The ending point.
    z: The noise.
    t: The time.
    interpolant_coefs: The interpolant coefficients.

Computes the deterministic interpolant at a given time.

Args:
    x0: The starting point.
    x1: The ending point.
    t: The time.
    interpolant_coefs: The interpolant coefficients.

compute_interpolant_diff

Computes the derivative of the interpolant at a given time.

Args:
    x0: The starting point.
    x1: The ending point.
    interpolant_coefs: The interpolant coefficients.
    t: The time.

diffusion_interpolant_coefs(trait::Models.Deterministic)

Returns the quadratic interpolant.

The quadratic interpolant is defined as:
α(t) = exp(-multiplier * t)
β(t) = sqrt(1 - exp(-2 * multiplier * t))

The derivatives are:
α'(t) = -multiplier * exp(-multiplier * t)
β'(t) = multiplier * exp(-multiplier * t) / sqrt(1 - exp(-2 * multiplier * t))

drift_term(
    model::FollmerGenerativeModel, 
)

Compute the drift term for a stochsatic interpolant.

drift_term(
    ::Models.Stochastic,
    model::Models.ConditionalFlowMatching,
)

drift_term(
    ::Models.Deterministic,
    model::Models.ConditionalFlowMatching,
)

Get the drift term for the conditional flow matching generative model.

drift_term(
    model::FollmerGenerativeModel, 
    diffusion_fn
)

Compute the drift term for a stochsatic interpolant.

drift_term(
    ::Models.Stochastic,
    model::StochasticInterpolant, 
    diffusion_fn, 
)

Compute the drift term for a stochsatic interpolant.

drift_term(
    ::Models.Stochastic,
    model::StochasticInterpolant, 
    diffusion_fn, 
)

Compute the drift term for a stochsatic interpolant.

drift_term(
    ::Models.Stochastic,
    model::StochasticInterpolant, 
    diffusion_fn, 
    x, 
    ps, 
    st
)

Compute the drift term for a stochastic interpolant.

get_interpolant_coefs(type::String)

Get the interpolant coefficients for a given type. 
Supported types are:
- "linear"
- "quadratic"

linear_interpolant_coefs(trait::Models.Deterministic)

Returns the linear interpolant.

The linear interpolant is defined as:
α(t) = 1 - t
β(t) = t
γ(t) = 1 - t, if trait == Models.Stochastic

The derivatives are:
α'(t) = -1
β'(t) = 1
γ'(t) = -1, if trait == Models.Stochastic

posterior_drift_term(
    model::FollmerGenerativeModel, 
    diffusion_fn
)

Compute the drift term for a stochsatic interpolant.

quadratic_interpolant_coefs(trait::Models.Deterministic)

Returns the quadratic interpolant.

The quadratic interpolant is defined as:
α(t) = 1 - t
β(t) = t^2
γ(t) = 1 - t, if trait == Models.Stochastic

The derivatives are:
α'(t) = -1
β'(t) = 2t
γ'(t) = -1, if trait == Models.Stochastic

MinMaxScaler

A min-max scaler for data.

Preprocessor

A preprocessor for data.

StandardScaler

A scaler for data.

get_dims_to_reduce(data)

Get the dimensions to reduce from the data. Outputs a tuple of dimensions to reduce.

min_max_scaler_inverse_transform(
    data,
    min_val::Float32,
    max_val::Float32
)

Inverse transform for the min-max scaler.

min_max_scaler_transform(
    data,
    min_val::Float32,
    max_val::Float32
)

A scaler for data.

reshape_scalar(scalar, new_dims)

Reshape a scalar to a new dimension.

standard_scaler_inverse_transform(
    data,
    mean,
    std
)

Inverse transform for the standard scaler.

standard_scaler_transform(
    data,
    mean,
    std
)

A scaler for data.

DenseNN(
    in_features::Tuple,
    out_features::Int, 
    hidden_features::Vector{Int};
    activation_function = DEFAULT_ACTIVATION_FUNCTION,
    dropout = DEFAULT_DROPOUT
)

in_features: A tuple of the number of input features
out_features: The number of output features.
hidden_features: The number of hidden features.
activation_function: The activation function.

A dense neural network with `in_features` input features, 
`out_features` output features, and `hidden_features` hidden features.
The input is a tuple that are concatenated to form a single input.

UNet(
    in_channels::Int,
    out_channels::Int,
    hidden_channels,
    time_embedding_dim::Int
    padding::String,
)

A UNet model.

# Arguments
- `in_channels::Int`: The number of input channels.
- `out_channels::Int`: The number of output channels.
- `hidden_channels::List{Int}`: The number of hidden channels in each layer.
- `time_embedding_dim::Int`: The dimension of the time embedding.
- `scalar_in_conditioning_dim::Int`: The dimension of the scalar input conditioning.
- `scalar_hidden_conditioning_dim::Int`: The dimension of the hidden scalar conditioning.
- `padding::String`: The padding type.

UNet(
    in_channels::Int,
    out_channels::Int,
    hidden_channels,
    time_embedding_dim::Int
    padding::String,
)

A UNet model.

# Arguments
- `in_channels::Int`: The number of input channels.
- `out_channels::Int`: The number of output channels.
- `hidden_channels::List{Int}`: The number of hidden channels in each layer.
- `time_embedding_dim::Int`: The dimension of the time embedding.
- `padding::String`: The padding type.
- `field_in_conditioning_dim::Int`: The dimension of the field input conditioning.
- `field_hidden_conditioning_dim::Int`: The dimension of the hidden field conditioning.
- `field_conditioning_combination::String`: The combination method for the field conditioning.

UNet(
    in_channels::Int,
    out_channels::Int,
    hidden_channels,
    time_embedding_dim::Int
    padding::String,
)

A UNet model.

# Arguments
- `in_channels::Int`: The number of input channels.
- `out_channels::Int`: The number of output channels.
- `hidden_channels::List{Int}`: The number of hidden channels in each layer.
- `time_embedding_dim::Int`: The dimension of the time embedding.
- `padding::String`: The padding type.

UNet(config::Config.UNetHyperparameters)

A UNet model.

# Arguments
- `config::Config.UNetHyperparameters`: The configuration for the UNet model.

get_model(
    in_features,
    out_features,
    hidden_features,
    activation_function = DEFAULT_ACTIVATION_FUNCTION,
    dropout = DEFAULT_DROPOUT
)

in_features: The number of input features.
out_features: The number of output features.
hidden_features: The number of hidden features.
activation_function: The activation function.
dropout: The dropout rate.

Returns a model that can be used to predict the output from the input.

conv_next_block(
    in_channels::Int, 
    out_channels::Int,
    multiplier::Int = 1,
    pars_embed_dim::Int = 1,
    imsize::Tuple{Int, Int} = (64, 128)
)

Create a conv next block with the given number of input and output channels. The block consists of two convolutional layers with kernel size kernel_size. The first layer has the same number of input and output channels, while the second layer has the same number of output channels as the block. The block also includes batch normalization and a skip connection.

https://arxiv.org/abs/2201.03545

Based on https://github.com/tum-pbs/autoreg-pde-diffusion/blob/b9b33913b99ede88d9452c5ab470c5d7f5da5c56/src/turbpred/modeldiffusionblocks.py#L60

conv_next_block_no_conditioning(
    in_channels::Int, 
    out_channels::Int,
    multiplier::Int = 1,
    imsize::Tuple{Int, Int} = (64, 128)
)

Create a conv next block with the given number of input and output channels. The block consists of two convolutional layers with kernel size kernel_size. The first layer has the same number of input and output channels, while the second layer has the same number of output channels as the block. The block also includes batch normalization and a skip connection.

https://arxiv.org/abs/2201.03545

Based on https://github.com/tum-pbs/autoreg-pde-diffusion/blob/b9b33913b99ede88d9452c5ab470c5d7f5da5c56/src/turbpred/modeldiffusionblocks.py#L60

get_padding(padding::String, padding_size::Int)

Get the padding function for the given padding type and size.

multiple_conv_next_blocks(
    in_channels::Int, 
    out_channels::Int,
    multiplier::Int = 1,
    conditioning_dim::Int = 1,
    imsize::Tuple{Int, Int} = (64, 128)
)

Create a chain of two conv next blocks with the given number of input and output channels. The first block has the same number of input and output

multiple_conv_next_blocks_no_conditioning(
    in_channels::Int, 
    out_channels::Int,
    multiplier::Int = 1,
    padding="constant"
)

Create a chain of two conv next blocks with the given number of input and output channels. The first block has the same number of input and output

sinusoidal_embedding(
    x::AbstractArray{AbstractFloat, 4},
    min_freq::AbstractFloat,
    max_freq::AbstractFloat,
    embedding_dims::Int
)

Embed the noise variances to a sinusoidal embedding with the given frequency range and embedding dimensions.

Based on https://yng87.page/en/blog/2022/lux-ddim/.

RK4_step(rhs, x, t, dt, ps, st)

Perform a 4th order Runge-Kutta step.

RK4_step(rhs, x, conditioning, t, dt, ps, st)

Perform a 4th order Runge-Kutta step.

Euler Maruyama step for SDEs.

Euler Maruyama step for SDEs.

forward_euler_step(model, x, t, dt)

Perform a forward Euler step.

forward_euler_step(model, x, scalar_conditioning, t, dt)

Perform a forward Euler step.

Heun step for SDEs.

Heun step for SDEs.

ode_integrator(
    stepper, 
    rhs, 
    x, 
    num_steps, 
    ps, 
    st;
    t_interval=[0.0, 1.0], 
    verbose::Bool = true
)

Generic ODE integrator.

ode_integrator(
    stepper, 
    rhs, 
    x, 
    scalar_conditioning,
    num_steps, 
    ps, 
    st;
    t_interval=[0.0, 1.0], 
    verbose::Bool = true
)

Generic ODE integrator.

load_knmi_data(
    num_steps
)

Loads the KNMI data for the given number of steps.

load_knmi_data(
    config::Config.KNMIDataHyperparameters;
    kwargs...
)

Loads the KNMI data for the given configuration.

load_kolmogorov_data(;
    num_trajectories = 5,
    num_skip_steps = 5,
    num_steps = 100,
    start_step = 500
)
Load the Kolmogorov data.

Outputs a named tuple with the following fields:
- `base`: The base data. [H, W, C, (T-1)*N]
- `target`: The target data. [H, W, C, (T-1)*N]
- `field_conditioning`: The field conditioning data. [H, W, C, (T-1)*N]

load_super_res_kolmogorov_data(;
    num_trajectories = 5,
    num_skip_steps = 5,
    num_steps = 100,
    start_step = 500
)
Load the super-resolution Kolmogorov data.

Outputs a named tuple with the following fields:
- `base`: The base data. [H, W, C, (T-1)*N]
- `target`: The target data. [H, W, C, (T-1)*N]
- `field_conditioning`: The field conditioning data. [H, W, C, (T-1)*N]

sample(
    model,
    args...; 
    kwargs...
)

Generic training function for all models.

sample(
    ::Models.Deterministic,
    model::Models.FlowMatching,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a stochastic interpolant generative model using the forward Euler method.

sample(
    ::Models.Deterministic,
    model::Models.FlowMatching,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a stochastic interpolant generative model using the forward Euler method.

sample(
    ::Models.Deterministic,
    model::Models.ScoreBasedDiffusionModel,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a denoising diffusion model using the forward Euler method.

sample(
    ::Models.Deterministic,
    model::Models.StochasticInterpolantGenerativeModel,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a stochastic interpolant generative model using the forward Euler method.

sample(
    ::Models.Stochastic,
    model::Models.FollmerStochasticInterpolant,
    scalar_conditioning,
    num_steps::Int,
    num_samples::Int,
    prior_samples = nothing,
    diffusion_fn = nothing,
    rng::Random.AbstractRNG = Random.default_rng(),
    verbose::Bool = true,
    stepper = TimeIntegrators.heun_step
)

Sample from a stochastic interpolant generative model using the forward Euler method.

sample(
    ::Models.Stochastic,
    model::Models.ScoreBasedDiffusionModel,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a denoising diffusion model using the forward Euler method.

sample(
    ::Models.Stochastic,
    model::Models.ScoreBasedDiffusionModel,
    num_samples::Int,
    num_steps::Int,
    num_samples::Int,
    prior_samples,
    diffusion_fn,
    rng::Random.AbstractRNG = Random.default_rng(),
    verbose::Bool = true,
    stepper = TimeIntegrators.heun_step
)

Sample from a denoising diffusion model using the forward Euler method.

sample(
    ::Models.Stochastic,
    model::Models.StochasticInterpolant,
    scalar_conditioning,
    num_steps::Int,
    num_samples::Int,
    prior_samples = nothing,
    diffusion_fn = nothing,
    rng::Random.AbstractRNG = Random.default_rng(),
    verbose::Bool = true,
    stepper = TimeIntegrators.heun_step
)

Sample from a stochastic interpolant generative model using the forward Euler method.

sample(
    ::Models.Stochastic,
    model::Models.StochasticInterpolantGenerativeModel,
    num_samples::Int,
    num_steps::Int,
    rng::Random.AbstractRNG = Random.default_rng()
)

Sample from a stochastic interpolant generative model using the forward Euler method.

Checkpoint

A checkpoint for a model and other data.

_train_step(
    ::Models.Deterministic,
    model::Models.ConditionalFlowMatching,
    base_samples,
    target_samples,
    scalar_conditioning_samples,
    train_state,
    rng
)

Train a stochastic interpolant generative model for one step.

_train_step(
    ::Models.Deterministic,
    model::Models.FlowMatching,
    base_samples,
    target_samples,
    train_state,
    rng
)

Train a flow matching generative model for one step.

_train_step(
    ::Models.Deterministic,
    model::Models.FlowMatching,
    base_samples,
    target_samples,
    field_conditioning,
    train_state,
    rng
)

Train a flow matching generative model for one step.

_train_step(
    ::Models.Deterministic,
    model::Models.StochasticInterpolant,
    base_samples,
    target_samples,
    train_state,
    rng
)

Train a stochastic interpolant generative model for one step.

_train_step(
    ::Models.Deterministic,
    model::Models.ScoreBasedDiffusionModel,
    base_samples,
    target_samples,
    train_state,
    rng
)

Train a score-based diffusion generative model for one step via flow matching.

_train_step(
    ::Models.Deterministic,
    model::Models.ScoreBasedDiffusionModel,
    base_samples,
    target_samples,
    conditioning,
    train_state,
    rng
)

Train a score-based diffusion generative model for one step via flow matching.

_train_step(
    ::Models.Stochastic,
    model::Models.FollmerStochasticInterpolant,
    base_samples,
    target_samples,
    field_conditioning,
    train_state,
    rng
)

Train a stochastic interpolant generative model for one step.

_train_step(
    ::Models.Stochastic,
    model::Models.StochasticInterpolant,
    base_samples,
    target_samples,
    train_state,
    rng
)

Train a stochastic interpolant generative model for one step.

_train_step(
    ::Models.Stochastic,
    model::Models.StochasticInterpolant,
    base_samples,
    target_samples,
    scalar_conditioning,
    train_state,
    rng
)

Train a stochastic interpolant generative model for one step.

compute_score_loss(model, ps, st, (x, z, gamma))

Compute the loss for a stochastic interpolant generative model.

compute_velocity_loss(model, ps, st, (x, y))

Compute the loss for a stochastic interpolant generative model.

get_dataloader(
    base, 
    target, 
    batch_size::Int, 
    match_base_and_target::Bool=false
)

Get a dataloader from a base and target.

get_dataloader(
    data::NamedTuple,
    batch_size::Int,
    match_base_and_target::Bool=false
)

Get a dataloader from a named tuple.

get_dataloader(
    base,
    target,
    scalar_conditioning,
    batch_size::Int,
    match_base_and_target::Bool=false
)

Get a dataloader from a base, target, and conditioning.

get_gradients(
    batch,
    train_state,
    loss_fn,
)

Get gradients from a batch, train state, and loss function.

get_interpolated_samples(
    base_samples,
    target_samples,
    t_samples,
    interpolant_coefs
)

Get interpolated samples from a base, target, and time.

get_interpolated_samples(
    base_samples,
    target_samples,
    z_samples,
    t_samples,
    interpolant_coefs

)

Get interpolated samples from a base, target, and time.

get_optimizer(type::String, lr, lambda)

Get an optimizer from a string.

load_train_state(
    checkpoint::Checkpoint
)

Load a train state from a checkpoint.

save_train_state(
    ps,
    st,
    checkpoint::Checkpoint
)

Save a train state to a checkpoint.

split_data(data, split_ratio, rng = Random.default_rng())

Split data into train and validation sets.

train(
    ::Models.Deterministic,
    model::Models.ScoreBasedDiffusionModel,
    data,
    config,
    rng = Random.default_rng();
    verbose = true,
    checkpoint = nothing
)

Train a score-based diffusion generative model vis flow matching.

train(
    ::Models.Deterministic,
    model::Models.FlowMatching,
    data,
    config,
    rng = Random.default_rng();
    verbose = true,
    checkpoint = nothing
)

Train a stochastic interpolant generative model.

train(
    ::Models.Stochastic,
    model::Models.FollmerStochasticInterpolant,
    data,
    config,
    rng = Random.default_rng();
    verbose = true,
    checkpoint = nothing
)

Train a Follmer stochastic interpolant generative model.

train(
    ::Models.Stochastic,
    model::Models.StochasticInterpolantGenerativeModel,
    data,
    config,
    rng = Random.default_rng();
    verbose = true
)

Train a stochastic interpolant generative model.

train(
    ::Models.Deterministic,
    model::Models.StochasticInterpolantGenerativeModel,
    data,
    config,
    rng = Random.default_rng();
    verbose = true,
    checkpoint = nothing
)

Train a stochastic interpolant generative model.

train(
    ::Models.Deterministic,
    model::Models.ConditionalFlowMatching,
    data,
    config,
    rng = Random.default_rng();
    verbose = true
)

Train a stochastic interpolant generative model.

train(
    model,
    args...; 
    kwargs...
)

Generic training function for all models.

ConditionalFlowMatchingHyperparameters

Hyperparameters for the conditional flow matching generative model.

DenseNeuralNetworkHyperparameters

Hyperparameters for the dense neural network.

FlowMatchingHyperparameters

Hyperparameters for the flow matching generative model.

FollmerStochasticInterpolantHyperparameters

Hyperparameters for the follmer stochastic interpolant generative model.

Hyperparameters

Hyperparameters for the architecture, training, and optimizer.

KNMIDataHyperparameters

Hyperparameters for the KNMI data.

KolmogorovDataHyperparameters

Hyperparameters for the kolmogorov data.

OptimizerHyperparameters

Hyperparameters for the optimizer.

PlaceholderDataHyperparameters

Hyperparameters for the placeholder.

ScoreBasedDiffusionModelHyperparameters

Hyperparameters for the score-based diffusion model generative model.

StochasticInterpolantHyperparameters

Hyperparameters for the stochastic interpolant generative model.

SuperResKolmogorovDataHyperparameters

Hyperparameters for the super res kolmogorov data.

TrainingHyperparameters

Hyperparameters for the training.

UNetHyperparameters

Hyperparameters for the UNet.

animate_field(field_list, filename, plot_titles)

Animate a list of fields.

animate_velocity_field(field_list, filename, plot_titles; velocity_channels = (1, 2))

Animate a list of velocity fields.