Documentation	Build Status	Julia	Testing

HybridDynamicModels.jl

Lux.jl layers and utilities to build and train hybrid dynamic models.

HybridDynamicModels.jl is a toolbox for easily building and training hybrid dynamic models which combine mechanistic and data driven components. Built on top of the deep learning framework Lux.jl, it enables both gradient descent optimization and Bayesian inference.

🚀 Key Features

Dynamic model layers

• ICLayer: For specifying initial conditions (ICs).
• ODEModel: For specifying (hybrid) ODEs
• ARModel: For specifying (hybrid) autoregressive models
• AnalyticModel: For specifying (hybrid) explicit dynamical models

Utility layers for hybrid modeling

• ParameterLayer: Learnable parameters, composable with optional Constraint layers
• BayesianLayer: Add probabilistic priors to any Lux layer

Data loaders

• SegmentedTimeSeries: Time series data loader with segmentation, implementing mini-batching.

Training API, with following backends

• SGDBackend: Gradient descent optimization with Optimisers.jl and Lux.jl training API
• MCSamplingBackend: Full Bayesian inference with uncertainty quantification using DynamicPPL.jl and Turing.jl.

📦 Installation

using Pkg
Pkg.add("HybridDynamicModels")

🔥 Quick Start

Autoregressive hybrid model

using Lux
using HybridDynamicModels
using Random

# Dense layer for interactions
interaction_layer = Dense(2, 2, tanh)

# Parameter layer for growth/decay rates
rate_params = ParameterLayer(init_value = (growth = [0.1], decay = [0.05]),
    constraint = NamedTupleConstraint((; growth = BoxConstraint([0.0], [1.0]),
        decay = BoxConstraint([0.0], [1.0]))
    )
)

# Simple hybrid dynamics: linear terms + neural interactions
function ar_step(layers, u, ps, t)
    # Linear terms from parameters
    params = layers.rates(ps.rates)
    growth = vcat(params.growth, -params.decay)

    # Neural network interactions
    interactions = layers.interaction(u, ps.interaction)

    return u .* (growth + interactions)
end

# Create autoregressive model
model = ARModel(
    (interaction = interaction_layer, rates = rate_params),
    ar_step;
    dt = 0.1)

ps, st = Lux.setup(Random.default_rng(), model)
tsteps = range(0, stop = 10.0, step = 0.1)
tspan = (tsteps[1], tsteps[end])

preds, _ = model(
    (; u0 = [1.0, 1.0],
        tspan,
        saveat = tsteps), ps, st)
size(preds)  # (2, 101)

We can predict batches of time series by providing a batch of initial conditions.

x = [(; u0 = rand(2),
    tspan = (tsteps[1], tsteps[end]),
    saveat = tsteps) for _ in 1:5]
batch_preds, _ = model(x, ps, st)
size(batch_preds)  # (2, 101, 5)

We may want to specify tspan, saveat or u0 once for all. This can be done when initiating the model. All hyperparameters are stored in the model's states.

# Create autoregressive model
model = ARModel(
    (interaction = interaction_layer, rates = rate_params),
    ar_step;
    dt = 0.1,
    tspan, 
    saveat = tsteps)
ps, st = Lux.setup(Random.default_rng(), model)
st.kwargs # (dt = 0.1, tspan = (0.0, 10.0), saveat = 0.0:0.1:10.0)
preds, _ = model((; u0 = [1.0, 1.0]), ps, st)

Specifying hyperparameters during model call will override the model's states.

We can bind our model with an additional layer predicting initial conditions from some other input data using the ICLayer.

ic_layer = ICLayer(Dense(10, 2, tanh))
model_with_ic = Chain(ic_layer, model)
ps, st = Lux.setup(Random.default_rng(), model_with_ic)
x = [(; u0 = rand(10),
         tspan = (tsteps[1], tsteps[end]),
         saveat = tsteps) for _ in 1:5]
batch_preds, _ = model_with_ic(x, ps, st)
size(batch_preds)  # (2, 101, 5)

A similar workwflow can be used with ODEModel and AnalyticModel.

Training with Optimisers.jl through the SGDBackend

⚠️ The default train function with the InferICs setup is opiniated and meant for demonstration purposes. You are encouraged to define your own training pipeline.

using Lux, Optimisers, ComponentArrays # conditional loading to use `SGDBackend`
using Zygote
data = rand(2, length(tsteps))
dataloader = SegmentedTimeSeries((data, tsteps); segment_length = 10, shift = 2)

backend = SGDBackend(Adam(1e-2), 100, AutoZygote(), MSELoss())
result = train(backend, model, dataloader, InferICs(false))

tspan = (tsteps[1], tsteps[end])
prediction, _ = model(
    (; u0 = result.ics[1].u0,
        tspan = tspan,
        saveat = tsteps), result.ps, result.st)

Bayesian inference with Turing.jl through the MCSamplingBackend

using Distributions, Turing, ComponentArrays # conditional loading to use `MCSamplingBackend`

rate_priors = (
    growth = arraydist([Normal(0.1, 0.05)]),
    decay = arraydist([Normal(0.05, 0.02)])
)
nn_priors = Normal(0, 1)  # Example prior for NN weights

# Create Bayesian model
bayesian_model = ARModel(
    (interaction = BayesianLayer(interaction_layer, nn_priors),
        rates = BayesianLayer(rate_params, rate_priors)),
    ar_step;
    dt = 0.1
)

datadistrib = Normal
mcmc_backend = MCSamplingBackend(NUTS(0.65), 500, datadistrib)
result = train(mcmc_backend,
    bayesian_model,
    dataloader,
    InferICs(false))

# Sample from posterior
chains = result.chains
posterior_samples = sample(bayesian_model, chains, 50)

📚 Documentation

Check out the tutorials and the API in the documentation.

🙏 Acknowledgments

Built on the excellent LuxDL, SciML and TuringLang ecosystems:

• Lux.jl for neural networks
• Turing.jl for Bayesian inference
• SciMLSensitivity.jl for automatic differentiation
• OrdinaryDiffEq.jl for differential equations