EasyHybrid.jl

Documentation for EasyHybrid.jl.

EasyHybrid.EasyHybrid Module

julia

EasyHybrid.jl

EasyHybrid is a Julia package for hybrid machine learning models, combining neural networks and traditional statistical methods. It provides tools for data preprocessing, model training, and evaluation, making it easier to build and deploy hybrid models.

The hybrid model combines a neural network h(x; θ), with inputs x and learnable parameters θ, together with a mechanistic model M(·, z; ϕ) driven by forcings z and parameterized by ϕ, where ϕ may be known, learned from data, or fixed.

Documentation https://earthyscience.github.io/EasyHybrid.jl/

EasyHybrid.DataConfig Type

Configuration for data preparation and loading.

Controls array types, observation shuffling, data splitting, cross-validation, and sequence construction for time-series training.

EasyHybrid.HybridModel Type

julia

HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)}

A unified hybrid model struct that handles both single and multi neural network architectures. It combines predictive neural networks (NNs) with a mechanistic_model to form a differentiable hybrid model.

EasyHybrid.HybridModel Method

julia

(m::HybridModel)(ds_k::Tuple, ps, st)

Forward pass of the hybrid model. Evaluates the neural networks to predict parameters, merges them with scaled global parameters and fixed parameters, and executes the mechanistic model. Returns a tuple (out, st_new).

EasyHybrid.InputBatchNorm Type

julia

InputBatchNorm(chs; kwargs...)

A wrapper around BatchNorm that handles 3D sequence input (features, timesteps, batch).

Lux's BatchNorm expects channels in the penultimate dimension, which works for 2D input (features, batch) but fails for 3D input (features, timesteps, batch) where features are in dim 1. This wrapper reshapes 3D input to 2D (features, timesteps * batch) before normalization, then reshapes back. For 2D input, delegates directly to BatchNorm.

EasyHybrid.LoggingLoss Type

julia

LoggingLoss

A structure to define a logging loss function for hybrid models.

Arguments

loss_types: A vector of loss specifications (Symbol, Function or Tuple)
- Symbol: predefined loss, e.g. :mse
- Function: custom loss function, e.g. custom_loss
- Tuple: function with args/kwargs:
  - (f, args): positional args, e.g. (weighted_loss, (0.5,))
  - (f, kwargs): keyword args, e.g. (scaled_loss, (scale=2.0,))
  - (f, args, kwargs): both, e.g. (complex_loss, (0.5,), (scale=2.0,))
training_loss: The loss specification to use during training (same format as above)
extra_loss: Optional function (ŷ, ps; kwargs...) -> NamedTuple (or splattable collection) added to training loss (default: nothing)
agg: Function to aggregate losses across targets, e.g. sum or mean
train_mode: If true, uses training_loss; otherwise uses loss_types.

Examples

julia

# Simple predefined loss
logging = LoggingLoss(
    loss_types=[:mse, :mae],
    training_loss=:mse
)

# Custom loss function
custom_loss(ŷ, y) = mean(abs2, ŷ .- y)
logging = LoggingLoss(
    loss_types=[:mse, custom_loss],
    training_loss=custom_loss
)

# With arguments/kwargs
weighted_loss(ŷ, y, w) = w * mean(abs2, ŷ .- y)
scaled_loss(ŷ, y; scale=1.0) = scale * mean(abs2, ŷ .- y)
logging = LoggingLoss(
    loss_types=[:mse, (weighted_loss, (0.5,)), (scaled_loss, (scale=2.0,))],
    training_loss=(weighted_loss, (0.5,))
)

EasyHybrid.ParameterContainer Type

julia

ParameterContainer{NT <: NamedTuple, T}

A container for holding the parameter definitions of a model, including their default values, lower bounds, and upper bounds.

EasyHybrid.PerTarget Type

julia

PerTarget(losses)

A wrapper to indicate that a tuple of losses should be applied on a per-target basis.

EasyHybrid.RecurrenceOutputDense Type

julia

RecurrenceOutputDense(in_dims => out_dims, [activation])

A layer that wraps a Dense layer to handle sequence outputs from Recurrence layers.

When a Recurrence layer has return_sequence=true, it outputs a tuple/vector of arrays (one per timestep). This layer broadcasts the Dense operation over each timestep and reshapes the result to (features, timesteps, batch) format.

Arguments

in_dims::Int: Input dimension (should match Recurrence output dimension)
out_dims::Int: Output dimension
activation: Activation function (default: identity)

Example

julia

# Instead of manually creating:
broadcast_layer = @compact(; layer = Dense(15 => 15)) do x
    y = map(layer, x)
    @return permutedims(stack(y; dims = 3), (1, 3, 2))
end

# Simply use:
Chain(
    Recurrence(LSTMCell(15 => 15), return_sequence = true),
    RecurrenceOutputDense(15 => 15)
)

EasyHybrid.TrainConfig Type

Configuration for training a hybrid model.

Controls all aspects of the training process including optimization, loss computation, data handling, output, and visualization.

EasyHybrid.TrainResults Type

Output of train, containing the full training history, model state, and diagnostics.

EasyHybrid.TrainingPaths Type

Paths to all output files produced during a training run.

EasyHybrid.WrappedTuples Type

julia

WrappedTuples(vec::Vector{<:NamedTuple})

Wraps a vector of named tuples to allow dot-access to each field as a vector.

EasyHybrid._apply_loss Function

julia

_apply_loss(ŷ, y, y_nan, loss_spec)

Helper function to apply the appropriate loss function based on the specification type.

Arguments

ŷ: Predictions for a single target
y: Target values for a single target
y_nan: NaN mask for a single target
loss_spec: Loss specification (Symbol, Function, or Tuple)

Returns

Computed loss value

EasyHybrid._build_optim_loss Method

Build the scalar loss closure consumed by Optimization.jl, called as loss_fn(p, data). data is an already device-placed / Array-converted batch (shape ((x, forcings), (y, mask))), produced once per (mini)batch by collect_dim_data in the caller — not inside this closure. Keeping the data prep out of the loss is important: L-BFGS line searches call the objective many times per iteration, so re-running collect_dim_data (NamedTuple rebuilds, Array copies, gdev transfers) on every evaluation was a major slowdown, especially on the minibatch path. It also keeps the closure trivially Zygote-differentiable (no pairs(...)/∇map in the AD tape).

EasyHybrid._compute_loss Function

julia

_compute_loss(ŷ, y, y_nan, targets, loss_spec, agg::Function)
_compute_loss(ŷ, y, y_nan, targets, loss_types::Vector, agg::Function)

Compute the loss for the given predictions and targets using the specified training loss (or vector of losses) type and aggregation function.

Arguments:

ŷ: Predicted values.
y: Target values.
y_nan: Mask for NaN values.
targets: The targets for which the loss is computed.
loss_spec: The loss type to use during training, e.g., :mse.
loss_types::Vector: A vector of loss types to compute, e.g., [:mse, :mae].
agg::Function: The aggregation function to apply to the computed losses, e.g., sum or mean.

Returns a single loss value if loss_spec is provided, or a NamedTuple of losses for each type in loss_types.

EasyHybrid._get_target_nan Function

julia

_get_target_nan(y_nan, target)

Helper function to extract target-specific values from y_nan. Supports NamedTuple, KeyedArray, AbstractDimArray, and callables (functions).

EasyHybrid._get_target_y Function

julia

_get_target_y(y, target)

Helper function to extract target-specific values from y, handling cases where y can be a tuple of (y_obs, y_sigma). Supports NamedTuple, KeyedArray, AbstractDimArray, Tuple, and callables (functions).

EasyHybrid._init_nn_params Method

julia

_init_nn_params(rng, m::HybridModel{<:Any, <:NamedTuple})

Initialize parameters for a multi-neural network architecture. Returns a NamedTuple containing the initialized parameters for each sub-network.

EasyHybrid._init_nn_params Method

julia

_init_nn_params(rng, m::HybridModel{<:Any, <:Vector})

Initialize parameters for a single-neural network architecture. Returns a NamedTuple containing a single ps field with the network's parameters.

EasyHybrid._init_nn_states Method

julia

_init_nn_states(rng, m::HybridModel{<:Any, <:NamedTuple})

Initialize states for a multi-neural network architecture. Returns a NamedTuple containing the initialized states for each sub-network.

EasyHybrid._init_nn_states Method

julia

_init_nn_states(rng, m::HybridModel{<:Any, <:Vector})

Initialize states for a single-neural network architecture. Returns a NamedTuple containing a single st_nn field with the network's states.

EasyHybrid._is_named_leaf Method

julia

_is_named_leaf(name, child, key) -> Bool

Return true when child is an array leaf whose parent field is key (e.g. a Dense layer's weight matrix).

EasyHybrid._run_minibatch! Method

Explicit "repeated minibatch" driver for the full_batch = false Optimization path (Le et al., 2011, ICML, §4.2). For each of cfg.nepochs outer passes we iterate a reshuffled DataLoader; on every fixed minibatch we run cfg.inner_maxiters optimizer iterations warm-started from the current ps, then resample. A validation EpochSnapshot (history / early-stopping / dashboard / checkpoint) is built once per outer pass, so patience is counted in outer passes — consistent with the Optimisers.jl loop.

maxiters / epochs from solve_kwargs are dropped (the per-minibatch budget is cfg.inner_maxiters and the pass count is cfg.nepochs); any remaining solve_kwargs (e.g. g_abstol, f_reltol) are forwarded to each inner solve.

EasyHybrid._run_nn Method

julia

_run_nn(m::HybridModel{<:Any, <:NamedTuple}, ds_k::Tuple, ps, st)

Execute the forward pass for a multi-neural network architecture. Applies each sub-network to its specific predictors, and applies scaling to the outputs if required. Returns scaled parameter values, updated states, and raw network outputs.

EasyHybrid._run_nn Method

julia

_run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st)

Execute the forward pass for a single-neural network architecture. Applies the neural network to the given predictors, slices the output for multiple predicted parameters, and scales them if required. Returns scaled parameter values, updated states, and raw network outputs.

EasyHybrid._train Method

julia

_train(model, data, train_cfg, data_cfg, solve_kwargs)

Dispatcher used by train(...): routes to the original 4-arg _train body (the Lux.Training / Optimisers.jl loop) when train_cfg.opt isa Optimisers.AbstractRule, or to _train_optimization (which delegates batch iteration to Optimization.jl) otherwise. solve_kwargs are forwarded to solve(...) on the Optimization.jl branch and warned about on the Optimisers.jl branch.

EasyHybrid._train_optimization Method

julia

_train_optimization(model, data, train_cfg, data_cfg, solve_kwargs)

Optimization.jl-based driver dispatched from _train whenever train_cfg.opt is not an Optimisers.AbstractRule (e.g. Optim.LBFGS() / Optimization.LBFGS()); see the SciML minibatching tutorial.

Two modes, selected by train_cfg.full_batch:

full_batch = true: pass the full training set as a single tuple to one OptimizationProblem and a single solve(...). Batch-method idiom (the recommended L-BFGS setup): the objective is a single consistent function. solve_kwargs (e.g. maxiters, g_abstol, f_reltol) are splatted into solve, and train_cfg.eval_every builds a validation EpochSnapshot every N solver iterations via the callback.
full_batch = false: explicit "repeated minibatch" loop grounded in Le et al., 2011 (On Optimization Methods for Deep Learning, ICML, §4.2). For each of train_cfg.nepochs outer passes we iterate a (reshuffled) DataLoader and, on each fixed minibatch, run train_cfg.inner_maxiters optimizer iterations (solve(...; maxiters = inner_maxiters)), warm-starting the next minibatch from the current parameters. Holding the minibatch fixed for a few iterations keeps the objective (and L-BFGS curvature pairs / line search) consistent — naive one-step-per-minibatch L-BFGS does not converge. A validation EpochSnapshot is built once per outer pass. Optimization.jl's own DataLoader iteration is not used here because it only applies to the Optimisers.jl-style solvers, not Optim.jl's L-BFGS.

Both modes honour train_cfg.promote_f64: promote ps to Float64 before optimization (workaround for Lux.jl#1260).

EasyHybrid.build_opt_state Method

julia

build_opt_state(opt, ps::NamedTuple; default_rule = Optimisers.Adam())

Build the optimizer state tree consumed by Lux.Training.TrainState / Optimisers.update!. Three forms of opt are accepted:

opt::Optimisers.AbstractRule — single rule applied to the whole parameter tree (delegates to Optimisers.setup(opt, ps)).
opt::NamedTuple of Optimisers.AbstractRules — each rule is wired to the matching top-level branch of ps via Optimisers.setup(rule, ps[name]). Branches missing from opt use default_rule.
opt::NamedTuple of pre-built state trees (already returned by a prior Optimisers.setup) — used as-is. Form 2 and 3 can be mixed in the same NamedTuple.

The returned state tree has the same top-level keys as ps.

Example

julia

ps, _ = LuxCore.setup(rng, hybrid_model)   # (; Rb = NN_ps, RUE = NN_ps, Q10 = [v])

# Form 2 — preferred, lets the framework call `Optimisers.setup`:
opt_state = build_opt_state(
    (; Rb = Adam(1e-3), RUE = Adam(1e-3), Q10 = Descent(1e-2)),
    ps,
)

EasyHybrid.build_parameter_matrix Method

julia

build_parameter_matrix(parameter_defaults_and_bounds::NamedTuple)

Build a ComponentArray matrix from a NamedTuple containing parameter defaults and bounds.

This function converts a NamedTuple where each value is a tuple of (default, lower, upper) bounds into a ComponentArray with named axes for easy parameter management in hybrid models.

Arguments

parameter_defaults_and_bounds::NamedTuple: A NamedTuple where each key is a parameter name and each value is a tuple of (default, lower, upper) for that parameter.

Returns

ComponentArray: A 2D ComponentArray with:
- Row axis: Parameter names (from the NamedTuple keys)
- Column axis: Bound types (:default, :lower, :upper)
- Data: The parameter values organized in a matrix format

Example

julia

# Define parameter defaults and bounds
parameter_defaults_and_bounds = (
    θ_s = (0.464f0, 0.302f0, 0.700f0),     # Saturated water content [cm³/cm³]
    h_r = (1500.0f0, 1500.0f0, 1500.0f0),  # Pressure head at residual water content [cm]
    α   = (log(0.103f0), log(0.01f0), log(7.874f0)),  # Shape parameter [cm⁻¹]
    n   = (log(3.163f0 - 1), log(1.100f0 - 1), log(20.000f0 - 1)),  # Shape parameter [-]
)

# Build the ComponentArray
parameter_matrix = build_parameter_matrix(parameter_defaults_and_bounds)

# Access specific parameter bounds
parameter_matrix.θ_s.default  # Get default value for θ_s
parameter_matrix[:, :lower]   # Get all lower bounds
parameter_matrix[:, :upper]   # Get all upper bounds

Notes

The function expects each value in the NamedTuple to be a tuple with exactly 3 elements
The order of bounds is always (default, lower, upper)
The resulting ComponentArray can be used for parameter optimization and constraint handling

EasyHybrid.compute_loss Method

julia

compute_loss(HM, x, (y_t, y_nan), ps, st, logging::LoggingLoss)

Main loss function for hybrid models that handles both training and evaluation modes.

Arguments

HM: The hybrid model (AbstractLuxContainerLayer or specific model type)
x: Input data for the model
(y_t, y_nan): Tuple containing target values and NaN mask functions/arrays
ps: Model parameters
st: Model state
logging: LoggingLoss configuration

Returns

In training mode (logging.train_mode = true):
- (loss_value, st): Single loss value and updated state
In evaluation mode (logging.train_mode = false):
- (loss_values, st, ŷ): NamedTuple of losses, state and predictions

EasyHybrid.constructHybridModel Method

julia

constructHybridModel(predictors::NamedTuple, forcing, targets, mechanistic_model, parameters, global_param_names; kwargs...)

Construct a HybridModel with multiple neural network architectures. A separate neural network is built for each key in the predictors NamedTuple.

Arguments:

predictors::NamedTuple: A NamedTuple where keys are network names, and values are vectors of predictor variables for that network.
forcing: Variables passed directly to the mechanistic model.
targets: The target variables to predict.
mechanistic_model: A function implementing the process-based model.
parameters: A parameter container defining defaults, lowers, and uppers.
global_param_names: Names of the parameters to be globally optimized.
kwargs: Additional configuration. hidden_layers and activation can also be NamedTuples to configure each network independently.

EasyHybrid.constructHybridModel Method

julia

constructHybridModel(predictors::Vector{Symbol}, forcing, targets, mechanistic_model, parameters, neural_param_names, global_param_names; kwargs...)

Construct a HybridModel with a single neural network architecture predicting all neural_param_names from the predictors.

Arguments:

predictors::Vector{Symbol}: Variables used as inputs to the neural network.
forcing: Variables passed directly to the mechanistic model.
targets: The target variables to predict.
mechanistic_model: A function implementing the process-based model.
parameters: A parameter container defining defaults, lowers, and uppers.
neural_param_names: Names of the parameters to be predicted by the neural network.
global_param_names: Names of the parameters to be globally optimized.
kwargs: Additional configuration like hidden_layers, activation, scale_nn_outputs, etc.

EasyHybrid.constructNNModel Method

julia

constructNNModel(predictors, targets; hidden_layers, activation, scale_nn_outputs)

Main constructor: hidden_layers can be either • a Vector{Int} of sizes, or • a Chain of hidden-layer Dense blocks.

EasyHybrid.evec Method

evec(nt::NamedTuple)

EasyHybrid.extract_weights Method

julia

extract_weights(ps; key=:weight) -> Vector{AbstractArray}

Walk the parameter tree ps (a ComponentArray, NamedTuple, or any nested combination of them) and return all leaf arrays whose immediate parent field name equals key.

Defaults to :weight, so you get the weight matrices of Dense/Conv layers and skip biases, BatchNorm scale/bias, running statistics in st, and any scalar global parameters.

The returned arrays are views/aliases into ps. When ps is the argument the autodiff is differentiating w.r.t., gradients of any function of these views flow back into the trainable weights.

EasyHybrid.filter_sequences Method

julia

filter_sequences(x, y) -> (x_filtered, y_filtered)

Drop 3rd-dim samples where any predictor is NaN or all targets are NaN. Accepts x as a 3D array or (x_array, forcings) Tuple.

EasyHybrid.get_layer_dim Method

julia

get_layer_dim(layer, type::Symbol)

Helper function to extract the input or output dimensions of a Lux layer. This function uses multiple dispatch to safely pull dimensions from various container types.

EasyHybrid.get_loss_value Method

julia

get_loss_value(losses, loss_spec, agg)

Extract loss value from losses based on the loss specification type.

Arguments

losses: NamedTuple containing loss values
loss_spec: Loss specification (Symbol, Function or Tuple)
agg: Aggregation function name as Symbol

Returns

Loss value for the specified loss function

Examples

julia

# Symbol case
val = get_loss_value(losses, :mse, :sum)

# Function case
custom_loss(ŷ, y) = mean(abs2, ŷ .- y)
val = get_loss_value(losses, custom_loss, :mean)

# Tuple case
weighted_loss(ŷ, y, w) = w * mean(abs2, ŷ .- y)
val = get_loss_value(losses, (weighted_loss, (0.5,)), :sum)

EasyHybrid.get_mechanistic_model_config Method

julia

get_mechanistic_model_config(f::Function)

Build an OrderedDict describing a function for YAML output: its name, and for each Method, its source file, starting line, and the full source text extracted from disk by parsing one complete expression starting at that line. The single-method case is flattened so the YAML stays compact.

Used to record mechanistic_model in the saved config so the exact function definition (not just its name) is preserved alongside the run.

EasyHybrid.get_parameters_config Method

julia

get_parameters_config(pc::ParameterContainer)

Serialize the parameter table (default, lower, upper per parameter) of a ParameterContainer into a nested OrderedDict suitable for YAML output. Without this, only the compact show of the struct (e.g. ParameterContainer(RUE, Rb, Q10)) would be written, which drops all of the actual default and bound values.

EasyHybrid.get_prediction_target_names Method

julia

get_prediction_target_names(hm)

Utility function to extract predictor/forcing and target names from a hybrid model.

Arguments:

hm: The Hybrid Model

Returns a tuple of (predictors_forcing, targets) names.

EasyHybrid.initialize_plotting_observables Method

julia

initialize_plotting_observables(init_ŷ_train, init_ŷ_val, y_train, y_val, l_init_train, l_init_val, training_loss, agg, monitor_names, target_names)

Initialize plotting observables for training visualization if the Makie extension is loaded.

EasyHybrid.is_optimisers_rule Method

julia

is_optimisers_rule(opt) -> Bool

Return true when opt originates from the Optimisers.jl package (e.g. Adam, AdamW, RMSProp, OptimiserChain), or when opt is a NamedTuple describing a per-branch optimizer (see is_per_branch_opt).

The check on single rules is by source package (nameof(parentmodule(typeof(opt))) === :Optimisers) rather than by isa Optimisers.AbstractRule, because in some package combinations Optim.jl optimizers were observed to satisfy the AbstractRule test and get misrouted to the Lux.Training loop.

The Lux.Training-based loop dispatches on this; everything else (including Optim.jl and Optimization.jl optimizers) is routed through the Optimization.jl driver.

EasyHybrid.is_per_branch_opt Method

julia

is_per_branch_opt(opt) -> Bool

Return true when opt is a NamedTuple describing a per-branch optimizer specification — i.e. one of:

a NamedTuple of Optimisers.AbstractRules (e.g. (; Rb = Adam(1e-3), Q10 = Descent(1e-2))), or
a NamedTuple of pre-built optimizer state trees as returned by Optimisers.setup(rule, ps_branch), or
a mix of the two.

Branches of the parameter tree not listed in opt fall back to the default rule Adam() (see build_opt_state).

This is detected purely by opt isa NamedTuple; the per-branch dispatch is checked again per-leaf when build_opt_state walks the spec.

EasyHybrid.kwargs_to_configs Method

julia

kwargs_to_configs(save_ps, kwargs) -> (TrainConfig, DataConfig, NamedTuple)

Build a fresh (TrainConfig, DataConfig) pair from a flat collection of kwargs. Kwargs are split between the two configs based on fieldnames(TrainConfig) and fieldnames(DataConfig); anything left over is returned as the third element and forwarded to solve(...) on the Optimization.jl path (or warned about on the Optimisers.jl path — see _train).

save_ps is the deprecated positional argument from train(model, data, save_ps; ...); when non-empty it is forwarded as tracked_params on the resulting TrainConfig.

EasyHybrid.load_timeseries_netcdf Method

julia

load_timeseries_netcdf(path::AbstractString; timedim::AbstractString = "time") -> DataFrame

Reads a NetCDF file where all data variables are 1D over the specified timedim and returns a tidy DataFrame with one row per time step.

Only includes variables whose sole dimension is timedim.
Does not attempt to parse or convert time units; all columns are read as-is.

EasyHybrid.loss_fn Function

julia

loss_fn(ŷ, y, y_nan, loss_type)

Compute the loss for given predictions and targets using various loss specifications.

Arguments

ŷ: Predicted values
y: Target values
y_nan: Mask for NaN values
loss_type: One of the following:
- Val(:rmse): Root Mean Square Error
- Val(:mse): Mean Square Error
- Val(:mae): Mean Absolute Error
- Val(:pearson): Pearson correlation coefficient
- Val(:r2): R-squared
- Val(:pearsonLoss): 1 - Pearson correlation coefficient
- Val(:nseLoss): 1 - NSE
- ::Function: Custom loss function with signature f(ŷ, y)
- ::Tuple{Function, Tuple}: Custom loss with args f(ŷ, y, args...)
- ::Tuple{Function, NamedTuple}: Custom loss with kwargs f(ŷ, y; kwargs...)
- ::Tuple{Function, Tuple, NamedTuple}: Custom loss with both f(ŷ, y, args...; kwargs...)

Examples

julia

# Predefined loss
loss = loss_fn(ŷ, y, y_nan, Val(:mse))

# Custom loss function
custom_loss(ŷ, y) = mean(abs2, ŷ .- y)
loss = loss_fn(ŷ, y, y_nan, custom_loss)

# With positional arguments
weighted_loss(ŷ, y, w) = w * mean(abs2, ŷ .- y)
loss = loss_fn(ŷ, y, y_nan, (weighted_loss, (0.5,)))

# With keyword arguments
scaled_loss(ŷ, y; scale=1.0) = scale * mean(abs2, ŷ .- y)
loss = loss_fn(ŷ, y, y_nan, (scaled_loss, (scale=2.0,)))

# With both args and kwargs
complex_loss(ŷ, y, w; scale=1.0) = scale * w * mean(abs2, ŷ .- y)
loss = loss_fn(ŷ, y, y_nan, (complex_loss, (0.5,), (scale=2.0,)))

You can define additional predefined loss functions by adding more methods:

julia

import EasyHybrid: loss_fn
function EasyHybrid.loss_fn(ŷ, y, y_nan, ::Val{:nse})
    return 1 - sum((ŷ[y_nan] .- y[y_nan]).^2) / sum((y[y_nan] .- mean(y[y_nan])).^2)
end

EasyHybrid.make_folds Method

julia

make_folds(df::DataFrame; k::Int=5, shuffle=true) -> Vector{Int}

Assigns each observation in the DataFrame df to one of k folds for cross-validation.

Arguments

df::DataFrame: The input DataFrame whose rows are to be split into folds.
k::Int=5: Number of folds to create.
shuffle=true: Whether to shuffle the data before assigning folds.

Returns

folds::Vector{Int}: A vector of length nrow(df) where each entry is an integer in 1:k indicating the fold assignment for that observation.

EasyHybrid.override_config Method

julia

override_config(cfg, overrides::NamedTuple)

Return a new cfg of the same type with the fields named in overrides replaced. Works with any @kwdef struct (e.g. TrainConfig, DataConfig).

EasyHybrid.override_configs Method

julia

override_configs(train_cfg, data_cfg, kwargs) -> (TrainConfig, DataConfig, NamedTuple)

Return (train_cfg′, data_cfg′, solve_kwargs) where any field present in kwargs overrides the corresponding field of train_cfg/data_cfg. Fields not mentioned in kwargs are kept as-is. Anything left over is returned as solve_kwargs and forwarded to solve(...) on the Optimization.jl path (or warned about on the Optimisers.jl path — see _train).

EasyHybrid.prepare_data Function

julia

prepare_data(hm, data::DataFrame; array_type=:KeyedArray, drop_missing_rows=true)
prepare_data(hm, data::KeyedArray)
prepare_data(hm, data::AbstractDimArray)
prepare_data(hm, data::Tuple)

Prepare data for training by extracting predictor/forcing and target variables based on the hybrid model's configuration.

Arguments:

hm: The Hybrid Model
data: The input data, which can be a DataFrame, KeyedArray, or DimensionalData array.
array_type: (DataFrame only) Output array type: :KeyedArray (default) or :DimArray.
drop_missing_rows: (DataFrame only) If true (default), drop rows where any predictor is NaN or all targets are NaN.

Returns:

If data is a DataFrame: a tuple of (predictors_forcing, targets) as KeyedArrays or DimArrays depending on array_type.
If data is a KeyedArray: a tuple of (predictors_forcing, targets) as KeyedArrays.
If data is an AbstractDimArray: a tuple of (predictors_forcing, targets) as DimArrays.
If data is already a Tuple, it is returned as-is.

EasyHybrid.scale_single_param Method

julia

scale_single_param(name, raw_val, parameters)

Scale a single parameter using the sigmoid scaling function.

EasyHybrid.scale_single_param_minmax Method

julia

scale_single_param_minmax(name, hm::AbstractHybridModel)

Scale a single parameter using the minmax scaling function.

EasyHybrid.select_cols Method

select_cols(df, predictors, x)

EasyHybrid.select_cols Method

select_cols(df::KeyedArray, predictors, x)

EasyHybrid.select_predictors Method

select_predictors(df, predictors)

EasyHybrid.select_predictors Method

select_predictors(dk::KeyedArray, predictors)

EasyHybrid.select_variable Method

select_variable(df::KeyedArray, x)

EasyHybrid.split_data Function

julia

split_data(data, hybridModel; split_by_id=nothing, shuffleobs=false, split_data_at=0.8, kwargs...)
split_data(data::Union{DataFrame, KeyedArray}, hybridModel; split_by_id=nothing, shuffleobs=false, split_data_at=0.8, folds=nothing, val_fold=nothing, kwargs...)
split_data(data::AbstractDimArray, hybridModel; split_by_id=nothing, shuffleobs=false, split_data_at=0.8, kwargs...)
split_data(data::Tuple, hybridModel; split_by_id=nothing, shuffleobs=false, split_data_at=0.8, kwargs...)
split_data(data::Tuple{Tuple, Tuple}, hybridModel; kwargs...)

Split data into training and validation sets, either randomly, by grouping by ID, or using external fold assignments.

Arguments:

data: The data to split, which can be a DataFrame, KeyedArray, AbstractDimArray, or Tuple
hybridModel: The hybrid model object used for data preparation
split_by_id=nothing: Either nothing for random splitting, a Symbol for column-based splitting, or an AbstractVector for custom ID-based splitting
shuffleobs=false: Whether to shuffle observations during splitting
split_data_at=0.8: Ratio of data to use for training
folds: Vector or column name of fold assignments (1..k), one per sample/column for k-fold cross-validation
val_fold: The validation fold to use when folds is provided
sequence_kwargs=nothing: NamedTuple of keyword arguments forwarded to split_into_sequences (e.g. (; input_window=10, output_window=1, output_shift=1, lead_time=2)). When set, data is windowed into 3D sequences before splitting.

Behavior:

For DataFrame/KeyedArray: Supports random splitting, ID-based splitting, and external fold assignments
For AbstractDimArray/Tuple: Random splitting only after data preparation
For pre-split Tuple{Tuple, Tuple}: Returns input unchanged

Returns:

((x_train, y_train), (x_val, y_val)): Tuple containing training and validation data pairs

EasyHybrid.split_data Method

split_data(df::DataFrame, target, xvars, seqID; f=0.8, batchsize=32, shuffle=true, partial=true)

EasyHybrid.split_data Method

split_data(df::DataFrame, target, xvars; f=0.8, batchsize=32, shuffle=true, partial=true)

EasyHybrid.split_into_sequences Method

julia

split_into_sequences(x, y; input_window=5, output_window=1, output_shift=1, lead_time=1)

Slide a (input_window + lead_time) window over 2D (feature, time) arrays to produce 3D (feature, time, batch) tensors for sequence-to-sequence training.

Arguments:

x: 2D input array (feature, time), or a (x_array, forcings) Tuple from prepare_data.
y: 2D target array (target, time), or a NamedTuple of 1D target vectors.
input_window: number of input time steps per sample.
output_window: number of target time steps per sample.
output_shift: stride between consecutive samples.
lead_time: gap between end of input window and end of output window.

Returns:

(X, Y) as 3D arrays. When x is a Tuple, returns ((X, forcings_windowed), Y).

EasyHybrid.split_seq2seq Method

julia

split_seq2seq(x, forcings, y; enc_window, dec_window, lead_time=1, stride=1)
split_seq2seq(x, y; enc_window, dec_window, lead_time=1, stride=1)

Slides a window over N-dimensional arrays (where time is the last dimension) to generate (Encoder Input, Decoder Input, Target) tuples for Sequence-to-Sequence Modeling.

Arguments:

x: Historical covariates array (..., Features, Time).
forcings: Known future covariates array (..., Features, Time). Can be omitted if not used.
y: Target array (..., Features, Time).
enc_window: Length of the historical past to feed the Encoder.
dec_window: Length of the future horizon to feed the Decoder and predict.
lead_time: Gap between the end of the encoder window and start of the decoder window. Use lead_time=1 for Forecasting (predicting next steps). Use lead_time=0 for Regression (predicting concurrently with the last encoder step).
stride: Step size between consecutive sliding windows.

Returns:

enc_x: Array of shape (..., Features, enc_window, Batch)
dec_x: Array of shape (..., Features, dec_window, Batch) (or nothing if forcings is omitted)
y_target: Array of shape (..., Features, dec_window, Batch)

EasyHybrid.toDataFrame Method

julia

toDataFrame(arr, target_names)

Extract specific target variables from a labeled array into a DataFrame with _pred suffix.

Arguments

arr: A labeled array or NamedTuple-like object with property access
target_names: Vector of target variable names to extract

Returns

DataFrame with columns named <target>_pred for each target

EasyHybrid.toDataFrame Method

julia

toDataFrame(nt::NamedTuple, target_names, y_ref::NamedTuple)

Extract targets from prediction NamedTuple nt into a DataFrame with _pred suffix. When the prediction has more rows than the reference y_ref (e.g. input_window > output_window), the last output_window rows are selected before flattening.

EasyHybrid.toDataFrame Method

julia

toDataFrame(nt::NamedTuple)

Convert a NamedTuple of arrays (vectors or matrices) into a DataFrame. Matrix values are flattened via vec so each key becomes a single column.

EasyHybrid.toDataFrame Method

julia

toDataFrame(arr::Union{KeyedArray{T, 2}, AbstractDimArray{T, 2}}, cols_dim=:variable, index_dim=:batch_size; index_col=:index)

Convert a 2D labeled array (KeyedArray or DimArray) to a DataFrame.

Arguments

arr: The 2D labeled array to convert
cols_dim: Dimension name to use as DataFrame columns (default: :variable)
index_dim: Dimension name to use as DataFrame row index (default: :batch_size)
index_col: Name for the index column in the result (default: :index)

Returns

DataFrame with columns from cols_dim keys and an index column from index_dim keys

EasyHybrid.toDataFrame Method

julia

toDataFrame(arr::AbstractLabeledArray{T, 3}, cols_dim=:variable, index_dim=:batch_size; slice_dim=:time, index_col=:index)

Convert a 3D labeled array (KeyedArray or DimArray) to a Dict of DataFrames, one per slice.

Arguments

arr: The 3D labeled array to convert
cols_dim: Dimension name to use as DataFrame columns (default: :variable)
index_dim: Dimension name to use as DataFrame row index (default: :batch_size)
slice_dim: Dimension name to slice along (default: :time)
index_col: Name for the index column in each result DataFrame (default: :index)

Returns

Dict{Any, DataFrame} mapping slice keys to DataFrames

EasyHybrid.toNamedTuple Method

toNamedTuple(ka::KeyedArray, variable::Symbol) Extract a single variable from a KeyedArray and return it as a vector.

Arguments:

ka: The KeyedArray or DimArray to unpack
variable: Symbol representing the variable to extract

Returns:

Vector containing the variable data

Example:

julia

# Extract just SW_IN from an array
sw_in = toNamedTuple(ds, :SW_IN)

EasyHybrid.toNamedTuple Method

julia

toNamedTuple(ka::Union{KeyedArray, AbstractDimArray}, variables::Vector{Symbol})

Extract specified variables from a KeyedArray or DimArray and return them as a NamedTuple of vectors.

Arguments:

ka: The KeyedArray or DimArray to unpack
variables: Vector of symbols representing the variables to extract

Returns:

NamedTuple with variable names as keys and vectors as values

Example:

julia

# Extract SW_IN and TA from an array
data = toNamedTuple(ds, [:SW_IN, :TA])
sw_in = data.SW_IN
ta = data.TA

EasyHybrid.toNamedTuple Method

toNamedTuple(ka::KeyedArray) Extract all variables from a KeyedArray and return them as a NamedTuple of vectors.

Arguments:

ka: The KeyedArray to unpack

Returns:

NamedTuple with all variable names as keys and vectors as values

Example:

julia

# Extract all variables from an array
data = toNamedTuple(ds)
# Access individual variables
sw_in = data.SW_IN
ta = data.TA
nee = data.NEE

EasyHybrid.to_dimArray Method

to_dimArray(df::DataFrame)

EasyHybrid.to_keyedArray Function

tokeyedArray(dfg::Union{Vector,GroupedDataFrame{DataFrame}}, vars=All())

EasyHybrid.to_keyedArray Method

tokeyedArray(df::DataFrame)

EasyHybrid.train Method

julia

train(model, data; train_cfg::TrainConfig = TrainConfig(), data_cfg::DataConfig = DataConfig())
train(model, data; kwargs...)

Train a hybrid model using the provided data.

Two equivalent calling styles are supported:

Typed configs — pass complete TrainConfig / DataConfig objects:

julia

train(model, data;
    train_cfg = TrainConfig(nepochs=100, batchsize=32),
    data_cfg  = DataConfig(split_data_at=0.8),
)

Flat kwargs — pass TrainConfig / DataConfig field names directly:

julia

train(model, data; nepochs=100, batchsize=32, split_data_at=0.8)

The two styles can also be mixed; flat kwargs override the corresponding fields of the supplied train_cfg / data_cfg:

julia

train(model, data; train_cfg = TrainConfig(nepochs=100), nepochs = 10)  # nepochs = 10

Returns nothing if data preparation fails (zero-size dimension in training or validation data).

Arguments

model: The hybrid model to train.
data: Training data, a single DimArray, a single DataFrame, or a single KeyedArray.

Keyword Arguments

train_cfg: Training configuration. See TrainConfig for all options.
data_cfg: Data preparation configuration. See DataConfig for all options.
Any other kwargs are forwarded as overrides to train_cfg / data_cfg.

Returns

A TrainResults with the following fields:

train_losses: Per-epoch training losses.
val_losses: Per-epoch validation losses.
snapshots: Model parameter snapshots taken during training.
train_obs_pred: Observed vs. predicted values on the training set.
val_obs_pred: Observed vs. predicted values on the validation set.
train_diffs: Additional diagnostic variables computed on the training set.
val_diffs: Additional diagnostic variables computed on the validation set.
ps: Final (or best) model parameters.
st: Final (or best) model state.
best_epoch: Epoch at which the best validation loss was achieved.
best_loss: Best validation loss recorded during training.

EasyHybrid.tune Method

julia

tune(hybrid_model, data, mspec::ModelSpec; kwargs...)
tune(hybrid_model, data; kwargs...)
tune(hybrid_model, data, train_cfg::TrainConfig; data_cfg::DataConfig = DataConfig(), kwargs...)

Construct a new hybrid model from hybrid_model plus hyperparameters, then call train.

Returns a TrainResults (or nothing if data preparation fails, as in train).

EasyHybrid.weight_l2 Method

julia

weight_l2(ps; key=:weight, normalize=false) -> Real

Sum of squared Frobenius norms over all parameter arrays in ps whose immediate parent field name is key (default :weight).

With normalize=true, returns the mean squared weight (sum divided by the number of scalar weights), so the value is independent of network width/depth.

Unlike sum(abs2, extract_weights(ps)), this fuses the tree walk with the reduction so it is safe to use inside Zygote-differentiated losses, e.g.:

julia

extra_loss = (ŷ, ps) -> (; l2_Rb = λ * weight_l2(ps.Rb; normalize=true),)

When ps is the loss function argument, gradients flow into the weight arrays.

EasyHybrid.@hybrid Macro

julia

@hybrid ModelName α β γ

Macro to define hybrid model structs with arbitrary numbers of physical parameters.

This defines a struct with:

Default fields: NN (neural network), predictors, forcing, targets.
Additional physical parameters, i.e., α β γ.

Examples

julia

@hybrid MyModel α β γ
@hybrid FluidModel (:viscosity, :density)
@hybrid SimpleModel :a :b

EasyHybrid.Transformers.EncoderDecoderModel Type

julia

EncoderDecoderModel(; in_features, dec_features, d_model, enc_layers, dec_layers, n_heads, n_kv_heads=n_heads, out_features, norm_eps=1.0f-5, dropout_rate=0.0f0, stem=nothing)

Creates a sequence-to-sequence Encoder-Decoder Transformer using GroupedQueryAttention for self-attention and MultiHeadSelfAttention for cross-attention.

Arguments

in_features: Number of input features for the encoder
dec_features: Number of input features for the decoder (shifted targets or covariates)
d_model: Dimensionality of the model's hidden states
enc_layers: Number of Transformer blocks in the encoder
dec_layers: Number of Transformer blocks in the decoder
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads (defaults to n_heads)
out_features: Dimensionality of the final output projection
norm_eps: Epsilon value for RMSNorm stability
dropout_rate: Dropout probability applied to attention and feedforward layers
stem: Optional Lux layer to apply as a feature extractor before encoder embedding

Returns

An EncoderDecoderModel container layer

EasyHybrid.Transformers.EncoderDecoderModel Method

julia

(m::EncoderDecoderModel)(enc_x, dec_x, ps, st; enc_causal=false, dec_causal=true)

Forward pass for the sequence-to-sequence EncoderDecoderModel.

Arguments

m: The EncoderDecoderModel
enc_x: Encoder input sequence data
dec_x: Decoder input sequence data (e.g. shifted targets)
ps: Model parameters
st: Model state
enc_causal: Boolean kwarg (default false). If true, applies causal masking to the encoder.
dec_causal: Boolean kwarg (default true). If true, applies causal masking to the decoder self-attention.

Returns

(out, new_st): A tuple containing the model predictions and updated state

EasyHybrid.Transformers.FeatureEmbedding Type

julia

FeatureEmbedding(in_features::Int, d_model::Int)

Creates a FeatureEmbedding layer for continuous/multivariate sequences. Unlike standard TokenEmbeddings (which use dictionaries for discrete tokens), this projects continuous features into the hidden d_model dimension via a Dense layer.

Arguments

in_features: Number of input features/covariates per timestep
d_model: Dimensionality of the model's hidden states

Returns

A FeatureEmbedding container layer

EasyHybrid.Transformers.FeatureEmbedding Method

julia

(m::FeatureEmbedding)(x, ps, st)

Forward pass for the continuous FeatureEmbedding.

Arguments

m: The FeatureEmbedding layer
x: Input sequence data of shape (in_features, seq_len, batch)
ps: Model parameters
st: Model state

Returns

(y, st_out): Projected features of shape (d_model, seq_len, batch) and updated state

EasyHybrid.Transformers.FeedForward Type

julia

FeedForward(dim::Int; multiple_of::Int=256, ffn_dim_multiplier::Union{Float64, Nothing}=nothing, dropout_rate=0.0f0)

Implements a SwiGLU FeedForward network. This block projects the input to a higher dimension, applies a Swish (SiLU) gating mechanism, and projects it back to the original dimension, offering better gradient flow than standard GELU MLPs.

Arguments

dim: Dimensionality of the model's hidden states (d_model)
multiple_of: Ensures the hidden dimension is a multiple of this value for hardware efficiency
ffn_dim_multiplier: Optional multiplier to explicitly scale the hidden dimension
dropout_rate: Dropout probability applied to the final projection

Returns

A FeedForward container layer

EasyHybrid.Transformers.FeedForward Method

julia

(m::FeedForward)(x, ps, st)

Forward pass for the SwiGLU FeedForward network.

Arguments

m: The FeedForward layer
x: Input sequence data of shape (d_model, seq_len, batch)
ps: Model parameters
st: Model state

Returns

(y, st_out): Projected sequence and updated state

EasyHybrid.Transformers.GroupedQueryAttention Type

julia

GroupedQueryAttention(dim::Int, n_heads::Int, n_kv_heads::Int; dropout_rate=0.0f0)

Implements Grouped Query Attention (GQA). GQA reduces memory and computational overhead during autoregressive generation by sharing key and value projections across multiple query heads.

EasyHybrid.Transformers.GroupedQueryAttention Method

julia

(m::GroupedQueryAttention)(x, ps, st; context=nothing, mask=nothing)

Standard forward pass for Grouped Query Attention without KVCache (e.g. for ViT or standard training).

Arguments

m: The GroupedQueryAttention layer
x: Input sequence data of shape (dim, seq_len, batch)
ps: Model parameters
st: Model state
context: Optional context sequence for cross-attention
mask: Optional attention mask

Returns

(out, st_out): Attended sequence and updated state

EasyHybrid.Transformers.GroupedQueryAttention Method

julia

GroupedQueryAttention(dim::Int, n_heads::Int, n_kv_heads::Int; dropout_rate=0.0f0)

Implements Grouped Query Attention (GQA). GQA reduces memory and computational overhead during autoregressive generation by sharing key and value projections across multiple query heads.

Arguments

dim: Dimensionality of the model's hidden states (d_model)
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads
dropout_rate: Dropout probability applied to the final attention projection

Returns

A GroupedQueryAttention container layer

EasyHybrid.Transformers.MultiHeadSelfAttention Type

julia

MultiHeadSelfAttention(d_model, n_heads; dropout_rate=0.0f0)

Standard Multi-Head Self Attention, primarily used for Cross-Attention in Encoder-Decoder architectures.

Arguments

d_model: Dimensionality of the model's hidden states
n_heads: Number of attention heads
dropout_rate: Dropout probability applied to the final attention projection

Returns

A MultiHeadSelfAttention container layer

EasyHybrid.Transformers.MultiHeadSelfAttention Method

julia

(m::MultiHeadSelfAttention)(x, ps, st; context=nothing, mask=nothing)

Forward pass for standard Multi-Head Self Attention.

Arguments

m: The MultiHeadSelfAttention layer
x: Input sequence data of shape (d_model, seq_len, batch) (used for Queries, and Keys/Values if context is nothing)
ps: Model parameters
st: Model state
context: Optional context sequence of shape (d_model, context_len, batch) for cross-attention (Keys/Values)
mask: Optional attention mask

Returns

(out, st_out): Attended sequence and updated state

EasyHybrid.Transformers.PatchEmbedding Type

julia

PatchEmbedding(patch_size::Tuple, in_channels::Int, d_model::Int; ndims::Int=2)

Creates a PatchEmbedding for Vision Transformers using a Convolutional layer.

Arguments

patch_size: Tuple defining the patch dimensions
in_channels: Number of input channels
d_model: Dimensionality of the model's hidden states
ndims: 2 for spatial grids (images), 3 for spatio-temporal grids (video/climate)

Returns

A PatchEmbedding container layer

EasyHybrid.Transformers.PatchEmbedding Method

julia

(m::PatchEmbedding)(x, ps, st)

Forward pass for PatchEmbedding.

Arguments

m: The PatchEmbedding layer
x: Input data of shape (W, H, C, B) for 2D or (W, H, T, C, B) for 3D
ps: Model parameters
st: Model state

Returns

(y, st_out): Flattened sequence of patches of shape (d_model, seq_len, batch) and updated state

EasyHybrid.Transformers.PatchUnEmbedding Type

julia

PatchUnEmbedding(patch_size::Tuple, d_model::Int, out_channels::Int, grid_size::Tuple; ndims::Int=2)

Creates a PatchUnEmbedding layer to reconstruct spatial or spatio-temporal grids from a sequence of patches. This is the reverse of PatchEmbedding.

Arguments

patch_size: Tuple defining the patch dimensions
d_model: Dimensionality of the model's hidden states
out_channels: Number of output channels for the reconstructed grid
grid_size: The number of patches in each spatial/temporal dimension (e.g., (W', H') for 2D).
ndims: 2 for spatial grids (images), 3 for spatio-temporal grids (video/climate)

Returns

A PatchUnEmbedding container layer

EasyHybrid.Transformers.PatchUnEmbedding Method

julia

(m::PatchUnEmbedding)(x, ps, st)

Forward pass for PatchUnEmbedding.

Arguments

m: The PatchUnEmbedding layer
x: Input sequence of patches of shape (d_model, seq_len, batch)
ps: Model parameters
st: Model state

Returns

(out, st_out): Reconstructed grid of shape (W, H, C, B) for 2D or (W, H, T, C, B) for 3D and updated state

EasyHybrid.Transformers.PositionEmbedding Type

julia

PositionEmbedding(max_positions::Integer, d_model::Integer; dim::Int=2)

Creates a learned additive Positional Embedding layer. Typically used when RoPE is disabled.

Arguments

max_positions: The maximum sequence length supported
d_model: Dimensionality of the model's hidden states
dim: The dimension along which to add the embeddings (default 2 for seq_len)

Returns

A PositionEmbedding layer

EasyHybrid.Transformers.PositionEmbedding Method

julia

(m::PositionEmbedding)(x, ps, st)

Forward pass for the additive PositionEmbedding.

Arguments

m: The PositionEmbedding layer
x: Hidden states of shape (d_model, seq_len, batch)
ps: Model parameters
st: Model state

Returns

(y, st): Hidden states with positional embeddings added, and updated state

EasyHybrid.Transformers.PrefixTokens Type

julia

PrefixTokens(embed_dim::Int; use_cls_token::Bool=true, n_register_tokens::Int=0)

Creates a layer that maintains learned [CLS] and [REGISTER] tokens. During the forward pass, these tokens are repeated for the batch size and prepended to the sequence.

Arguments

embed_dim: Dimensionality of the model's hidden states
use_cls_token: If true, includes a single [CLS] token
n_register_tokens: Number of [REGISTER] (storage) tokens to include

Returns

A PrefixTokens layer

EasyHybrid.Transformers.TransformerBlock Type

julia

TransformerBlock(dim, n_heads, n_kv_heads; norm_eps=1.0f-5, cross_attention=false, dropout_rate=0.0f0)

Creates a State-of-the-Art Transformer Block using RMSNorm, SwiGLU FeedForward, and GroupedQueryAttention. If cross_attention is true, an additional MultiHeadSelfAttention block is added for Encoder-Decoder architectures.

Arguments

dim: Dimensionality of the model's hidden states (d_model)
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads
norm_eps: Epsilon value for RMSNorm stability
cross_attention: If true, adds a cross-attention layer (e.g. for decoders)
dropout_rate: Dropout probability applied to attention and feedforward layers
layer_scale_init: Initial value for LayerScale (e.g., 1e-5). If nothing, LayerScale is not used.

Returns

A TransformerBlock container layer

EasyHybrid.Transformers.TransformerBlock Method

julia

(m::TransformerBlock)(x, ps, st; cosf=nothing, sinf=nothing, context=nothing, mask=nothing)

Forward pass for a single TransformerBlock.

Arguments

m: The TransformerBlock layer
x: Input sequence data
ps: Model parameters
st: Model state
cosf: Optional precomputed cosine frequencies for RoPE
sinf: Optional precomputed sine frequencies for RoPE
context: Optional context sequence for cross-attention
mask: Optional attention mask

Returns

(y, st_out): Processed sequence and updated state

EasyHybrid.Transformers.TransformerModel Type

julia

TransformerModel(; in_features, d_model, n_layers, n_heads, n_kv_heads=n_heads, out_features, dropout_rate=0.0f0, stem=nothing, max_positions=nothing, norm_eps=1.0f-5)

Creates a continuous sequence TransformerModel using GroupedQueryAttention and RMSNorm. Optionally accepts a stem (e.g. a CNN or LSTM) to act as a feature extractor before embedding.

Arguments

in_features: Number of input features/covariates per timestep
d_model: Dimensionality of the model's hidden states
n_layers: Number of Transformer blocks
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads (defaults to n_heads)
out_features: Dimensionality of the final output projection
dropout_rate: Dropout probability applied to attention and feedforward layers
stem: Optional Lux layer to apply as a feature extractor before embedding
max_positions: Unused placeholder for additive embeddings if needed later
norm_eps: Epsilon value for RMSNorm stability
layer_scale_init: Initial value for LayerScale (e.g., 1e-5). If nothing, LayerScale is not used.

Returns

A TransformerModel container layer

EasyHybrid.Transformers.TransformerModel Method

julia

(m::TransformerModel)(x, ps, st; causal=false)

Forward pass for the continuous sequence TransformerModel.

Arguments

m: The TransformerModel model
x: Input sequence data of shape (in_features, seq_len, batch) or (spatial..., batch) if stem is used
ps: Model parameters
st: Model state
causal: Boolean kwarg (default false). If true, applies a causal upper-triangular mask to prevent peeking into the future.

Returns

(y, st_out): A tuple containing the model predictions and updated state

EasyHybrid.Transformers.TransformerStack Type

julia

TransformerStack(layers::Union{Vector, Tuple})

A sequential container for Transformer blocks. While Lux.Chain passes inputs through a sequence of layers, TransformerStack is specifically designed to correctly propagate arbitrary keyword arguments (such as mask, cache, and RoPE frequencies) down to each individual block, which is required for advanced attention mechanisms.

Arguments

layers: A Tuple or Vector of TransformerBlock layers

Returns

A TransformerStack container layer

EasyHybrid.Transformers.TransformerStack Method

julia

(m::TransformerStack)(x, ps, st; kwargs...)

Forward pass for the TransformerStack. Passes x sequentially through each TransformerBlock in the stack. Arbitrary keyword arguments (like mask, cache, start_pos, cosf, sinf, context) are correctly propagated to every block.

Arguments

m: The TransformerStack container
x: Input sequence data
ps: Model parameters
st: Model state
kwargs...: Optional arguments passed down to the blocks

Returns

(y, st_new): Processed sequence and updated state

EasyHybrid.Transformers.VisionToVisionModel Type

julia

VisionToVisionModel(; patch_size, grid_size, in_channels, out_channels, d_model, n_layers, n_heads, n_kv_heads=n_heads, max_positions, ndims=2, use_rope=false, dropout_rate=0.0f0, stem=nothing, norm_eps=1.0f-5)

Creates a Vision-to-Vision Transformer (e.g. for Image-to-Image regression or Grid-to-Grid forecasting). The input grid is processed via PatchEmbedding and the output sequence is reconstructed back into a grid via PatchUnEmbedding.

Arguments

patch_size: Tuple defining the spatial/spatio-temporal dimensions of each patch
grid_size: Tuple defining the number of patches in each dimension (e.g., (W', H'))
in_channels: Number of input channels
out_channels: Number of output channels for the reconstructed grid
d_model: Dimensionality of the model's hidden states
n_layers: Number of Transformer blocks
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads (defaults to n_heads)
max_positions: Maximum number of positions for additive positional embeddings
ndims: Number of spatial/spatio-temporal dimensions (2 for spatial, 3 for spatio-temporal)
use_rope: If true, uses Rotary Positional Embeddings instead of additive embeddings
dropout_rate: Dropout probability
stem: Optional Lux layer to apply before patch embedding
norm_eps: Epsilon value for RMSNorm stability
use_cls_token: If true, prepends a [CLS] token.
n_register_tokens: Number of [REGISTER] tokens to prepend.
layer_scale_init: Initial value for LayerScale (e.g., 1e-5). If nothing, LayerScale is not used.

Returns

A VisionToVisionModel container layer

EasyHybrid.Transformers.VisionToVisionModel Method

julia

(m::VisionToVisionModel)(x, ps, st)

Forward pass for the Vision-to-Vision Transformer model.

Arguments

m: The VisionToVisionModel model
x: Input grid of shape (W, H, C, B) for 2D or (W, H, T, C, B) for 3D
ps: Model parameters
st: Model state

Returns

(y, st_out): Predicted output grid of shape (W, H, out_channels, B) (or 3D equivalent) and updated state

EasyHybrid.Transformers.VisionTransformer Type

julia

VisionTransformer(; patch_size, in_channels, d_model, n_layers, n_heads, n_kv_heads=n_heads, max_positions, num_classes, ndims=2, use_rope=false, dropout_rate=0.0f0, stem=nothing, norm_eps=1.0f-5)

Creates a Vision Transformer. If stem is provided, it acts as a Hybrid feature extractor before the PatchEmbedding.

Arguments

patch_size: Tuple defining the spatial/spatio-temporal dimensions of each patch
in_channels: Number of input channels (e.g., 1 for single variable, 3 for RGB)
d_model: Dimensionality of the model's hidden states
n_layers: Number of Transformer blocks
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads (defaults to n_heads for standard attention, less for GQA)
max_positions: Maximum number of positions for additive positional embeddings
num_classes: Dimensionality of the final output (e.g., number of classes for classification)
ndims: Number of spatial/spatio-temporal dimensions (2 for spatial, 3 for spatio-temporal)
use_rope: If true, uses Rotary Positional Embeddings instead of additive embeddings
dropout_rate: Dropout probability applied to attention and feedforward layers
stem: Optional Lux layer to apply before patch embedding (e.g., a CNN for Hybrid ViT)
norm_eps: Epsilon value for RMSNorm stability
use_cls_token: If true, prepends a [CLS] token and uses it for the final output instead of GAP.
n_register_tokens: Number of [REGISTER] tokens to prepend.
layer_scale_init: Initial value for LayerScale (e.g., 1e-5). If nothing, LayerScale is not used.

Returns

A VisionTransformer container layer

EasyHybrid.Transformers.VisionTransformer Method

julia

(m::VisionTransformer)(x, ps, st)

Forward pass for the Vision Transformer model.

Arguments

m: The VisionTransformer model
x: Input data of shape (W, H, C, B) for 2D or (W, H, T, C, B) for 3D
ps: Model parameters
st: Model state

Returns

(y, st_out): A tuple containing the model predictions (logits) and updated state

EasyHybrid.Transformers.VisionEncoderDecoderModel Method

julia

VisionEncoderDecoderModel(; patch_size, in_channels, dec_features, d_model, enc_layers, dec_layers, n_heads, n_kv_heads=n_heads, out_features, ndims=2, norm_eps=1.0f-5, dropout_rate=0.0f0, stem=nothing)

Creates a sequence-to-sequence Encoder-Decoder Transformer where the Encoder processes spatial or spatio-temporal data (via PatchEmbedding) and the Decoder processes continuous sequential covariates (via FeatureEmbedding).

Arguments

patch_size: Tuple defining the spatial/spatio-temporal dimensions of each encoder patch
in_channels: Number of input channels for the encoder grid
dec_features: Number of input features for the decoder (shifted targets or covariates)
d_model: Dimensionality of the model's hidden states
enc_layers: Number of Transformer blocks in the encoder
dec_layers: Number of Transformer blocks in the decoder
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads (defaults to n_heads)
out_features: Dimensionality of the final output projection
ndims: Number of spatial/spatio-temporal dimensions for the encoder (2 or 3)
norm_eps: Epsilon value for RMSNorm stability
dropout_rate: Dropout probability applied to attention and feedforward layers
stem: Optional Lux layer to apply as a feature extractor before patch embedding

Returns

A EncoderDecoderModel configured for vision-to-sequence tasks

EasyHybrid.Transformers.VisionToVisionEncoderDecoderModel Method

julia

VisionToVisionEncoderDecoderModel(; patch_size, grid_size, in_channels, dec_channels, out_channels, d_model, enc_layers, dec_layers, n_heads, n_kv_heads=n_heads, ndims=2, norm_eps=1.0f-5, dropout_rate=0.0f0, stem=nothing)

Creates a sequence-to-sequence Encoder-Decoder Transformer where BOTH inputs and outputs are spatial or spatio-temporal grids. The Encoder processes the historical maps, the Decoder processes known future covariate maps, and the Output predicts target future maps.

Arguments

patch_size: Tuple defining the spatial/spatio-temporal dimensions of each patch
grid_size: Tuple defining the number of patches in each dimension (e.g., (W', H'))
in_channels: Number of input channels for the encoder grid
dec_channels: Number of input channels for the decoder grid (covariates)
out_channels: Number of output channels for the reconstructed prediction grid
d_model: Dimensionality of the model's hidden states
enc_layers: Number of Transformer blocks in the encoder
dec_layers: Number of Transformer blocks in the decoder
n_heads: Number of query attention heads
n_kv_heads: Number of key/value attention heads
ndims: Number of spatial/spatio-temporal dimensions (2 or 3)
norm_eps: Epsilon value for RMSNorm stability
dropout_rate: Dropout probability
stem: Optional Lux layer to apply as a feature extractor before patch embedding

Returns

A EncoderDecoderModel configured for grid-to-grid forecasting tasks

EasyHybrid.Transformers.apply_rotary_embeddings Method

julia

apply_rotary_embeddings(x::AbstractArray{T, 4}, cosf, sinf) where {T}

Applies the precomputed rotary positional embeddings to a 4D tensor x.

How it works: It splits the head_dim (feature dimension) in half and applies a 2D rotation. Because it directly rotates the hidden representations in the complex plane, the dot product between any Query and Key during attention will naturally decay based on their relative distance in the sequence. This injects highly effective, shift-invariant positional context without adding any learned parameters to the network.

Arguments

x: Input tensor of shape (head_dim, n_heads, seq_len, batch) (either Queries or Keys)
cosf: Precomputed cosine frequencies from precompute_rope_freqs
sinf: Precomputed sine frequencies from precompute_rope_freqs

Returns

A rotated tensor of the same shape as x

EasyHybrid.Transformers.extract_features Method

julia

extract_features(m::Union{VisionTransformer, VisionToVisionModel}, x, ps, st; n_blocks::Union{Int, Nothing} = 1, blocks::Union{AbstractVector{Int}, Nothing} = nothing)

Extracts the spatial/spatio-temporal features from intermediate blocks of the VisionTransformer. This is particularly useful when using a pretrained LingBot-Vision or DINOv2 model as a frozen feature extractor for downstream dense prediction tasks (like depth estimation or segmentation).

Arguments

m: The VisionTransformer or VisionToVisionModel layer
x: Input data of shape (W, H, C, B) for 2D or (W, H, T, C, B) for 3D
ps: Model parameters
st: Model state
n_blocks: The number of intermediate blocks to extract from the end of the transformer (default: 1).
blocks: Specific block indices to extract (e.g. [1, 3, 5]). If provided, n_blocks is ignored.

Returns

A tuple of tensors, one for each extracted block, reshaped back into spatial grids of shape (W', H', d_model, B) for 2D or (W', H', T', d_model, B) for 3D.

EasyHybrid.Transformers.precompute_rope_freqs Method

julia

precompute_rope_freqs(x::AbstractArray, head_dim::Int, max_seq_len::Int; theta::Float32 = 10_000f0)

Precomputes the trigonometric frequencies (cosine and sine) needed for Rotary Positional Embeddings (RoPE).

Why RoPE matters: Unlike standard additive positional embeddings (which encode absolute positions as fixed, learned vectors), RoPE encodes relative positional information directly into the attention mechanism's Queries and Keys using complex rotations. This allows the model to:

Extrapolate to sequence lengths (or spatial grids) longer than those seen during training.
Intuitively understand the relative distance between tokens, making it incredibly robust for spatial-temporal grids and continuous time-series where boundaries might shift.

Arguments

head_dim: The dimensionality of each attention head
max_seq_len: The maximum sequence length to precompute frequencies for
theta: The base for the inverse frequency computation (default 10_000f0)

Returns

(cosf, sinf): A tuple of cosine and sine frequency matrices, each of shape (head_dim/2, max_seq_len)