API

SegmentedTimeSeries(data; segment_length=2, shift=nothing, batchsize=1, shuffle=false, partial_segment=false, partial_batch=false, rng=GLOBAL_RNG)

An object that iterates over mini-batches of segments of data, each segment containing segment_length data points, each mini-batch containing batchsize segments. The last dimension in each tensor is the time dimension.

Arguments

• data: Input data (array, tuple, or named tuple).
• segment_length: Number of time points in each segment.
• shift: Step size between consecutive segments (default: segment_length).
• batchsize: Number of segments per batch.
• shuffle: Whether to shuffle segment order.
• partial_segment: Allow shorter final segments.
• partial_batch: Allow smaller final batches.
• rng: Random number generator for shuffling.

Inputs

• data: The time series data to segment.

Outputs

• Iterator yielding batches of data segments.

Behavior

Creates overlapping or non-overlapping segments from time series data for training dynamical models. Segments can be shuffled and batched for efficient training.

Example

julia> Xtrain = rand(10, 100)
julia> sdl = SegmentedTimeSeries(Xtrain; segment_length=2, batchsize=1)
julia> for batch in sdl
           println("Batch: ", summary(batch))
       end

!!!warning Undefined behavior when data dimensions are incompatible

tokenize(sdl::SegmentedTimeSeries)

Convert a SegmentedTimeSeries to use token-based indexing.

Arguments

• sdl: The SegmentedTimeSeries to tokenize.

Inputs

• sdl: SegmentedTimeSeries object.

Outputs

• Tokenized SegmentedTimeSeries with integer-based segment access.

Behavior

Transforms segment indices into a token-based system for easier access to individual segments.

Example

julia> sdl = SegmentedTimeSeries(rand(10, 100); segment_length=2)
julia> tokenized_sdl = tokenize(sdl)
julia> tokens(tokenized_sdl) # Returns available tokens

tokens(sdl::SegmentedTimeSeries)

Get the available tokens for a tokenized SegmentedTimeSeries.

Arguments

• sdl: A tokenized SegmentedTimeSeries.

Inputs

• sdl: Tokenized SegmentedTimeSeries object.

Outputs

• Range of available tokens (1 to number of segments).

Behavior

Returns the range of tokens that can be used to access individual segments in a tokenized SegmentedTimeSeries.

Example

julia> sdl = SegmentedTimeSeries(rand(10, 100); segment_length=2)
julia> tokenized_sdl = tokenize(sdl)
julia> collect(tokens(tokenized_sdl)) # [1, 2, 3, ...]

create_train_val_loaders(data; segment_length, valid_length, kwargs...)

Create separate training and validation SegmentedTimeSeries loaders from a dataset with the same number of batches.

This function splits the data into non-overlapping training and validation segments. The training data uses segments with gaps equal to valid_length to leave space for validation segments. The validation data starts after the first training segment and uses segments of length valid_length. Both loaders are guaranteed to have the same number of batches, with tokens referring to the same ordering.

Arguments

• data: Input data (can be an array, tuple, or named tuple)
• segment_length: Size of training segments
• valid_length: Size of validation segments
• kwargs...: Additional arguments passed to SegmentedTimeSeries constructors

Returns

• dataloader_train: SegmentedTimeSeries for training data
• dataloader_valid: SegmentedTimeSeries for validation data

Examples

With array data

data = rand(10, 100)  # 10 features, 100 time steps
train_loader, val_loader = create_train_val_loaders(data; 
    segment_length=20, valid_length=5, batchsize=4)
# Both loaders will have the same number of batches
@assert length(train_loader) == length(val_loader)

With tuple data (data, time steps)

data = rand(10, 100)
tsteps = 1:100
train_loader, val_loader = create_train_val_loaders((data, tsteps); 
    segment_length=20, valid_length=5, batchsize=4)

With named tuple data

dataset = (observations = rand(10, 100), times = 1:100, metadata = rand(5, 100))
train_loader, val_loader = create_train_val_loaders(dataset; 
    segment_length=20, valid_length=5)

Notes

• Training segments are spaced segment_length + valid_length apart to avoid overlap with validation
• Validation segments start at position segment_length + 1 to avoid overlap with first training segment
• Both loaders have partial_segment = false and partial_batch = false to ensure consistent sizes
• Both loaders are guaranteed to have the same number of batches for synchronized training/validation

ParameterLayer(;constraint::AbstractConstraint = NoConstraint(), 
                init_value = (;), 
                init_state_value = (;))

A layer representing parameters, optionally with constraints.

Arguments

• constraint: Constraint to transform parameters.
• init_value: Initial parameter values (NamedTuple or AbstractArray).
• init_state_value: Internal state (NamedTuple).

Inputs

• ps: Parameters of the layer.
• st: States of the layer.

Outputs

• Constrained parameter values merged with states.

Behavior

Applies constraints to parameters during forward pass. Parameters are transformed from unconstrained to constrained space.

Example

julia> param = ParameterLayer(; constraint = NoConstraint(),
                            init_value = (;a = ones(2)), 
                            init_state_value = (;b = (0.0, 1.0)))
julia> ps, st = Lux.setup(Random.default_rng(), param)
julia> x, _ = param(ps, st)
julia> x == (a = [1.0, 1.0], b = (0.0, 1.0))
true

NoConstraint()

Applies no transformation to parameters.

Arguments

• None.

Inputs

• x: Parameter values.
• st: States.

Outputs

• Unmodified parameter values and states.

Behavior

Identity transformation - parameters remain unconstrained.

Example

julia> constraint = NoConstraint()
julia> x, st = constraint([1.0, 2.0], (;))
([1.0, 2.0], (;))

BoxConstraint(lb::AbstractArray, ub::AbstractArray)

Constrains parameters to lie within specified bounds using sigmoid transformation.

Arguments

• lb: Lower bounds array.
• ub: Upper bounds array.

Inputs

• y: Unconstrained parameter values.
• st: States containing bounds.

Outputs

• Constrained parameter values within bounds.

Behavior

Transforms unconstrained parameters to constrained space using sigmoid function.

Example

julia> constraint = BoxConstraint([0.0], [1.0])
julia> x, st = constraint([0.0], (;lb=[0.0], ub=[1.0]))
([0.5], (;lb=[0.0], ub=[1.0]))

NamedTupleConstraint(constraints::NamedTuple)

Applies different constraints to different fields of a NamedTuple.

Arguments

• constraints: NamedTuple of constraint objects.

Inputs

• x: NamedTuple of parameter values.
• st: NamedTuple of states.

Outputs

• Constrained parameter values.

Behavior

Applies field-specific constraints to NamedTuple parameters.

Example

julia> constraints = (a = BoxConstraint([0.0], [1.0]), b = NoConstraint())
julia> constraint = NamedTupleConstraint(constraints)

ODEModel(layers::NamedTuple, dudt::Function; kwargs...)

Wraps an ODE model for simulation using Lux layers.

Arguments

• layers: NamedTuple of Lux layers representing the layers of the model.
• dudt: Function that computes the derivative of the state, with signature dudt(layers, u, ps, t).
• kwargs: Additional keyword arguments passed to the solver (e.g., tspan, saveat, alg).

Inputs

• (x, ps, st)
  • x: a NamedTuple or AbstractVector{NamedTuple} (batch mode).
  • ps: Parameters of the model.
  • st: States of the model.
• A tuple of (x, ps, st): batch mode.
• (ps, st): If x not provided, defaults to kwargs.

Outputs

• (sol, st)
  • sol: Solution of the ODE problem, with second dimension corresponding to time and batches stacked along the third dimension, if applicable.
  • st: Updated states of the model.

Behavior

layers are wrapped in StatefulLuxLayers to maintain their states. The derivative function dudt should be defined as dudt(layers, u, ps, t) where u is the current state and t is the current time. The function returns the derivative of the state.

Example

julia> layers = (; layer1 = Lux.Dense(10, 10, relu))
julia> dudt(layers, u, ps, t) = layers.layer1(u, ps.layer1)[1]
julia> ode_model = ODEModel(layers, dudt, tspan = (0f0, 1f0), saveat = range(0f0, stop=1f0, length=100), alg = Tsit5())
julia> ps, st = Lux.setup(Random.default_rng(), ode_model)
julia> ode_model((; u0 = ones(Float32, 10)), ps, st)

!!!warning Undefined behavior when ps is not a NamedTuple

AnalyticModel(layers::NamedTuple, fun::Function; kwargs...)

Wraps an analytic model for direct evaluation using Lux layers.

Arguments

• layers: NamedTuple of Lux layers representing the layers of the model.
• fun: Function that computes the analytic solution, with signature fun(layers, u0, t0, ps, t).
• kwargs: Additional keyword arguments (e.g., default values for u0, tspan, saveat).

Inputs

• (x, ps, st)
  • x: a NamedTuple or AbstractVector{NamedTuple} (batch mode).
  • ps: Parameters of the model.
  • st: States of the model.
• A tuple of (x, ps, st): batch mode.
• (ps, st): If x not provided, defaults to kwargs.

Outputs

• (sol, st)
  • sol: Solution array evaluated at specified time points, with second dimension corresponding to time and batches stacked along the third dimension, if applicable.
  • st: Updated states of the model.

Behavior

layers are wrapped in StatefulLuxLayer to maintain their states. The analytic function fun should be defined as fun(layers, u0, t0, ps, t) where t can be a vector of time points and t0 is extracted from tspan.

Example

julia> layers = (; params = ParameterLayer(init_value = (a = 1.0, b = 0.5)))
julia> analytic_solution(layers, u0, t0, ps, t) = u0 .* exp.(layers.params(ps.params)[1].a .* (t .- t0))
julia> model = AnalyticModel(layers, analytic_solution; u0 = [1.0], tspan = (0.0, 1.0), saveat = 0:0.1:1.0)
julia> ps, st = Lux.setup(Random.default_rng(), model)
julia> model((; u0 = [1.0]), ps, st)

!!!warning Undefined behavior when ps is not a NamedTuple

ARModel(layers::NamedTuple, fun::Function; kwargs...)

Wraps an autoregressive (AR) model.

Arguments

• layers: NamedTuple of Lux layers representing the layers associated with the model.
• fun: Function that computes the next time step, with signature fun(layers, u, ps, t).
• kwargs: Additional keyword arguments (e.g., default values for u0, tspan, saveat, dt).

Inputs

• (x, ps, st)
  • x: a NamedTuple or AbstractVector{NamedTuple} (batch mode).
  • ps: Parameters of the model.
  • st: States of the model.
• A tuple of (x, ps, st): batch mode.
• (ps, st): If x not provided, defaults to kwargs.

Outputs

• (sol, st)
  • sol: Solution array with iterative predictions, with second dimension corresponding to time and batches stacked along the third dimension, if applicable.
  • st: Updated states of the model.

Behavior

layers are wrapped in StatefulLuxLayer to maintain their states. The AR function fun should be defined as fun(layers, u, ps, t) where:

• u is the current state
• t is the current time
• The function returns the next state

The model iteratively applies the function to generate a time series from initial conditions.

Example

julia> using HybridDynamicModels, Lux, Random

julia> layers = (; 
           predictor = Dense(2, 2), 
           params = ParameterLayer(init_value = (; decay = [95f-2],))
       );

julia> ar_step(layers, u, ps, t) = layers.predictor(u, ps.predictor) .* layers.params(ps.params).decay;

julia> model = ARModel(layers, ar_step; dt = 1f-1, u0 = [1f0, 5f-1], tspan = (0f0, 1f0), saveat = 0:1f-1:1f0);

julia> ps, st = Lux.setup(Random.default_rng(), model);

julia> x = (; u0 = [1f0, 5f-1]);

julia> y, st = model(x, ps, st);

julia> size(y) # 2 state variables, 11 time points
(2, 11) 

!!!warning Undefined behavior when x is not a NamedTuple

ICLayer(ics::AbstractLuxLayer)
ICLayer(ics::<:ParameterLayer)
ICLayer(ics::Vector{<:ParameterLayer})

Initial condition layer.

Arguments

• ics: Lux layer, ParameterLayer, or vector of ParameterLayer.

Inputs

• (x, ps, st) with x a NamedTuple or AbstractVector{NamedTuple} (batch mode).
• (ps, st) when ics is a ParameterLayer.

Outputs

• Initial conditions merged with other fields.
• Updated states.

Behavior

Processes initial conditions through wrapped layers and merges with input data.

Example

julia> ic_layer = ICLayer(ParameterLayer(init_value = (;u0 = [1.0])))

train(backend::AbstractOptimBackend, model, dataloader::SegmentedTimeSeries, infer_ics::InferICs, rng=Random.default_rng(); pstype=Lux.f64)

Train a dynamical model using segmented time series data.

Arguments

• backend: Training configuration and optimization settings.
• model: Lux model to train.
• dataloader: Time series data split into segments.
• infer_ics: Initial condition inference configuration.
• rng: Random number generator.
• pstype: Precision type for parameters.

Inputs

• backend: Training backend (SGDBackend, MCSamplingBackend, etc.).
• model: Model to train.
• dataloader: SegmentedTimeSeries data.
• infer_ics: InferICs configuration.

Outputs

• NamedTuple with training results (varies by backend).

Behavior

Trains model using specified backend on segmented time series data.

Example

julia> backend = SGDBackend(Adam(1e-3), 1000, AutoZygote(), MSELoss())
julia> dataloader = SegmentedTimeSeries(data, segment_length=20)
julia> infer_ics = InferICs(true)
julia> result = train(backend, model, dataloader, infer_ics)

InferICs(infer::Bool, u0_constraint=NoConstraint())

Configuration for initial condition inference in training.

Arguments

• infer: Whether to treat initial conditions as learnable parameters.
• u0_constraint: Constraint for initial condition optimization.

Inputs

• infer: Boolean flag for inference.
• u0_constraint: Constraint object.

Outputs

• InferICs configuration object.

Behavior

Controls whether initial conditions are learned or fixed during training.

Example

julia> infer_ics = InferICs(true, NoConstraint())

SGDBackend(opt, n_epochs, adtype, loss_fn, callback)

Training backend using Lux.jl for mode estimation.

!!! warning Conditional loading You need to load Optimisers, ComponentArrays and Lux before loading HybridDynamicModels to use SGDBackend.

Arguments

• opt: Optimizers.jl rule for parameter updates.
• n_epochs: Number of training epochs.
• adtype: Automatic differentiation backend.
• loss_fn: Loss function for training.
• callback: User-defined callback function.

Inputs

• opt: Optimization rule (e.g., Adam(1e-3)).
• n_epochs: Total training epochs.
• adtype: AD backend (e.g., AutoZygote()).
• loss_fn: Loss function.
• callback: Optional callback.

Outputs

• NamedTuple with trained parameters and states.

Behavior

Uses stochastic gradient descent for maximum likelihood estimation.

Example

julia> backend = SGDBackend(Adam(1e-3), 1000, AutoZygote(), MSELoss())

BayesianLayer(layer, priors)

Wrapper layer that adds Bayesian priors to Lux layers for probabilistic modeling.

Arguments

• layer: Lux layer to make Bayesian.
• priors: Prior distributions for parameters.

Inputs

• layer: Any Lux layer.
• priors: Distribution or NamedTuple of distributions.

Outputs

• Layer with Bayesian priors for MCMC inference.

Behavior

Enables probabilistic modeling by attaching priors to layer parameters.

Example

julia> dense_layer = Dense(10, 5)
julia> bayesian_dense = BayesianLayer(dense_layer, Normal(0, 1))

getpriors(layer)

Extract prior distributions from Bayesian layers.

Arguments

• layer: Layer or model containing BayesianLayer components.

Inputs

• layer: Bayesian or composite layer.

Outputs

• NamedTuple of prior distributions.

Behavior

Recursively extracts priors from Bayesian layers in model hierarchy.

Example

julia> priors = getpriors(bayesian_model)

create_turing_model(ps_priors, data_distrib, st_model)

Create a Turing model for Bayesian inference from a BayesianLayer model.

Arguments

• ps_priors: A nested structure (typically a NamedTuple) containing prior distributions for model parameters. Each leaf should be a Distributions.Distribution.
• data_distrib: A function or distribution constructor that creates the likelihood distribution for observed data points.
• st_model: A stateful Lux model that can be called with parameters to generate predictions.

Returns

A function (xs, ys) -> Turing.Model that creates a Turing model when given input data xs and observed data ys.

MCSamplingBackend(sampler, n_iterations, datadistrib; kwargs...)

Training backend for Bayesian inference using Monte Carlo sampling. !!! warning Conditional loading You need to load Turing, ComponentArrays and Lux before loading HybridDynamicModels to use MCSamplingBackend.

Arguments

• sampler: Turing.jl MCMC sampler.
• n_iterations: Number of MCMC samples.
• datadistrib: Data distribution for likelihood.
• kwargs: Additional sampler options.

Inputs

• sampler: MCMC sampling algorithm.
• n_iterations: Total posterior samples.
• datadistrib: Distribution for data likelihood.

Outputs

• NamedTuple with MCMC chains and model state.

Behavior

Performs Bayesian inference using MCMC sampling on models with priors.

Example

julia> backend = MCSamplingBackend(NUTS(0.65), 1000, LogNormal)