NeuralDELayer

Abstract interface for DiffEqFlux neural differential equation layers with a single Lux model field.

Interface

Concrete subtypes are Lux layers and are callable as:

solution, new_state = layer(x, ps, st)

where x is the initial condition or layer input, ps are Lux parameters, and st is Lux state. Implementations must return the SciML solution produced by solve and the updated Lux state. The wrapped model is stored in a field named model, matching the AbstractLuxWrapperLayer{:model} interface.

Rules

• The call must not mutate ps.
• Solver keyword arguments supplied to the constructor are forwarded to solve.
• State updates from the wrapped Lux model must be returned as the second tuple value.

Implementations

NeuralODE, NeuralCDDE, NeuralDAE, and NeuralODEMM implement this interface.

NeuralSDELayer

Abstract interface for DiffEqFlux neural stochastic differential equation layers with separate drift and diffusion Lux models.

Interface

Concrete subtypes are Lux container layers and are callable as:

solution, new_state = layer(x, ps, st)

where ps and st contain drift and diffusion fields. Implementations build an SDEProblem, call solve, and return the solution plus updated drift and diffusion states.

Rules

• drift(x, ps.drift) must return the deterministic drift vector.
• diffusion(x, ps.diffusion) must return the noise-rate object required by the concrete layer.
• Solver keyword arguments supplied to the constructor are forwarded to solve.

Implementations

NeuralDSDE implements diagonal noise. NeuralSDE implements a general noise-rate matrix with a fixed number of Brownian processes.

NeuralODE(model, tspan, alg = nothing, args...; kwargs...)

Constructs a continuous-time recurrent neural network, also known as a neural ordinary differential equation (neural ODE), with a fast gradient calculation via adjoints [1]. At a high level this corresponds to solving the forward differential equation, using a second differential equation that propagates the derivatives of the loss backwards in time.

Arguments

• model: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the ̇x.
• tspan: The timespan to be solved on.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl.
• sensealg: The choice of differentiation algorithm used in the backpropagation. Defaults to an adjoint method. See the Local Sensitivity Analysis documentation for more details.
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• model: Lux layer used as the ODE right-hand side.
• tspan: Integration time span.
• args: Positional solver arguments, usually including the ODE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralDELayer. Calling the layer as node(x, ps, st) returns (sol, new_state), where sol is the SciML ODE solution.

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
model = Lux.Chain(Lux.Dense(2 => 8, tanh), Lux.Dense(8 => 2))
node = NeuralODE(model, (0.0f0, 1.0f0); saveat = 0.1f0)
ps, st = Lux.setup(rng, node)
sol, st = node(Float32[1, 0], ps, st)

References

[1] Pontryagin, Lev Semenovich. Mathematical theory of optimal processes. CRC press, 1987.

NeuralDSDE(drift, diffusion, tspan, alg = nothing, args...; sensealg = TrackerAdjoint(),
    kwargs...)

Constructs a neural stochastic differential equation (neural SDE) with diagonal noise.

Arguments

• drift: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the drift function.
• diffusion: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the diffusion function. Should output a vector of the same size as the input.
• tspan: The timespan to be solved on.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl.
• sensealg: The choice of differentiation algorithm used in the backpropagation.
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• drift: Lux layer used as the SDE drift function.
• diffusion: Lux layer used as the diagonal diffusion function.
• tspan: Integration time span.
• args: Positional solver arguments, usually including the SDE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralSDELayer. Calling the layer as nsde(x, ps, st) returns (sol, new_state), where sol is the SciML SDE solution.

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
drift = Lux.Dense(2 => 2)
diffusion = Lux.Dense(2 => 2)
layer = NeuralDSDE(drift, diffusion, (0.0f0, 1.0f0))
ps, st = Lux.setup(rng, layer)
sol, st = layer(Float32[1, 0], ps, st)

NeuralSDE(drift, diffusion, tspan, nbrown, alg = nothing, args...;
    sensealg=TrackerAdjoint(), kwargs...)

Constructs a neural stochastic differential equation (neural SDE).

Arguments

• drift: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the drift function.
• diffusion: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the diffusion function. Should output a matrix that is nbrown x size(x, 1).
• tspan: The timespan to be solved on.
• nbrown: The number of Brownian processes.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl.
• sensealg: The choice of differentiation algorithm used in the backpropagation.
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• drift: Lux layer used as the SDE drift function.
• diffusion: Lux layer used as the full noise-rate function.
• tspan: Integration time span.
• nbrown: Number of Brownian processes.
• args: Positional solver arguments, usually including the SDE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralSDELayer. Calling the layer as nsde(x, ps, st) returns (sol, new_state), where sol is the SciML SDE solution.

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
drift = Lux.Dense(2 => 2)
diffusion = Lux.Dense(2 => 4)
layer = NeuralSDE(drift, diffusion, (0.0f0, 1.0f0), 2)
ps, st = Lux.setup(rng, layer)
sol, st = layer(Float32[1, 0], ps, st)

NeuralCDDE(model, tspan, hist, lags, alg = nothing, args...;
    sensealg = TrackerAdjoint(), kwargs...)

Constructs a neural delay differential equation (neural DDE) with constant delays.

Arguments

• model: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the derivative function. Should take an input of size [x; x(t - lag_1); ...; x(t - lag_n)] and produce and output shaped like x.
• tspan: The timespan to be solved on.
• hist: Defines the history function h(u, p, t) for values before the start of the integration. Note that u is supposed to be used to return a value that matches the size of u.
• lags: Defines the lagged values that should be utilized in the neural network.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl.
• sensealg: The choice of differentiation algorithm used in the backpropagation. Defaults to using reverse-mode automatic differentiation via Tracker.jl
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• model: Lux layer used as the delayed derivative model.
• tspan: Integration time span.
• hist: History function for times before the start of the integration.
• lags: Constant delays used by the DDE.
• args: Positional solver arguments, usually including the DDE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralDELayer. Calling the layer as ndde(x, ps, st) returns (sol, new_state), where sol is the SciML DDE solution.

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
model = Lux.Dense(4 => 2)
hist(u, p, t) = u
layer = NeuralCDDE(model, (0.0f0, 1.0f0), hist, [0.1f0])
ps, st = Lux.setup(rng, layer)
sol, st = layer(Float32[1, 0], ps, st)

NeuralDAE(model, constraints_model, tspan, args...; differential_vars = nothing,
    sensealg = TrackerAdjoint(), kwargs...)

Constructs a neural differential-algebraic equation (neural DAE).

Arguments

• model: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the derivative function. Should take an input of size x and produce the residual of f(dx,x,t) for only the differential variables.
• constraints_model: A function constraints_model(u,p,t) for the fixed constraints to impose on the algebraic equations.
• tspan: The timespan to be solved on.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl.
• sensealg: The choice of differentiation algorithm used in the backpropagation. Defaults to using reverse-mode automatic differentiation via Tracker.jl
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• model: Lux layer used for differential-variable residuals.
• constraints_model: Function returning residuals for algebraic constraints.
• tspan: Integration time span.
• args: Positional solver arguments, usually including the DAE algorithm.
• differential_vars: Boolean mask marking differential variables.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralDELayer. Calling the layer as ndae((u0, du0), ps, st) returns (sol, new_state), where sol is the SciML DAE solution.

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
model = Lux.Dense(4 => 1)
constraints(u, p, t) = [sum(u) - 1]
layer = NeuralDAE(model, constraints, (0.0f0, 1.0f0);
    differential_vars = [true, false])
ps, st = Lux.setup(rng, layer)
sol, st = layer((Float32[1, 0], Float32[0, 0]), ps, st)

NeuralODEMM(model, constraints_model, tspan, mass_matrix, alg = nothing, args...;
    sensealg = InterpolatingAdjoint(autojacvec = ZygoteVJP()), kwargs...)

Constructs a physically-constrained continuous-time recurrent neural network, also known as a neural differential-algebraic equation (neural DAE), with a mass matrix and a fast gradient calculation via adjoints [1]. The mass matrix formulation is:

\[Mu' = f(u,p,t)\]

where M is semi-explicit, i.e. singular with zeros for rows corresponding to the constraint equations.

Arguments

• model: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the ̇f(u,p,t)
• constraints_model: A function constraints_model(u,p,t) for the fixed constraints to impose on the algebraic equations.
• tspan: The timespan to be solved on.
• mass_matrix: The mass matrix associated with the DAE.
• alg: The algorithm used to solve the ODE. Defaults to nothing, i.e. the default algorithm from DifferentialEquations.jl. This method requires an implicit ODE solver compatible with singular mass matrices. Consult the DAE solvers documentation for more details.
• sensealg: The choice of differentiation algorithm used in the backpropagation. Defaults to an adjoint method. See the Local Sensitivity Analysis documentation for more details.
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• model: Lux layer used for the differential rows of f(u, p, t).
• constraints_model: Function returning algebraic constraint rows.
• tspan: Integration time span.
• mass_matrix: Mass matrix passed to ODEFunction.
• args: Positional solver arguments, usually including the implicit ODE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A NeuralDELayer. Calling the layer as node(x, ps, st) returns (sol, new_state), where sol is the SciML ODE solution for the mass-matrix formulation.

Examples

using DiffEqFlux, Lux, Random, LinearAlgebra

rng = Random.default_rng()
model = Lux.Dense(2 => 1)
constraints(u, p, t) = [sum(u) - 1]
mass_matrix = Diagonal([1.0f0, 0.0f0])
layer = NeuralODEMM(model, constraints, (0.0f0, 1.0f0), mass_matrix)
ps, st = Lux.setup(rng, layer)
sol, st = layer(Float32[1, 0], ps, st)

AugmentedNDELayer(model, adim::Int)

Constructs an Augmented Neural Differential Equation Layer.

Arguments

• model: Any Neural Differential Equation Layer.
• adim: The number of dimensions the initial conditions should be lifted.

Returns

A Lux.Chain that first augments the input with adim zero-valued dimensions and then calls nde.

Examples

using DiffEqFlux, Lux

nde = NeuralODE(Lux.Dense(3 => 3), (0.0f0, 1.0f0))
augmented = AugmentedNDELayer(nde, 1)

References

[1] Dupont, Emilien, Arnaud Doucet, and Yee Whye Teh. "Augmented neural ODEs." In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 3140-3150. 2019.

DimMover(; from = -2, to = -1)

Constructs a Dimension Mover Layer.

We can have Lux's conventional order (data, channel, batch) by using it as the last layer of AbstractLuxLayer to swap the batch-index and the time-index of the Neural DE's output considering that each time point is a channel.

Keywords

• from: Source dimension. Negative values are counted from the end, so -2 refers to the second-to-last dimension.
• to: Destination dimension. Negative values are counted from the end.

Fields

• from: Stored source dimension.
• to: Stored destination dimension.

Returns

A Lux layer. Calling DimMover(; from, to)(x, ps, st) returns (moved_x, st).

Examples

using DiffEqFlux

layer = DimMover(; from = -2, to = -1)
y, st = layer(rand(2, 3, 4), nothing, NamedTuple())