SciML Sensitivity

https://docs.sciml.ai/SciMLSensitivity/stable/
Admin
SciML Sensitivity is the automatic differentiation and adjoints system for the SciML ecosystem, used...
Tokens:71,953
Snippets:488
Trust Score:9.1
Update:4 months ago
Show doc for...
Tokens:
Raw
### Setup and Model Definition for Universal Differential Equations

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/optimal_control/feedback_control

Initializes the Lux model, parameters, and system dynamics for a universal differential equation problem. This setup is required before defining the differential equation itself.

```julia
import Lux
import Optimization as OPT
import OptimizationPolyalgorithms as OPA
import ComponentArrays as CA
import SciMLSensitivity as SMS
import Zygote
import OrdinaryDiffEq as ODE
import Plots
import Random

rng = Random.default_rng()
u0 = [1.1]
tspan = (0.0, 25.0)
tsteps = 0.0:1.0:25.0

model_univ = Lux.Chain(Lux.Dense(2, 16, tanh), Lux.Dense(16, 16, tanh), Lux.Dense(16, 1))
ps, st = Lux.setup(Random.default_rng(), model_univ)
p_model = CA.ComponentArray(ps)

# Parameters of the second equation (linear dynamics)
p_system = [0.5, -0.5]
p_all = CA.ComponentArray(; p_model, p_system)
θ = CA.ComponentArray(; u0, p_all)

function dudt_univ!(du, u, p, t)
    # Destructure the parameters
    model_weights = p.p_model
    α, β = p.p_system

    # The neural network outputs a control taken by the system
    # The system then produces an output
    model_control, system_output = u

    # Dynamics of the control and system
    dmodel_control = first(model_univ(u, model_weights, st))[1]
    dsystem_output = α * system_output + β * model_control

    # Update in place
    du[1] = dmodel_control
    du[2] = dsystem_output
end

prob_univ = ODE.ODEProblem(dudt_univ!, [0.0, u0[1]], tspan, p_all)
sol_univ = ODE.solve(prob_univ, ODE.Tsit5(), abstol = 1e-8, reltol = 1e-6)

function predict_univ(θ)
    return Array(ODE.solve(prob_univ, ODE.Tsit5(), u0 = [0.0, θ.u0[1]], p = θ.p_all,
        sensealg = SMS.InterpolatingAdjoint(autojacvec = SMS.ReverseDiffVJP(true)),
        saveat = tsteps))
end

loss_univ(θ) = sum(abs2, predict_univ(θ)[2, :] .- 1)
l = loss_univ(θ)
```

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

### Full PDE Constrained Optimization Example

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/pde/pde_constrained

This snippet shows the complete workflow for solving a PDE-constrained optimization problem, including problem setup, ODE definition, solving, and optimization.

```julia
import SciMLSensitivity as SMS
import DelimitedFiles, Plots
import OrdinaryDiffEq as ODE, Optimization as OPT, OptimizationPolyalgorithms as OPA, Zygote

# Problem setup parameters:
Lx = 10.0
x = 0.0:0.01:Lx
dx = x[2] - x[1]
Nx = size(x)

u0 = exp.(-(x .- 3.0) .^ 2) # I.C

## Problem Parameters
p = [1.0, 1.0]    # True solution parameters
const xtrs = [dx, Nx]      # Extra parameters
dt = 0.40 * dx^2    # CFL condition
t0, tMax = 0.0, 1000 * dt
tspan = (t0, tMax)
t = t0:dt:tMax;

## Definition of Auxiliary functions
function ddx(u, dx)
    """
    2nd order Central difference for 1st degree derivative
    """
    return [[zero(eltype(u))]; (u[3:end] - u[1:(end - 2)]) ./ (2.0 * dx); [zero(eltype(u))]]
end

function d2dx(u, dx)
    """
    2nd order Central difference for 2nd degree derivative
    """
    return [zero(eltype(u));
            (@view(u[3:end]) .- 2.0 .* @view(u[2:(end - 1)]) .+ @view(u[1:(end - 2)])) ./
            (dx^2)
            zero(eltype(u))]
end

## ODE description of the Physics:
function heat(u, p, t, xtrs)
    # Model parameters
    a0, a1 = p
    dx, Nx = xtrs #[1.0,3.0,0.125,100]
    return 2.0 * a0 .* u + a1 .* d2dx(u, dx)
end
heat_closure(u, p, t) = heat(u, p, t, xtrs)

# Testing Solver on linear PDE
prob = ODE.ODEProblem(heat_closure, u0, tspan, p)
sol = ODE.solve(prob, ODE.Tsit5(); dt, saveat = t);
arr_sol = Array(sol)

Plots.plot(x, sol.u[1], lw = 3, label = "t0", size = (800, 500))
Plots.plot!(x, sol.u[end], lw = 3, ls = :dash, label = "tMax")

ps = [0.1, 0.2];   # Initial guess for model parameters
function predict(θ)
    Array(ODE.solve(prob, ODE.Tsit5(); p = θ, dt, saveat = t))
end

## Defining Loss function
function loss(θ)
    pred = predict(θ)
    return sum(abs2.(predict(θ) .- arr_sol)) # Mean squared error
end

l = loss(ps)
size(sol), size(t) # Checking sizes

LOSS = []                              # Loss accumulator
PRED = []                              # prediction accumulator
PARS = []                              # parameters accumulator

cb = function (st, l) #callback function to observe training
    display(l)
    pred = predict(st.u)
    append!(PRED, [pred])
    append!(LOSS, l)
    append!(PARS, [st.u])
    false
end

cb((; u = ps), loss(ps)) # Testing callback function

# Let see prediction vs. Truth
Plots.scatter(sol[:, end], label = "Truth", size = (800, 500))
Plots.plot!(PRED[end][:, end], lw = 2, label = "Prediction")

adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss(x), adtype)

optprob = OPT.OptimizationProblem(optf, ps)
res = OPT.solve(optprob, OPA.PolyOpt(), callback = cb)
@show res.u # returns [0.999999999613485, 0.9999999991343996]

```

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

### Lotka-Volterra ODE and Parameter Estimation Setup

Source: https://docs.sciml.ai/SciMLSensitivity/stable/tutorials/parameter_estimation_ode

This snippet sets up the Lotka-Volterra ODE problem, solves it, defines a loss function, and configures an optimization problem to estimate parameters. It includes the ODE definition, initial conditions, parameter values, and the optimization setup.

```julia
import OrdinaryDiffEq as ODE
import Optimization as OPT
import OptimizationPolyalgorithms as OPA
import SciMLSensitivity as SMS
import Zygote
import Plots

function lotka_volterra!(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = dx = α * x - β * x * y
    du[2] = dy = -δ * y + γ * x * y
end

# Initial condition
u0 = [1.0, 1.0]

# Simulation interval and intermediary points
tspan = (0.0, 10.0)
tsteps = 0.0:0.1:10.0

# LV equation parameter. p = [α, β, δ, γ]
p = [1.5, 1.0, 3.0, 1.0]

# Setup the ODE problem, then solve
prob = ODE.ODEProblem(lotka_volterra!, u0, tspan, p)
sol = ODE.solve(prob, ODE.Tsit5())

# Plot the solution
Plots.plot(sol)
Plots.savefig("LV_ode.png")

function loss(p)
    sol = ODE.solve(prob, ODE.Tsit5(); p, saveat = tsteps)
    loss = sum(abs2, sol .- 1)
    return loss
end

callback = function (state, l)
    display(l)
    pred = ODE.solve(prob, ODE.Tsit5(), p = state.u, saveat = tsteps)
    plt = Plots.plot(pred, ylim = (0, 6))
    display(plt)
    # Tell Optimization.solve to not halt the optimization. If return true, then
    # optimization stops.
    return false
end

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

result_ode = OPT.solve(optprob, OPA.PolyOpt(); callback, maxiters = 1000)

```

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

### ODE Forward Sensitivity Problem Setup

Source: https://docs.sciml.ai/SciMLSensitivity/stable/tutorials/adjoint_continuous_functional

Sets up an ODE problem for forward sensitivity analysis. Ensure `dense=true` in the `solve` call for continuous solutions.

```julia
import OrdinaryDiffEq as ODE
import SciMLSensitivity as SMS

function f(du, u, p, t)
    du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = dy = -p[3] * u[2] + u[1] * u[2]
end

p = [1.5, 1.0, 3.0]
prob = SMS.ODEForwardSensitivityProblem(f, [1.0; 1.0], (0.0, 10.0), p)
sol = ODE.solve(prob, ODE.DP8())

```

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

### Define and Setup Neural Network

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/sde/SDE_control

Defines a simple neural network using Lux for control and sets up its initial parameters and state.

```julia
# Define Neural Network

# state-aware
nn = Lux.Chain(Lux.Dense(4, 32, Lux.relu),
    Lux.Dense(32, 1, tanh))

p_nn, st = Lux.setup(rng, nn)
p_nn = CA.ComponentArray(p_nn)
```

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

### Full Example: Lotka-Volterra Model Optimization

Source: https://docs.sciml.ai/SciMLSensitivity/stable/tutorials/training_tips/divergence

A complete example demonstrating the Lotka-Volterra model, solving it, plotting results, and setting up an optimization problem with a loss function that handles divergent trajectories.

```julia
import OrdinaryDiffEq as ODE, SciMLSensitivity as SMS, SciMLBase, Optimization as OPT,
       OptimizationOptimisers as OPO, Plots

function lotka_volterra!(du, u, p, t)
    rab, wol = u
    α, β, γ, δ = p
    du[1] = drab = α * rab - β * rab * wol
    du[2] = dwol = γ * rab * wol - δ * wol
    nothing
end

u0 = [1.0, 1.0]
tspan = (0.0, 10.0)
p = [1.5, 1.0, 3.0, 1.0]
prob = ODE.ODEProblem(lotka_volterra!, u0, tspan, p)
sol = ODE.solve(prob, ODE.Tsit5(); saveat = 0.1)
Plots.plot(sol)

dataset = Array(sol)
Plots.scatter!(sol.t, dataset')

tmp_prob = ODE.remake(prob, p = [1.2, 0.8, 2.5, 0.8])
tmp_sol = ODE.solve(tmp_prob, ODE.Tsit5())
Plots.plot(tmp_sol)
Plots.scatter!(sol.t, dataset')

function loss(p)
    tmp_prob = ODE.remake(prob; p)
    tmp_sol = ODE.solve(tmp_prob, ODE.Tsit5(), saveat = 0.1)
    if tmp_sol.retcode == SciMLBase.ReturnCode.Success
        return sum(abs2, Array(tmp_sol) - dataset)
    else
        return Inf
    end
end

pinit = [1.2, 0.8, 2.5, 0.8]
adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss(x), adtype)

optprob = OPT.OptimizationProblem(optf, pinit)
res = OPT.solve(optprob, OPO.Adam(), maxiters = 1000)

# res = OPT.solve(optprob,NLopt.LD_LBFGS(), maxiters = 1000) ### errors!
```

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

### Multithreaded Batching Example

Source: https://docs.sciml.ai/SciMLSensitivity/stable/tutorials/data_parallel

This example demonstrates how to use multithreading for parallelizing ODE solves. It sets up an ODE problem, defines a function to remake the problem for different trajectories, and then solves it using both serial and multithreaded ensemble methods for comparison.

```julia
import OrdinaryDiffEq as ODE
import Optimization as OPT
import OptimizationOptimisers as OPO
import SciMLSensitivity as SMS
pa = [1.0]
u0 = [3.0]
θ = [u0; pa]

function model1(θ, ensemble)
    prob = ODE.ODEProblem((u, p, t) -> 1.01u .* p, [θ[1]], (0.0, 1.0), [θ[2]])

    function prob_func(prob, i, repeat)
        ODE.remake(prob, u0 = 0.5 .+ i / 100 .* prob.u0)
    end

    ensemble_prob = ODE.EnsembleProblem(prob; prob_func)
    sim = ODE.solve(ensemble_prob, ODE.Tsit5(), ensemble, saveat = 0.1, trajectories = 100)
end

# loss function
loss_serial(θ) = sum(abs2, 1.0 .- Array(model1(θ, ODE.EnsembleSerial())))
loss_threaded(θ) = sum(abs2, 1.0 .- Array(model1(θ, ODE.EnsembleThreads())))

callback = function (θ, l) # callback function to observe training
    @show l
    false
end

opt = OPO.Adam(0.1)
l1 = loss_serial(θ)

adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss_serial(x), adtype)
optprob = OPT.OptimizationProblem(optf, θ)

res_serial = OPT.solve(optprob, opt; callback, maxiters = 100)

optf2 = OPT.OptimizationFunction((x, p) -> loss_threaded(x), adtype)
optprob2 = OPT.OptimizationProblem(optf2, θ)

res_threads = OPT.solve(optprob2, opt; callback, maxiters = 100)
```

```text
retcode: Default
u: 2-element Vector{Float64}:
  1.094120061965145
 -0.41178007289731716
```

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

### Problem Setup Parameters for Heat Equation

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/pde/pde_constrained

Defines spatial and temporal parameters, initial conditions, and physical constants for the 1D Heat Equation.

```julia
# Problem setup parameters:
Lx = 10.0
x = 0.0:0.01:Lx
dx = x[2] - x[1]
Nx = size(x)

u0 = exp.(-(x .- 3.0) .^ 2) # I.C

## Problem Parameters
p = [1.0, 1.0]    # True solution parameters
const xtrs = [dx, Nx]      # Extra parameters
dt = 0.40 * dx^2    # CFL condition
t0, tMax = 0.0, 1000 * dt
tspan = (t0, tMax)
t = t0:dt:tMax;
```

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

### Training Loop Setup and Execution

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/sde/SDE_control

Sets up and runs the optimization loop using the Adam optimizer. It defines the optimization function, problem, and solves it with a callback for visualization.

```julia
# training loop
@info "Start Training.."

# optimize the parameters for a few epochs with Adam on time span
# Setup and run the optimization
adtype = OPT.AutoForwardDiff()
optf = OPT.OptimizationFunction((x, p) -> loss(x), adtype)

optprob = OPT.OptimizationProblem(optf, p_nn)
res = OPT.solve(optprob, OPO.Adam(myparameters.lr),
    callback = visualization_callback,
    maxiters = 100)

# plot optimized control
visualization_callback(res, loss(res.u); doplot = true)
```

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

### ODE Problem Setup and Solution

Source: https://docs.sciml.ai/SciMLSensitivity/stable/manual/direct_adjoint_sensitivities

Sets up and solves an Ordinary Differential Equation problem using specified parameters and initial conditions. This provides the solution object `sol` required for subsequent sensitivity analysis.

```julia
p = [1.5,1.0,3.0]
prob = ODEProblem(f,[1.0;1.0],(0.0,10.0),p)
sol = solve(prob,Vern9(),abstol=1e-10,reltol=1e-10)
```

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

### Install SciMLSensitivity.jl

Source: https://docs.sciml.ai/SciMLSensitivity/stable

Use the Julia package manager to add SciMLSensitivity.jl to your project. This command installs the package and its dependencies.

```julia
import Pkg
Pkg.add("SciMLSensitivity")
```

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

### Pendulum Dynamics Simulation Setup

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/ode/prediction_error_method

Defines the time span, initial conditions, and the differential equation for a simple pendulum system. This setup is used for generating simulation data and for the prediction error method.

```julia
import OrdinaryDiffEq as ODE
import Optimization as OPT
import OptimizationPolyalgorithms as OPA
import Plots
import Statistics
import DataInterpolations as DI
import ForwardDiff as FD
import SciMLBase as SMB

tspan = (0.1, 20.0)
tsteps = range(tspan[1], tspan[2], length = 1000)

u0 = [0.0, 3.0] # Initial angle and angular velocity

function simulator(du, u, p, t) # Pendulum dynamics
    g = 9.82 # Gravitational constant
    L = p isa Number ? p : p[1] # Length of the pendulum
    gL = g / L
    θ, dθ = u
    du[1] = dθ
    du[2] = -gL * sin(θ)
end
```

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

### ODEProblem Setup for UDE

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/pde/brusselator

This snippet shows the resulting ODEProblem object after defining the UDE function and initial conditions. It confirms the problem's in-place nature, time span, and the structure of the initial state (u0).

```julia
ODEProblem with uType Array{Float32, 3} and tType Float32. In-place: true
Non-trivial mass matrix: false
timespan: (0.0f0, 11.5f0)
u0: 16×16×2 Array{Float32, 3}:
[:, :, 1] =
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0

[:, :, 2] =
 0.0       0.0       0.0       0.0       …  0.0       0.0       0.0
 0.419066  0.419066  0.419066  0.419066     0.419066  0.419066  0.419066
 1.0606    1.0606    1.0606    1.0606       1.0606    1.0606    1.0606
 1.728     1.728     1.728     1.728        1.728     1.728     1.728
```

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

### Using a Closure for Exogenous Input

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/ode/exogenous_input

This example demonstrates using a closure to pass an exogenous input signal as an extra argument to a differential equation function, making it interface-compliant.

```Julia
function f(du, u, p, t, I)
    du[1] = I(t)
    du[2] = u[1]
end

_f(du, u, p, t) = f(du, u, p, t, x -> x^2)
```

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

### Example Output of Estimated Parameters

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/ode/prediction_error_method

Provides an example of the numerical output for the estimated parameters (u) from the prediction-error method. This is a 2-element vector of Float64 values.

```julia
2-element Vector{Float64}:
  0.9999999985077761
 -2.3087300880395805e-7
```

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

### Neural ODE with Discrete Exogenous Input

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/ode/exogenous_input

This example sets up and solves a Neural ODE system that incorporates a discrete exogenous input signal. It defines the system, trains the neural network, and visualizes the results.

```Julia
import SciMLSensitivity as SMS
import OrdinaryDiffEq as ODE
import Lux
import ComponentArrays as CA
import Optimization as OPT
import OptimizationPolyalgorithms as OPA
import OptimizationOptimisers as OPO
import Plots
import Random

rng = Random.default_rng()
tspan = (0.1, 10.0)
tsteps = range(tspan[1], tspan[2], length = 100)
t_vec = collect(tsteps)
ex = vec(ones(Float64, length(tsteps), 1))
f(x) = (atan(8.0 * x - 4.0) + atan(4.0)) / (2.0 * atan(4.0))

function hammerstein_system(u)
    y = zeros(size(u))
    for k in 2:length(u)
        y[k] = 0.2 * f(u[k - 1]) + 0.8 * y[k - 1]
    end
    return y
end

y = hammerstein_system(ex)
Plots.plot(collect(tsteps), y, ticks = :native)

nn_model = Lux.Chain(Lux.Dense(2, 8, tanh), Lux.Dense(8, 1))
p_model, st = Lux.setup(rng, nn_model)

u0 = Float64.([0.0])

function dudt(u, p, t)
    global st
    #input_val = u_vals[Int(round(t*10)+1)]
    out, st = nn_model(vcat(u[1], ex[Int(round(10 * 0.1))]), p, st)
    return out
end

prob = ODE.ODEProblem(dudt, u0, tspan, nothing)

function predict_neuralode(p)
    _prob = ODE.remake(prob; p)
    Array(ODE.solve(_prob, ODE.Tsit5(), saveat = tsteps, abstol = 1e-8, reltol = 1e-6))
end

function loss(p)
    sol = predict_neuralode(p)
    N = length(sol)
    return sum(abs2.(y[1:N] .- sol')) / N
end

adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = OPT.OptimizationProblem(optf, CA.ComponentArray{Float64}(p_model))

res0 = OPT.solve(optprob, OPA.PolyOpt(), maxiters = 100)

sol = predict_neuralode(res0.u)
Plots.plot(tsteps, sol')
N = length(sol)
Plots.scatter!(tsteps, y[1:N])
```

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

### Hybrid Differential Equation Training Example

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/hybrid_jump/hybrid_diffeq

This snippet sets up and trains a neural network to model the unknown dynamics in a hybrid differential equation system. It defines the ODE problem, the neural network, the loss function, and uses an optimizer with a callback to monitor training progress. Requires Lux, OrdinaryDiffEq, SciMLSensitivity, Optimization, and DiffEqCallbacks.

```julia
import ComponentArrays as CA
import Random
import SciMLSensitivity as SMS
import Lux
import OrdinaryDiffEq as ODE
import Plots
import Optimization as OPT
import OptimizationOptimisers as OPO
import DiffEqCallbacks as DEC

u0 = Float32[2.0; 0.0]
datasize = 100
tspan = (0.0f0, 10.5f0)
dosetimes = [1.0, 2.0, 4.0, 8.0]

function affect!(integrator)
    integrator.u = integrator.u .+ 1
end
cb_ = DEC.PresetTimeCallback(dosetimes, affect!, save_positions = (false, false))
function trueODEfunc(du, u, p, t)
    du .= -u
end
t = range(tspan[1], tspan[2], length = datasize)

prob = ODE.ODEProblem(trueODEfunc, u0, tspan)
ode_data = Array(ODE.solve(prob, ODE.Tsit5(), callback = cb_, saveat = t))

dudt2 = Lux.Chain(Lux.Dense(2, 50, tanh), Lux.Dense(50, 2))
ps, st = Lux.setup(Random.default_rng(), dudt2)

function dudt(du, u, p, t)
    du[1:2] .= -u[1:2]
    du[3:end] .= first(dudt2(u[1:2], p, st))
end
z0 = Float32[u0; u0]
prob = ODE.ODEProblem(dudt, z0, tspan)

affect!(integrator) = integrator.u[1:2] .= integrator.u[3:end]
cb = DEC.PresetTimeCallback(dosetimes, affect!, save_positions = (false, false))

function predict_n_ode(p)
    _prob = ODE.remake(prob; p)
    Array(ODE.solve(_prob, ODE.Tsit5(); u0 = z0, p, callback = cb, saveat = t,
        sensealg = SMS.ReverseDiffAdjoint()))[1:2, :]
    #Array(solve(prob,Tsit5();u0=z0,p,saveat=t))[1:2,:]
end

function loss_n_ode(p, _)
    pred = predict_n_ode(p)
    loss = sum(abs2, ode_data .- pred)
    loss
end
loss_n_ode(ps, nothing)

cba = function (state, l; doplot = false) #callback function to observe training
    pred = predict_n_ode(state.u)
    display(sum(abs2, ode_data .- pred))
    # plot current prediction against data
    pl = Plots.scatter(t, ode_data[1, :], label = "data")
    Plots.scatter!(pl, t, pred[1, :], label = "prediction")
    display(Plots.plot(pl))
    return false
end

res = OPT.solve(
    OPT.OptimizationProblem(OPT.OptimizationFunction(loss_n_ode, OPT.AutoZygote()),
        CA.ComponentArray(ps)),
    OPO.Adam(0.05); callback = cba, maxiters = 1000)

```

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

### Defining and Solving an ODE Forward Sensitivity Problem

Source: https://docs.sciml.ai/SciMLSensitivity/stable/manual/direct_forward_sensitivity

Example demonstrating the creation of an `ODEForwardSensitivityProblem` for the Lotka-Volterra equations and solving it. The solution object contains both the ODE states and their sensitivities.

```julia
function f(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end

p = [1.5,1.0,3.0]
prob = ODEForwardSensitivityProblem(f,[1.0;1.0],(0.0,10.0),p)
```

```julia
sol = solve(prob,DP8())
```

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

### ODE Solve with Forward Sensitivity using Zygote Gradient

Source: https://docs.sciml.ai/SciMLSensitivity/stable/manual/differential_equation_sensitivities

This example shows how to compute the gradient of a loss function for an ODE solve using Zygote.jl, but explicitly specifies `ForwardSensitivity()` for the `sensealg`. This forces the ODE solve's reversal to use forward sensitivity analysis, mixing forward and reverse differentiation.

```julia
function loss(u0, p)
    sum(ODE.solve(prob, ODE.Tsit5(); u0, p, saveat = 0.1,
        sensealg = SMS.ForwardSensitivity()))
end
du0, dp = Zygote.gradient(loss, u0, p)
```

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

### Training a Neural Network for a Second Order ODE

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/ode/second_order_neural

This snippet demonstrates the setup and training process for a neural network designed to solve second-order ODEs. It includes defining the problem, the neural network model, the loss function, and the optimization procedure.

```julia
import SciMLSensitivity as SMS
import OrdinaryDiffEq as ODE
import Lux
import Optimization as OPT
import OptimizationOptimisers as OPO
import RecursiveArrayTools
import Random
import ComponentArrays as CA

u0 = Float32[0.0; 2.0]
du0 = Float32[0.0; 0.0]
tspan = (0.0f0, 1.0f0)
t = range(tspan[1], tspan[2], length = 20)

model = Lux.Chain(Lux.Dense(2, 50, tanh), Lux.Dense(50, 2))
ps, st = Lux.setup(Random.default_rng(), model)
ps = CA.ComponentArray(ps)
model = Lux.StatefulLuxLayer{true}(model, ps, st)

ff(du, u, p, t) = model(u, p)
prob = ODE.SecondOrderODEProblem{false}(ff, du0, u0, tspan, ps)

function predict(p)
    Array(ODE.solve(prob, ODE.Tsit5(); p, saveat = t))
end

correct_pos = Float32.(transpose(hcat(collect(0:0.05:1)[2:end], collect(2:-0.05:1)[2:end])))

function loss_n_ode(p)
    pred = predict(p)
    sum(abs2, correct_pos .- pred[1:2, :])
end

l1 = loss_n_ode(ps)

callback = function (state, l)
    println(l)
    l < 0.01
end

adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss_n_ode(x), adtype)
optprob = OPT.OptimizationProblem(optf, ps)

res = OPT.solve(optprob, OPO.Adam(0.01); callback, maxiters = 1000)
```

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

### ODE Definition for Sensitivity Example

Source: https://docs.sciml.ai/SciMLSensitivity/stable/manual/direct_adjoint_sensitivities

Defines the differential equations for a system, used in the sensitivity analysis example. This function is a prerequisite for solving the ODE problem.

```julia
function f(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end
```

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

### Brusselator PDE Solution Data Example

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/pde/brusselator

This is an example of the output data structure from solving the Brusselator PDE. It shows a 4D array representing the state variables over space and time.

```text
16×16×2×24 Array{Float32, 4}:
[:, :, 1, 1] =
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251     1.408  0.864189  0.341461  0.0
 0.0  0.341461  0.864189  1.408  1.90251  …  1.408  0.864189  0.341461  0.0

[:, :, 2, 1] =
 0.0       0.0       0.0       0.0       …  0.0       0.0       0.0
 0.419066  0.419066  0.419066  0.419066     0.419066  0.419066  0.419066
 1.0606    1.0606    1.0606    1.0606       1.0606    1.0606    1.0606
 1.728     1.728     1.728     1.728        1.728     1.728     1.728
 2.3349    2.3349    2.3349    2.3349       2.3349    2.3349    2.3349
 2.82843   2.82843   2.82843   2.82843   …  2.82843   2.82843   2.82843
 3.17454   3.17454   3.17454   3.17454      3.17454   3.17454   3.17454
 3.35252   3.35252   3.35252   3.35252      3.35252   3.35252   3.35252
 3.35252   3.35252   3.35252   3.35252      3.35252   3.35252   3.35252
 3.17454   3.17454   3.17454   3.17454      3.17454   3.17454   3.17454
 2.82843   2.82843   2.82843   2.82843   …  2.82843   2.82843   2.82843
 2.3349    2.3349    2.3349    2.3349       2.3349    2.3349    2.3349
 1.728     1.728     1.728     1.728        1.728     1.728     1.728
 1.0606    1.0606    1.0606    1.0606       1.0606    1.0606    1.0606
 0.419066  0.419066  0.419066  0.419066     0.419066  0.419066  0.419066
 0.0       0.0       0.0       0.0       …  0.0       0.0       0.0

[:, :, 1, 2] =
 0.899116  0.899099  0.899066  0.899022  …  0.899065  0.899098  0.899116
 0.899118  0.899101  0.89907   0.899028     0.899069  0.8991    0.899118
 0.899121  0.899105  0.899076  0.899039     0.899076  0.899104  0.899121
 0.899124  0.899109  0.899083  0.899052     0.899083  0.89911   0.899125
 0.899127  0.899114  0.899092  0.899065     0.899092  0.899115  0.899128
 0.899132  0.899119  0.899099  0.899077  …  0.8991    0.89912   0.899133
 0.899136  0.899124  0.899105  0.899087     0.899106  0.899125  0.899136
 0.899137  0.899125  0.899108  0.899092     0.89911   0.899127  0.899138
 0.899137  0.899126  0.899109  0.899091     0.89911   0.899127  0.899137
 0.899136  0.899124  0.899104  0.899086     0.899107  0.899125  0.899136
 0.899132  0.899119  0.899099  0.899077  …  0.8991    0.89912   0.899133
 0.899127  0.899114  0.899092  0.899065     0.899092  0.899115  0.899128
 0.899124  0.899109  0.899083  0.899052     0.899083  0.89911   0.899125
 0.899121  0.899105  0.899076  0.899039     0.899076  0.899104  0.899121
 0.899118  0.899101  0.89907   0.899028     0.899069  0.8991    0.899118
 0.899116  0.899099  0.899066  0.899022  …  0.899065  0.899098  0.899116
```

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

### Set up Optimization Problem

Source: https://docs.sciml.ai/SciMLSensitivity/stable/examples/pde/pde_constrained

Configures the optimization problem using Zygote for automatic differentiation and PolyOpt for solving.

```julia
adtype = OPT.AutoZygote()
optf = OPT.OptimizationFunction((x, p) -> loss(x), adtype)

optprob = OPT.OptimizationProblem(optf, ps)
res = OPT.solve(optprob, OPA.PolyOpt(), callback = cb)
@show res.u # returns [0.999999999613485, 0.9999999991343996]
```
Tokens:
Raw