DiffEqFlux.jl

DiffEqFlux.jl (also known as DiffEq(For)Lux.jl) is a package designed to fuse differential equations with machine learning. It allows users to embed differential equation solvers directly into neural network architectures, facilitating research in Scientific Machine Learning (SciML).

It is built upon DifferentialEquations.jl and Lux.jl and provides architectures that match the interfaces of standard machine learning libraries, making it easy to integrate continuous-time machine learning layers into larger applications.

DiffEqFlux.jl is a high-level implicit deep learning library within the SciML ecosystem. It provides pre-built architectures and utility functions designed to mix differential equations with machine learning. It is specifically optimized for the easy and efficient training of Neural Ordinary Differential Equations (Neural ODEs) and their variants.

Note: DiffEqFlux.jl focuses on pre-built architectures. For low-level details regarding automatic differentiation of equation solvers, adjoint techniques, model calibration, nonlinear optimal control, or PDE-constrained optimization, refer to SciMLSensitivity.jl.

Smoothed collocation (also known as the two-stage method) is a technique for fitting differential equations to time series data without using a numerical differential equation solver.

How it works:

1. It builds a smoothed collocating polynomial from the data.
2. It uses this polynomial to estimate the true (u', u) pairs.
3. It estimates the residual u' - f(u, p, t) directly as a loss to determine the parameters p.

Pros and Cons:

• Pros: Extremely fast and robust to noise.
• Cons: Not as exact as other methods because it does not accumulate errors through time.

Note: For a more comprehensive set of collocation methods, refer to JuliaSimModelOptimizer.

A Neural Ordinary Differential Equation (Neural ODE) is an ODE where a neural network defines the derivative function ($u' = NN(u)$). In DiffEqFlux.jl, this is implemented using the NeuralODE struct. You can pass a neural network (such as a Lux.Chain) as the first argument to NeuralODE to define the dynamics.

To define a NeuralODE, you need:

1. A neural network structure (e.g., Lux.Chain).
2. A time span (tspan).
3. An ODE solver (e.g., Tsit5()).
4. An optional saveat argument to specify the time steps for output.

using Lux, DiffEqFlux, OrdinaryDiffEq

# 1. Define the neural network
dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))

# 2. Setup parameters and state
rng = Xoshiro(0)
p, st = Lux.setup(rng, dudt2)

# 3. Define the NeuralODE
tspan = (0.0f0, 1.5f0)
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)

DiffEqFlux.jl provides several continuous-time machine learning layers that match the interfaces of libraries like Flux.jl and Lux.jl. Supported architectures include:

• Neural Ordinary Differential Equations (Neural ODEs)
• Collocation-Based Neural ODEs (Fastest method; does not use a solver)
• Multiple Shooting Neural Ordinary Differential Equations
• Neural Stochastic Differential Equations (Neural SDEs)
• Neural Differential-Algebraic Equations (Neural DAEs)
• Neural Delay Differential Equations (Neural DDEs)
• Augmented Neural ODEs
• Hamiltonian Neural Networks (Includes specialized second-order and symplectic integrators)
• Continuous Normalizing Flows (CNF) and FFJORD

DiffEqFlux.jl provides several implicit layer architectures for continuous-time machine learning, including:

• Neural Ordinary Differential Equations (Neural ODEs)
• Collocation-Based Neural ODEs: Neural ODEs without a solver (optimized for speed).
• Multiple Shooting Neural Ordinary Differential Equations
• Neural Stochastic Differential Equations (Neural SDEs)
• Neural Differential-Algebraic Equations (Neural DAEs)
• Neural Delay Differential Equations (Neural DDEs)
• Augmented Neural ODEs
• Hamiltonian Neural Networks: Includes specialized second-order and symplectic integrators.
• Continuous Normalizing Flows (CNF) and FFJORD

These architectures support high-order, adaptive, implicit, GPU-accelerated, and Newton-Krylov methods.

When using NeuralODE within a Lux.Chain, you may need to manually bridge the initialparameters and initialstates calls if the NeuralODE wrapper does not automatically delegate them to the underlying model.

Note: This is a workaround for a known issue where Lux.setup might not correctly traverse the NeuralODE structure to find the parameters of the internal model.

import Lux: initialparameters, initialstates

# Manual delegation to the internal model
initialparameters(rng::AbstractRNG, node::NeuralODE) = initialparameters(rng, node.model)
initialstates(rng::AbstractRNG, node::NeuralODE) = initialstates(rng, node.model)

The multiple_shoot functionality in DiffEqFlux.jl is a specialized implementation designed for data shooting in implicit layer deep learning.

Key Assumptions & Constraints:

• It assumes full observability of the underlying dynamics.
• It assumes a lack of noise in the data.

If your use case requires a more general implementation of multiple shooting, or if you are performing parameter estimation against noisy data, do not use the functions in DiffEqFlux.jl. Instead, use the following specialized libraries:

• General Multiple Shooting: Use JuliaSimModelOptimizer.
• Parameter Estimation against Data: Use DiffEqParamEstim.jl.

The AugmentedNDELayer is used to improve the expressive power of Neural ODEs by augmenting the input with zeros. When using this layer, the input dimension to the underlying DE layer must be increased by the augment_dim.

For example, if your original input dimension is input_dim and you want to add an augmentation of size augment_dim, the internal Neural ODE layer should be constructed to accept an input of size input_dim + augment_dim.

# If input_dim is 2 and augment_dim is 1, the internal layer needs 3 inputs
input_dim = 2
augment_dim = 1

# The internal NeuralODE is constructed with the augmented dimension
node = NeuralODE(
    Chain(Dense(input_dim + augment_dim, hidden_dim, relu), ...),
    (0.0f0, 1.0f0),
    Tsit5()
)

# Wrap the node with the augmentation layer
node = AugmentedNDELayer(node, augment_dim)

DiffEqFlux.jl provides specialized helper functions for building neural differential equation architectures designed for Continuous Normalizing Flows (CNF). These layers are optimized for the task of density estimation.

Key components include:

• CNFLayer: A helper function to construct CNF-specific layers.
• FFJORD: An implementation of the Free-form Jacobian of Reversible Deterministic flows.
• FFJORDDistribution: A distribution wrapper for FFJORD-based models.

# Note: Specific API usage depends on the underlying DiffEqFlux implementation
# of CNFLayer, FFJORD, and FFJORDDistribution.