### Full Neural ODE Training Example on GPU in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/GPUs.md

This comprehensive example demonstrates the full workflow of defining, training, and visualizing a Neural ODE on a GPU. It includes generating synthetic data, setting up a `Lux.Chain` model, defining a loss function, and using `Optimization.jl` with `Zygote.jl` for training. The `gpu_device` function ensures compatibility across CPU and GPU environments.

```Julia
using Lux, Optimization, OptimizationOptimisers, Zygote, OrdinaryDiffEq, Plots, LuxCUDA,
      SciMLSensitivity, Random, ComponentArrays
import DiffEqFlux: NeuralODE
const cdev = cpu_device()
const gdev = gpu_device()

CUDA.allowscalar(false) # Makes sure no slow operations are occurring

#rng for Lux.setup
rng = Xoshiro(0)
# Generate Data
u0 = Float32[2.0; 0.0]
datasize = 30
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)
function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end
prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
# Make the data into a GPU-based array if the user has a GPU
ode_data = solve(prob_trueode, Tsit5(); saveat = tsteps)
ode_data = Array(ode_data) |> gdev

dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
u0 = Float32[2.0; 0.0] |> gdev
p, st = Lux.setup(rng, dudt2)
p = p |> ComponentArray |> gdev
st = st |> gdev

prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)

predict_neuralode(p) = reduce(hcat, first(prob_neuralode(u0, p, st)).u)
function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss
end
# Callback function to observe training
list_plots = []
iter = 0
callback = function (state, l; doplot = false)
    p = state.u
    global list_plots, iter
    pred = predict_neuralode(p)
    if iter == 0
        list_plots = []
    end
    iter += 1
    display(l)
    # plot current prediction against data
    plt = scatter(tsteps, Array(ode_data[1, :]); label = "data")
    scatter!(plt, tsteps, Array(pred[1, :]); label = "prediction")
    push!(list_plots, plt)
    if doplot
        display(plot(plt))
    end
    return false
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)
result_neuralode = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.05); callback, maxiters = 300)
```

--------------------------------

### Setting Up Lux Model Parameters for Training in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_gde.md

This snippet initializes the model's parameters and states using Lux's `setup` function. It sets a random seed for reproducibility, converts the parameters to a `ComponentArray` for compatibility with sensitivity analysis tools, and moves both parameters and states to the specified computation device (CPU/GPU).

```Julia
rng = Random.default_rng()
Random.seed!(rng, 0)

ps, st = Lux.setup(rng, model)
ps = ComponentArray(ps) |> device
st = st |> device
```

--------------------------------

### Defining and Training a FFJORD Model in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/normalizing_flows.md

This comprehensive example demonstrates the full workflow for Continuous Normalizing Flows (CNF) using DiffEqFlux.jl. It covers model definition with a Chain neural network, FFJORD layer setup, parameter initialization, training with Adam and LBFGS optimizers, and evaluation of the learned density against the actual distribution. It also shows how to generate new data from the learned distribution.

```Julia
using ComponentArrays, DiffEqFlux, OrdinaryDiffEq, Optimization, Distributions, Random,
      OptimizationOptimisers, OptimizationOptimJL

nn = Chain(Dense(1, 3, tanh), Dense(3, 1, tanh))
tspan = (0.0f0, 10.0f0)

ffjord_mdl = FFJORD(nn, tspan, (1,), Tsit5(); ad = AutoZygote())
ps, st = Lux.setup(Xoshiro(0), ffjord_mdl)
ps = ComponentArray(ps)
model = StatefulLuxLayer{true}(ffjord_mdl, nothing, st)

# Training
data_dist = Normal(6.0f0, 0.7f0)
train_data = Float32.(rand(data_dist, 1, 100))

function loss(θ)
    logpx, λ₁, λ₂ = model(train_data, θ)
    return -mean(logpx)
end

function cb(state, l)
    @info "FFJORD Training" loss=l
    return false
end

adtype = Optimization.AutoForwardDiff()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ps)

res1 = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.01); maxiters = 20, callback = cb)

optprob2 = Optimization.OptimizationProblem(optf, res1.u)
res2 = Optimization.solve(optprob2, Optim.LBFGS(); allow_f_increases = false, callback = cb)

# Evaluation
using Distances

st_ = (; st..., monte_carlo = false)

actual_pdf = pdf.(data_dist, train_data)
learned_pdf = exp.(ffjord_mdl(train_data, res2.u, st_)[1][1])
train_dis = totalvariation(learned_pdf, actual_pdf) / size(train_data, 2)

# Data Generation
ffjord_dist = FFJORDDistribution(ffjord_mdl, ps, st)
new_data = rand(ffjord_dist, 100)
```

--------------------------------

### Setting Up Environment and Defining True ODE in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/collocation.md

This snippet initializes the Julia environment by importing necessary packages for differential equations, neural networks, optimization, and plotting. It then defines the initial conditions, data size, time span, and the `trueODEfunc` which represents the ground truth ODE system used to generate synthetic data for the example.

```Julia
using ComponentArrays, Lux, DiffEqFlux, OrdinaryDiffEq, SciMLSensitivity, Optimization,
      OptimizationOptimisers, Plots

using Random
rng = Xoshiro(0)

u0 = Float32[2.0; 0.0]
datasize = 300
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end
```

--------------------------------

### Loading Core Packages and Device Setup (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/augmented_neural_ode.md

This snippet loads essential Julia packages required for implementing and training Neural Ordinary Differential Equations, including DiffEqFlux for neural ODEs, OrdinaryDiffEq for ODE solvers, LuxCUDA for GPU acceleration, and Optimization.jl for the training loop. It also initializes `cpu_device()` and `gpu_device()` for explicit device placement.

```Julia
using DiffEqFlux, OrdinaryDiffEq, Statistics, LinearAlgebra, Plots, LuxCUDA, Random
using MLUtils, ComponentArrays
using Optimization, OptimizationOptimisers, IterTools

const cdev = cpu_device()
const gdev = gpu_device()
```

--------------------------------

### Integrating Lux.Chain with GPU Arrays for ODEs in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/GPUs.md

This example demonstrates how to define a `Lux.Chain` neural network and ensure that its parameters and state are moved to the GPU along with the initial condition. This setup allows the `dudt2_` function, which wraps the `Lux.Chain` model, to perform computations entirely on the GPU when solving the `ODEProblem`.

```Julia
dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))

u0 = Float32[2.0; 0.0] |> gdev
p, st = Lux.setup(rng, dudt2) |> gdev

dudt2_(u, p, t) = first(dudt2(u, p, st))

# Simulation interval and intermediary points
tspan = (0.0f0, 10.0f0)
tsteps = 0.0f0:1.0f-1:10.0f0

prob_gpu = ODEProblem(dudt2_, u0, tspan, p)

# Runs on a GPU
sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)
```

--------------------------------

### Initializing Parameters and Testing Callback Function (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode.md

This snippet initializes the neural ODE's parameters using `ComponentArray` from an existing `p` variable. It then immediately invokes the defined `callback` function with these initial parameters and their corresponding loss, primarily to check its functionality and display the starting loss value without plotting.

```Julia
pinit = ComponentArray(p)
callback((; u = pinit), loss_neuralode(pinit))
```

--------------------------------

### Example of Lux.Chain for Neural ODE - Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode.md

This snippet provides a standalone example of defining a neural network using `Lux.Chain`, highlighting how Lux.jl's `Chain` can be directly used as the derivative function within a `NeuralODE`. It emphasizes the flexibility of incorporating structured assumptions into the model.

```Julia
dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
```

--------------------------------

### Defining and Training a Neural ODE (Full Example) - Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode.md

This comprehensive example demonstrates the full workflow for defining, training, and optimizing a Neural Ordinary Differential Equation (Neural ODE) in Julia. It covers data generation from a true ODE, setting up a Lux.jl neural network, defining the NeuralODE problem, creating prediction and loss functions, implementing a callback for monitoring, and using Optimization.jl with Adam and BFGS optimizers.

```Julia
using ComponentArrays, Lux, DiffEqFlux, OrdinaryDiffEq, Optimization, OptimizationOptimJL,
      OptimizationOptimisers, Random, Plots

rng = Xoshiro(0)
u0 = Float32[2.0; 0.0]
datasize = 30
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end

prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
ode_data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps))

dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
p, st = Lux.setup(rng, dudt2)
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)

function predict_neuralode(p)
    Array(prob_neuralode(u0, p, st)[1])
end

function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss
end

# Do not plot by default for the documentation
# Users should change doplot=true to see the plots callbacks
function callback(state, l; doplot = false)
    println(l)
    # plot current prediction against data
    if doplot
        pred = predict_neuralode(state.u)
        plt = scatter(tsteps, ode_data[1, :]; label = "data")
        scatter!(plt, tsteps, pred[1, :]; label = "prediction")
        display(plot(plt))
    end
    return false
end

pinit = ComponentArray(p)
callback((; u = pinit), loss_neuralode(pinit); doplot = true)

# use Optimization.jl to solve the problem
adtype = Optimization.AutoZygote()

optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pinit)

result_neuralode = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.05); callback = callback, maxiters = 300)

optprob2 = remake(optprob; u0 = result_neuralode.u)

result_neuralode2 = Optimization.solve(
    optprob2, Optim.BFGS(; initial_stepnorm = 0.01); callback, allow_f_increases = false)

callback((; u = result_neuralode2.u), loss_neuralode(result_neuralode2.u); doplot = true)
```

--------------------------------

### Using DiffEqFlux NeuralODE Layer for GPU in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/GPUs.md

This code snippet shows a more concise way to define a neural ODE using the `DiffEqFlux.NeuralODE` layer function. It directly wraps the previously defined `model` and `tspan` into a `NeuralODE` object, simplifying the setup for GPU-accelerated neural ODEs.

```Julia
using DiffEqFlux: NeuralODE
prob_neuralode_gpu = NeuralODE(model, tspan, Tsit5(); saveat = tsteps)
```

--------------------------------

### Visualizing Initial Neural SDE Prediction (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_sde.md

This snippet evaluates the untrained `neuralsde` model and visualizes its initial prediction against the true SDE data. It first gets a single prediction, then reconstructs the SDE problem from the neural network components for ensemble simulation. An ensemble of 100 trajectories is run, and the mean/variance are summarized and plotted, along with a direct comparison of the first component of the prediction versus data.

```Julia
# Get the prediction using the correct initial condition
prediction0 = neuralsde(u0, ps, st)[1]

drift_model = StatefulLuxLayer{true}(drift_dudt, ps.drift, st.drift)
diffusion_model = StatefulLuxLayer{true}(diffusion_dudt, ps.diffusion, st.diffusion)

drift_(u, p, t) = drift_model(u, p.drift)
diffusion_(u, p, t) = diffusion_model(u, p.diffusion)

prob_neuralsde = SDEProblem(drift_, diffusion_, u0, (0.0f0, 1.2f0), ps)

ensemble_nprob = EnsembleProblem(prob_neuralsde; safetycopy = false)
ensemble_nsol = solve(ensemble_nprob, SOSRI(); trajectories = 100, saveat = tsteps)
ensemble_nsum = EnsembleSummary(ensemble_nsol)

plt1 = plot(ensemble_nsum; title = "Neural SDE: Before Training")
scatter!(plt1, tsteps, sde_data'; lw = 3)

scatter(tsteps, sde_data[1, :]; label = "data")
scatter!(tsteps, prediction0[1, :]; label = "prediction")
```

--------------------------------

### Training Hamiltonian Neural Network for 1D Spring-Mass System (Full Example) - Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/hamiltonian_nn.md

This comprehensive Julia code snippet demonstrates the end-to-end process of training a Hamiltonian Neural Network (HNN) to model a 1D spring-mass system. It includes data generation for position, momentum, and their derivatives, setting up the HNN with Lux.jl, defining a loss function, and using Optimization.jl to train the model. Finally, it visualizes the learned trajectory against the original data using NeuralODE.

```Julia
using Lux, DiffEqFlux, OrdinaryDiffEq, Statistics, Plots, Zygote, ForwardDiff, Random,
      ComponentArrays, Optimization, OptimizationOptimisers, MLUtils

t = range(0.0f0, 1.0f0; length = 1024)
π_32 = Float32(π)
q_t = reshape(sin.(2π_32 * t), 1, :)
p_t = reshape(cos.(2π_32 * t), 1, :)
dqdt = 2π_32 .* p_t
dpdt = -2π_32 .* q_t

data = cat(q_t, p_t; dims = 1)
target = cat(dqdt, dpdt; dims = 1)
B = 256
NEPOCHS = 500
dataloader = DataLoader((data, target); batchsize = B)

hnn = Layers.HamiltonianNN{true}(Layers.MLP(2, (1028, 1)); autodiff = AutoZygote())
ps, st = Lux.setup(Xoshiro(0), hnn)
ps_c = ps |> ComponentArray

opt = OptimizationOptimisers.Adam(0.01f0)

function loss_function(ps, databatch)
    data, target = databatch
    pred, st_ = hnn(data, ps, st)
    return mean(abs2, pred .- target)
end

function callback(state, loss)
    println("[Hamiltonian NN] Loss: ", loss)
    return false
end

opt_func = OptimizationFunction(loss_function, Optimization.AutoForwardDiff())
opt_prob = OptimizationProblem(opt_func, ps_c, dataloader)

res = Optimization.solve(opt_prob, opt; callback, epochs = NEPOCHS)

ps_trained = res.u

model = NeuralODE(
    hnn, (0.0f0, 1.0f0), Tsit5(); save_everystep = false, save_start = true, saveat = t)

pred = Array(first(model(data[:, 1], ps_trained, st)))
plot(data[1, :], data[2, :]; lw = 4, label = "Original")
plot!(pred[1, :], pred[2, :]; lw = 4, label = "Predicted")
xlabel!("Position (q)")
ylabel!("Momentum (p)")
```

--------------------------------

### Setting Up Training Data Environment (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_sde.md

This snippet initializes the environment for building training data for a neural SDE. It loads necessary Julia packages like `Plots`, `Statistics`, `DiffEqFlux`, and `StochasticDiffEq`. It then defines the initial condition `u0`, the `datasize`, the `tspan` (time span), and `tsteps` (time points) for data generation.

```Julia
using Plots, Statistics, ComponentArrays, Optimization, OptimizationOptimisers, DiffEqFlux,
      StochasticDiffEq, SciMLBase.EnsembleAnalysis, Random

u0 = Float32[2.0; 0.0]
datasize = 30
tspan = (0.0f0, 1.0f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)
```

--------------------------------

### Refining CNF Optimization with LBFGS in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/normalizing_flows.md

Following the initial Adam optimization, this snippet continues training the Continuous Normalizing Flow model using the LBFGS optimizer. It starts from the parameters found by Adam (res1.u) to further refine the optimization, aiming for a more precise minimum.

```Julia
optprob2 = Optimization.OptimizationProblem(optf, res1.u)
res2 = Optimization.solve(optprob2, Optim.LBFGS(); allow_f_increases = false, callback = cb)
```

--------------------------------

### Constructing Neural ODE and Full Model (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/mnist_conv_neural_ode.md

Initializes the `NeuralODE` using the `dudt` chain, `Tsit5` solver, and `BacksolveAdjoint` for sensitivity analysis. It then combines `down`, `nn_ode`, a `solution_to_array` helper, and `fc` into the complete `m` model. Parameters are set up using `Lux.setup` and moved to the GPU for efficient computation.

```Julia
nn_ode = NeuralODE(dudt, (0.0f0, 1.0f0), Tsit5(); save_everystep = false,
    sensealg = BacksolveAdjoint(; autojacvec = ZygoteVJP()),
    reltol = 1e-5, abstol = 1e-6, save_start = false)

solution_to_array(sol) = sol.u[end]

m = Chain(
    down,
    nn_ode,
    solution_to_array,
    fc
)

ps, st = Lux.setup(Xoshiro(0), m);
ps = ComponentArray(ps) |> gdev;
st = st |> gdev;

# To understand the intermediate NN-ODE layer, we can examine it's dimensionality
img, lab = first(dataloader);
x_m, _ = m(img, ps, st);
```

--------------------------------

### Displaying Direct Package Dependencies (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/index.md

This snippet uses Julia's Pkg module to display the status of direct package dependencies, providing a concise list of packages and their versions used to build the documentation. The `# hide` comments indicate these lines are hidden in the final documentation output.

```Julia
using Pkg # hide
Pkg.status() # hide
```

--------------------------------

### Displaying Julia Version and System Information (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/index.md

This snippet utilizes Julia's InteractiveUtils module to print detailed information about the Julia version, operating system, and hardware, which is crucial for ensuring reproducibility across different environments. The `# hide` comments ensure these lines are not visible in the final documentation.

```Julia
using InteractiveUtils # hide
versioninfo() # hide
```

--------------------------------

### Defining Prediction and Loss Functions for Neural ODE - Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode.md

This snippet defines the `predict_neuralode` function to get predictions from the Neural ODE model and the `loss_neuralode` function, which calculates the L2 loss between the model's predictions and the true ODE data. These functions are crucial for the optimization process.

```Julia
function predict_neuralode(p)
    Array(prob_neuralode(u0, p, st)[1])
end

function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss
end
```

--------------------------------

### Downloading and Loading Weather Data in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode_weather_forecast.md

This snippet defines a `download_data` function to fetch daily weather data from a specified URL, load it into a DataFrame using CSV.jl, and concatenate training and testing datasets. It then calls this function to load the data and displays the first few rows for initial inspection.

```Julia
using Random, Dates, Optimization, ComponentArrays, Lux, OptimizationOptimisers, DiffEqFlux,
      OrdinaryDiffEq, CSV, DataFrames, Dates, Statistics, Plots, DataDeps

function download_data(
        data_url = "https://raw.githubusercontent.com/SebastianCallh/neural-ode-weather-forecast/master/data/",
        data_local_path = "./delhi")
    function load(file_name)
        data_dep = DataDep("delhi/train", "", "$data_url/$file_name")
        Base.download(data_dep, data_local_path; i_accept_the_terms_of_use = true)
        CSV.read(joinpath(data_local_path, file_name), DataFrame)
    end

    train_df = load("DailyDelhiClimateTrain.csv")
    test_df = load("DailyDelhiClimateTest.csv")
    return vcat(train_df, test_df)
end

df = download_data()
first(df, 5) # hide
```

--------------------------------

### Generating Dynamic Links to Project and Manifest Files (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/index.md

This Julia snippet reads `Project.toml` to extract the package version and name, then constructs dynamic GitHub links to the `Manifest.toml` and `Project.toml` files. It uses the `Markdown` module to format these links into a user-friendly message, allowing users to easily download the full dependency files for reproducibility.

```Julia
using TOML
using Markdown
version = TOML.parse(read("../../Project.toml", String))["version"]
name = TOML.parse(read("../../Project.toml", String))["name"]
link_manifest = "https://github.com/SciML/" *
                name *
                ".jl/tree/gh-pages/v" *
                version *
                "/assets/Manifest.toml"
link_project = "https://github.com/SciML/" *
               name *
               ".jl/tree/gh-pages/v" *
               version *
               "/assets/Project.toml"
Markdown.parse("""You can also download the
[manifest]($link_manifest)
file and the
[project]($link_project)
file.
""")
```

--------------------------------

### Defining Classification and Accuracy Metrics (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/mnist_conv_neural_ode.md

Defines a `classify` helper function to get the predicted class from model outputs and an `accuracy` function to compute the model's accuracy over a given number of batches. The `accuracy` function moves data to the CPU for comparison and sets the model to test mode.

```Julia
classify(x) = argmax.(eachcol(x))

function accuracy(model, data, ps, st; n_batches = 10)
    total_correct = 0
    total = 0
    st = Lux.testmode(st)
    for (x, y) in collect(data)[1:min(n_batches, length(data))]
        target_class = classify(cdev(y))
        predicted_class = classify(cdev(first(model(x, ps, st))))
        total_correct += sum(target_class .== predicted_class)
        total += length(target_class)
    end
    return total_correct / total
end

# burn in accuracy
accuracy(m, ((img, lab),), ps, st)
```

--------------------------------

### Initializing TensorLayer for Legendre Basis Expansion in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/tensor_layer.md

This code initializes a TensorProductLayer from DiffEqFlux.Layers, configured for 10th order Legendre Basis expansions. It sets up the neural network (nn) and its parameters (ps, st) using Lux.setup, then wraps it in a StatefulLuxLayer for stateful operations, preparing it for integration into the physics-informed model.

```Julia
A = [Basis.Legendre(10), Basis.Legendre(10)]
nn = Layers.TensorProductLayer(A, 1)
ps, st = Lux.setup(Xoshiro(0), nn)
ps = ComponentArray(ps)
nn = StatefulLuxLayer{true}(nn, nothing, st)
```

--------------------------------

### Generating Smoothed Collocation Data in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/collocation.md

This snippet demonstrates how to generate synthetic noisy data from a true ODE function and then use `collocate_data` with an `EpanechnikovKernel` to smooth this data and estimate the underlying dynamics (`u`) and their derivatives (`du`). It initializes the ODE problem, solves it to get data, adds noise, and then performs the collocation.

```Julia
using ComponentArrays, Lux, DiffEqFlux, Optimization, OptimizationOptimisers,
      OrdinaryDiffEq, Plots

using Random
rng = Xoshiro(0)

u0 = Float32[2.0; 0.0]
datasize = 300
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end

prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps)) .+ 0.1randn(2, 300)

du, u = collocate_data(data, tsteps, EpanechnikovKernel())

scatter(tsteps, data')
plot!(tsteps, u'; lw = 5)
```

--------------------------------

### Setting Up Multiple Shooting for Neural ODE in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/multiple_shooting.md

This Julia code sets up a demonstration of Multiple Shooting for a Neural Ordinary Differential Equation (NODE). It defines initial conditions, time steps, and generates synthetic ODE data using a 'trueODEfunc'. It then defines a Lux neural network, initializes its parameters, and constructs a 'NeuralODE' problem, preparing for training with multiple shooting. Key dependencies include ComponentArrays, Lux, DiffEqFlux, Optimization, OptimizationPolyalgorithms, OrdinaryDiffEq, and Plots.

```Julia
using ComponentArrays, Lux, DiffEqFlux, Optimization, OptimizationPolyalgorithms,
      OrdinaryDiffEq, Plots
using DiffEqFlux: group_ranges

using Random
rng = Xoshiro(0)

# Define initial conditions and time steps
datasize = 30
u0 = Float32[2.0, 0.0]
tspan = (0.0f0, 5.0f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)

# Get the data
function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end
prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
ode_data = Array(solve(prob_trueode, Tsit5(); saveat = tsteps))

# Define the Neural Network
nn = Chain(x -> x .^ 3, Dense(2, 16, tanh), Dense(16, 2))
p_init, st = Lux.setup(rng, nn)

ps = ComponentArray(p_init)
pd, pax = getdata(ps), getaxes(ps)

neuralode = NeuralODE(nn, tspan, Tsit5(); saveat = tsteps)
prob_node = ODEProblem((u, p, t) -> nn(u, p, st)[1], u0, tspan, ComponentArray(p_init))
```

--------------------------------

### Training Augmented Neural ODE Model in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/augmented_neural_ode.md

This snippet shows the training configuration for an Augmented Neural ODE, which is similar to the standard Neural ODE but with an augmented input (e.g., a single zero). This augmentation increases the problem's dimensionality, allowing the neural ODE to express more complex functions. The setup uses the same `OptimizationFunction` and `OptimizationProblem` structure.

```Julia
model, ps, st = construct_model(1, 2, 64, 1)

optfunc = OptimizationFunction(
    (x, data) -> loss_node(model, data, x, st),
    Optimization.AutoZygote())
optprob = OptimizationProblem(optfunc, ComponentArray(ps |> cdev) |> gdev, dataloader)
res = solve(optprob, opt; callback = cb, epochs = 100)

plot_contour(model, res.u, st)
```

--------------------------------

### Loading Required Packages for DAE Modeling in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/physical_constraints.md

This snippet loads the essential Julia packages required for defining, solving, and training differential-algebraic equations (DAEs) with neural networks. Key packages include `DiffEqFlux` for neural differential equations, `Lux` for neural network layers, `ComponentArrays` for parameter handling, `Optimization` and `OptimizationOptimJL` for optimization, `OrdinaryDiffEq` for ODE/DAE solvers, and `Plots` for visualization. `Random` is used for neural network initialization.

```Julia
using DiffEqFlux
using Lux, ComponentArrays, Optimization, OptimizationOptimJL, OrdinaryDiffEq, Plots

using Random
rng = Random.default_rng()
```

--------------------------------

### Displaying Full Package Manifest Status (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/index.md

This snippet uses Julia's Pkg module with `mode = PKGMODE_MANIFEST` to display the complete manifest of all package dependencies, including transitive ones, providing a comprehensive overview of the entire dependency tree. The `# hide` comments indicate these lines are hidden in the final documentation output.

```Julia
using Pkg # hide
Pkg.status(; mode = PKGMODE_MANIFEST) # hide
```

--------------------------------

### Resuming Neural SDE Training with Increased 'n' in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_sde.md

This code resumes the neural SDE training, but with an increased `n` value of 100 for the `loss_neuralsde` function and a reduced learning rate of 0.001. It continues training for another 100 iterations, starting from the parameters obtained from the first training phase (`result1.u`), to refine the model's behavior.

```Julia
opt = OptimizationOptimisers.Adam(0.001)
optf2 = Optimization.OptimizationFunction((x, p) -> loss_neuralsde(x; n = 100), adtype)
optprob2 = Optimization.OptimizationProblem(optf2, result1.u)
result2 = Optimization.solve(optprob2, opt; callback, maxiters = 100)
```

--------------------------------

### Implementing Single Shooting Optimization in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/multiple_shooting.md

This snippet demonstrates the implementation of single shooting, which is a special case of multiple shooting where `group_size` equals the total data size. It reuses the `callback` and `loss_function` from the multiple shooting setup but defines a new `loss_single_shooting` function. The optimization problem is then configured and solved similarly to the multiple shooting case, resulting in a GIF illustrating the single shooting optimization.

```Julia
anim = Plots.Animation()
iter = 0
group_size = 30

ps = ComponentArray(p_init)
pd, pax = getdata(ps), getaxes(ps)

function loss_single_shooting(p)
    ps = ComponentArray(p, pax)
    loss, currpred = multiple_shoot(ps, ode_data, tsteps, prob_node, loss_function,
        Tsit5(), group_size; continuity_term)
    global preds = currpred
    return loss
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_single_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pd)
res_ms = Optimization.solve(optprob, PolyOpt(); callback = callback, maxiters = 300)
gif(anim, "single_shooting.gif"; fps = 15)
```

--------------------------------

### Pretraining Neural Network with Smoothed Collocation Data in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/collocation.md

This section sets up and executes the pretraining phase of the neural network. It initializes the network parameters, defines a no-op callback, and configures an `OptimizationProblem` using `AutoZygote` for automatic differentiation. The network is then trained using the Adam optimizer against the collocated data (`du`, `u`) to quickly find a good initial parameter set before full ODE integration.

```Julia
pinit, st = Lux.setup(rng, dudt2)

callback = function (p, l)
    return false
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))

result_neuralode = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.05); callback, maxiters = 10000)

prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
nn_sol, st = prob_neuralode(u0, result_neuralode.u, st)
scatter(tsteps, data')
plot!(nn_sol)
savefig("colloc_trained.png")
```

--------------------------------

### Defining Callback Function for Training in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/physical_constraints.md

This snippet defines a `callback` function used to observe the training process. It takes the current `state` and `loss` as arguments, displays the current loss, and returns `false` to indicate that the optimization should continue. This function is typically passed to the optimization solver to provide real-time feedback.

```Julia
callback = function (state, l) #callback function to observe training
    display(l)
    return false
end
```

--------------------------------

### Initializing Adam Optimizer in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_gde.md

This snippet initializes the Adam optimizer from the `Optimisers.jl` package with a learning rate of `0.01`. It then sets up the optimizer state (`st_opt`) for the model parameters (`ps`), preparing it for the training process.

```Julia
opt = Optimisers.Adam(0.01f0)
st_opt = Optimisers.setup(opt, ps)
```

--------------------------------

### Defining and Training a Neural DAE Model in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/physical_constraints.md

This comprehensive snippet illustrates the full workflow for defining a differential-algebraic equation (DAE) with a singular mass matrix and subsequently training a `NeuralODEMM` model to fit data generated by this DAE. It encompasses the DAE function definition, problem setup, solution, neural network architecture, prediction function, loss calculation, and optimization process using `Optimization.jl`.

```Julia
using DiffEqFlux
using Lux, ComponentArrays, Optimization, OptimizationOptimJL, OrdinaryDiffEq, Plots

using Random
rng = Random.default_rng()

function f!(du, u, p, t)
    y₁, y₂, y₃ = u
    k₁, k₂, k₃ = p
    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
    du[2] = k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2
    du[3] = y₁ + y₂ + y₃ - 1
    return nothing
end

u₀ = [1.0, 0, 0]
M = [1.0 0 0
     0 1.0 0
     0 0 0]

tspan = (0.0, 1.0)
p = [0.04, 3e7, 1e4]

stiff_func = ODEFunction(f!; mass_matrix = M)
prob_stiff = ODEProblem(stiff_func, u₀, tspan, p)
sol_stiff = solve(prob_stiff, Rodas5(); saveat = 0.1)

nn_dudt2 = Lux.Chain(Lux.Dense(3, 64, tanh), Lux.Dense(64, 2))

pinit, st = Lux.setup(rng, nn_dudt2)

model_stiff_ndae = NeuralODEMM(nn_dudt2, (u, p, t) -> [u[1] + u[2] + u[3] - 1],
    tspan, M, Rodas5(; autodiff = false); saveat = 0.1)

function predict_stiff_ndae(p)
    return model_stiff_ndae(u₀, p, st)[1]
end

function loss_stiff_ndae(p)
    pred = predict_stiff_ndae(p)
    loss = sum(abs2, Array(sol_stiff) .- pred)
    return loss
end

# callback = function (state, l, pred) #callback function to observe training
#   display(l)
#   return false
# end

l1 = first(loss_stiff_ndae(ComponentArray(pinit)))

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_stiff_ndae(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
result_stiff = Optimization.solve(optprob, OptimizationOptimJL.BFGS(); maxiters = 100)
```

--------------------------------

### Importing Libraries and Setting Up Devices (Julia)

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/mnist_conv_neural_ode.md

Imports necessary libraries for DiffEqFlux, Lux, CUDA, Zygote, MLDatasets, OrdinaryDiffEq, Optimization, and MLUtils. It defines `cpu_device` and `gpu_device` for device management and sets up `logitcrossentropy` for loss calculation. `CUDA.allowscalar(false)` and `ENV["DATADEPS_ALWAYS_ACCEPT"] = true` are configured for GPU performance and data loading.

```Julia
using DiffEqFlux, ComponentArrays, CUDA, Zygote, MLDatasets, OrdinaryDiffEq,
      Printf, LuxCUDA, Random, MLUtils, OneHotArrays
using Optimization, OptimizationOptimisers
using MLDatasets: MNIST

const cdev = cpu_device()
const gdev = gpu_device()

logitcrossentropy = CrossEntropyLoss(; logits = Val(true))

CUDA.allowscalar(false)
ENV["DATADEPS_ALWAYS_ACCEPT"] = true
```

--------------------------------

### Using NeuralODE Struct with Lux.Chain for GPU in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/GPUs.md

This snippet illustrates the direct use of the `NeuralODE` struct with a `Lux.Chain` model that has its parameters and state already on the GPU. It shows how to call the `prob_neuralode_gpu` object with GPU-allocated initial conditions, parameters, and state, enabling GPU-accelerated predictions.

```Julia
prob_neuralode_gpu = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
prob_neuralode_gpu(u0, p, st)
```

--------------------------------

### Defining a Neural ODE Manually for GPU in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/GPUs.md

This snippet demonstrates how to manually define a neural ODE for GPU execution. It initializes a `Lux.Chain` model, moves its parameters and state to the GPU, and then defines a `dudt` function that uses this model. The `ODEProblem` is set up with a GPU-allocated initial condition, ensuring the `Tsit5` solver runs on the GPU.

```Julia
using OrdinaryDiffEq, Lux, LuxCUDA, SciMLSensitivity, ComponentArrays, Random
rng = Xoshiro(0)

const cdev = cpu_device()
const gdev = gpu_device()

model = Chain(Dense(2, 50, tanh), Dense(50, 2))
ps, st = Lux.setup(rng, model)
ps = ps |> ComponentArray |> gdev
st = st |> gdev
dudt(u, p, t) = model(u, p, st)[1]

# Simulation interval and intermediary points
tspan = (0.0f0, 10.0f0)
tsteps = 0.0f0:1.0f-1:10.0f0

u0 = Float32[2.0; 0.0] |> gdev
prob_gpu = ODEProblem(dudt, u0, tspan, ps)

# Runs on a GPU
sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)
```

--------------------------------

### Training Neural ODE Model in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/augmented_neural_ode.md

This snippet demonstrates how to set up and train a Neural ODE model. It defines an `OptimizationFunction` using `loss_node` and `AutoZygote` for automatic differentiation, then creates an `OptimizationProblem` with the model parameters and `dataloader`. The model is solved for 100 epochs with a callback, and the results are used to plot a contour.

```Julia
model, ps, st = construct_model(1, 2, 64, 0)

optfunc = OptimizationFunction(
    (x, data) -> loss_node(model, data, x, st),
    Optimization.AutoZygote())
optprob = OptimizationProblem(optfunc, ComponentArray(ps |> cdev) |> gdev, dataloader)
res = solve(optprob, opt; callback = cb, epochs = 100)

plot_contour(model, res.u, st)
```

--------------------------------

### Defining Neural Network and NeuralODE Problem - Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/neural_ode.md

This section defines the neural network architecture using `Lux.Chain` and then constructs the `NeuralODE` problem. It initializes the network parameters and state, preparing the model for integration with the differential equation solver.

```Julia
dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
p, st = Lux.setup(rng, dudt2)
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
```

--------------------------------

### Training Full Neural ODE with Original Data in Julia

Source: https://github.com/sciml/diffeqflux.jl/blob/master/docs/src/examples/collocation.md

This final snippet performs the second stage of training. It sets up a new `OptimizationProblem` using the `loss_neuralode` function, which involves solving the full Neural ODE. The network parameters, initialized from the pretraining stage, are further refined by optimizing against the original `data`. The final trained Neural ODE solution is then plotted against the original data.

```Julia
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))

numerical_neuralode = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.05); callback, maxiters = 300)

nn_sol, st = prob_neuralode(u0, numerical_neuralode.u, st)
scatter(tsteps, data')
plot!(nn_sol; lw = 5)
```