# Neural Differential Equation Layer Functions

The following layers are helper functions for easily building neural differential equation architectures in the currently most efficient way. As demonstrated in the tutorials, they do not have to be used since automatic differentiation will just work over solve, but these cover common use cases and choose what's known to be the optimal mode of AD for the respective equation type.

DiffEqFlux.NeuralODEType

Constructs a continuous-time recurrant 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.

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

Arguments:

• model: A Chain or FastChain 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 algorthm used in the backpropogation. Defaults to an adjoint method, and with FastChain it defaults to utilizing a tape-compiled ReverseDiff vector-Jacobian product for extra efficiency. Seee 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.

References:

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

DiffEqFlux.NeuralDSDEType

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

NeuralDSDE(model1,model2,tspan,alg=nothing,args...;
NeuralDSDE(model1::FastChain,model2::FastChain,tspan,alg=nothing,args...;
sensealg=TrackerAdjoint(),kwargs...)

Arguments:

• model1: A Chain or FastChain neural network that defines the drift function.
• model2: A Chain or FastChain 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 algorthm used in the backpropogation. 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.
DiffEqFlux.NeuralSDEType

Constructs a neural stochastic differential equation (neural SDE).

NeuralSDE(model1,model2,tspan,nbrown,alg=nothing,args...;
NeuralSDE(model1::FastChain,model2::FastChain,tspan,nbrown,alg=nothing,args...;
sensealg=TrackerAdjoint(),kwargs...)

Arguments:

• model1: A Chain or FastChain neural network that defines the drift function.
• model2: A Chain or FastChain 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 algorthm used in the backpropogation. 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.
DiffEqFlux.NeuralCDDEType

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

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

Arguments:

• model: A Chain or FastChain 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(t) for values before the start of the integration.
• 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 algorthm used in the backpropogation. 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.
DiffEqFlux.NeuralDAEType

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

NeuralDAE(model,constraints_model,tspan,alg=nothing,args...;
NeuralDAE(model::FastChain,constraints_model,tspan,alg=nothing,args...;
sensealg=TrackerAdjoint(),kwargs...)

Arguments:

• model: A Chain or FastChain 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 constaints 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 algorthm used in the backpropogation. 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.
DiffEqFlux.NeuralODEMMType

Constructs a physically-constrained continuous-time recurrant 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.

NeuralODEMM(model,constraints_model,tspan,mass_matrix,alg=nothing,args...;kwargs...)
NeuralODEMM(model::FastChain,tspan,mass_matrix,alg=nothing,args...;
kwargs...)

Arguments:

• model: A Chain or FastChain neural network that defines the ̇f(u,p,t)
• constraints_model: A function constraints_model(u,p,t) for the fixed constaints 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 algorthm used in the backpropogation. Defaults to an adjoint method, and with FastChain it defaults to utilizing a tape-compiled ReverseDiff vector-Jacobian product for extra efficiency. Seee 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.
DiffEqFlux.AugmentedNDELayerType

Constructs an Augmented Neural Differential Equation Layer.

AugmentedNDELayer(nde, adim::Int)

Arguments:

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

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.

# Helper Layer Functions

DiffEqFlux.FluxBatchOrderFunction

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

julia FluxBatchOrder = DimMover(-2, -1)`