Merge pull request #1128 from SciML/docsoptv4

Update docs remove extra returns from loss and extra args from callback

Author: ChrisRackauckas

Project.toml

@@ -81,7 +81,7 @@ PreallocationTools = "0.4.4"
 QuadGK = "2.9.1"
 Random = "1.10"
 RandomNumbers = "1.5.3"
-RecursiveArrayTools = "3.18.1"
+RecursiveArrayTools = "3.27.2"
 Reexport = "1.0"
 ReverseDiff = "1.15.1"
 SafeTestsets = "0.1.0"

docs/Project.toml

@@ -14,6 +14,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
@@ -23,6 +24,7 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -45,6 +47,7 @@ Enzyme = "0.12, 0.13"
 Flux = "0.14"
 ForwardDiff = "0.10"
 IterTools = "1"
+MLUtils = "0.4"
 Lux = "1"
 LuxCUDA = "0.3"
 Optimization = "3.9, 4"
@@ -56,6 +59,7 @@ Plots = "1.36"
 QuadGK = "2.6"
 RecursiveArrayTools = "2.32, 3"
 ReverseDiff = "1.14"
+SciMLBase = "2.58"
 SciMLSensitivity = "7.11"
 SimpleChains = "0.4"
 StaticArrays = "1"

docs/src/examples/dde/delay_diffeq.md

@@ -38,7 +38,7 @@ end
 loss_dde(p) = sum(abs2, x - 1 for x in predict_dde(p))
 
 using Plots
-callback = function (state, l...; doplot = false)
+callback = function (state, l; doplot = false)
     display(loss_dde(state.u))
     doplot &&
         display(plot(
@@ -60,7 +60,7 @@ We define a callback to display the solution at the current parameters for each
 
 ```@example dde
 using Plots
-callback = function (state, l...; doplot = false)
+callback = function (state, l; doplot = false)
     display(loss_dde(state.u))
     doplot &&
         display(plot(

docs/src/examples/neural_ode/minibatch.md

@@ -2,7 +2,7 @@
 
 ```@example
 using SciMLSensitivity
-using DifferentialEquations, Flux, Random, Plots
+using DifferentialEquations, Flux, Random, Plots, MLUtils
 using IterTools: ncycle
 
 rng = Random.default_rng()
@@ -46,7 +46,7 @@ ode_data = Array(solve(true_prob, Tsit5(), saveat = t))
 prob = ODEProblem{false}(dudt_, u0, tspan, θ)
 
 k = 10
-train_loader = Flux.Data.DataLoader((ode_data, t), batchsize = k)
+train_loader = DataLoader((ode_data, t), batchsize = k)
 
 for (x, y) in train_loader
     @show x
@@ -96,7 +96,7 @@ When training a neural network, we need to find the gradient with respect to our
 For this example, we will use a very simple ordinary differential equation, newtons law of cooling. We can represent this in Julia like so.
 
 ```@example minibatch
-using SciMLSensitivity
+using SciMLSensitivity, MLUtils
 using DifferentialEquations, Flux, Random, Plots
 using IterTools: ncycle
 
@@ -152,7 +152,7 @@ ode_data = Array(solve(true_prob, Tsit5(), saveat = t))
 prob = ODEProblem{false}(dudt_, u0, tspan, θ)
 
 k = 10
-train_loader = Flux.Data.DataLoader((ode_data, t), batchsize = k)
+train_loader = DataLoader((ode_data, t), batchsize = k)
 
 for (x, y) in train_loader
     @show x

docs/src/examples/neural_ode/neural_ode_flux.md

@@ -114,22 +114,23 @@ end
 function loss_n_ode(θ)
     pred = predict_n_ode(θ)
     loss = sum(abs2, ode_data .- pred)
-    loss, pred
+    return loss
 end
 
 loss_n_ode(θ)
 
-callback = function (θ, l, pred; doplot = false) #callback function to observe training
+callback = function (state, l; doplot = false) #callback function to observe training
     display(l)
     # plot current prediction against data
+    pred = predict_n_ode(state.u)
     pl = scatter(t, ode_data[1, :], label = "data")
     scatter!(pl, t, pred[1, :], label = "prediction")
     display(plot(pl))
     return false
 end
 
 # Display the ODE with the initial parameter values.
-callback(θ, loss_n_ode(θ)...)
+callback((; u = θ), loss_n_ode(θ)...)
 
 # use Optimization.jl to solve the problem
 adtype = Optimization.AutoZygote()
@@ -143,7 +144,7 @@ result_neuralode = Optimization.solve(optprob,
     maxiters = 300)
 ```
 
-Notice that the advantage of this format is that we can use Optim's optimizers, like
+Notice that the advantage of this format is that we can use other optimizers, like
 `LBFGS` with a full `Chain` object, for all of Flux's neural networks, like
 convolutional neural networks.
 

docs/src/examples/neural_ode/simplechains.md

@@ -54,19 +54,20 @@ end
 function loss_neuralode(p)
     pred = predict_neuralode(p)
     loss = sum(abs2, data .- pred)
-    return loss, pred
+    return loss
 end
 ```
 
 ## Training
 
 The next step is to minimize the loss, so that the NeuralODE gets trained. But in order to be able to do that, we have to be able to backpropagate through the NeuralODE model. Here the backpropagation through the neural network is the easy part, and we get that out of the box with any deep learning package(although not as fast as SimpleChains for the small nn case here). But we have to find a way to first propagate the sensitivities of the loss back, first through the ODE solver and then to the neural network.
 
-The adjoint of a neural ODE can be calculated through the various AD algorithms available in SciMLSensitivity.jl. But working with [StaticArrays](https://docs.sciml.ai/StaticArrays/stable/) in SimpleChains.jl requires a special adjoint method as StaticArrays do not allow any mutation. All the adjoint methods make heavy use of in-place mutation to be performant with the heap allocated normal arrays. For our statically sized, stack allocated StaticArrays, in order to be able to compute the ODE adjoint we need to do everything out of place. Hence, we have specifically used `QuadratureAdjoint(autojacvec=ZygoteVJP())` adjoint algorithm in the solve call inside `predict_neuralode(p)` which computes everything out-of-place when u0 is a StaticArray. Hence, we can move forward with the training of the NeuralODE
+The adjoint of a neural ODE can be calculated through the various AD algorithms available in SciMLSensitivity.jl. But working with [StaticArrays](https://juliaarrays.github.io/StaticArrays.jl/stable/) in SimpleChains.jl requires a special adjoint method as StaticArrays do not allow any mutation. All the adjoint methods make heavy use of in-place mutation to be performant with the heap allocated normal arrays. For our statically sized, stack allocated StaticArrays, in order to be able to compute the ODE adjoint we need to do everything out of place. Hence, we have specifically used `QuadratureAdjoint(autojacvec=ZygoteVJP())` adjoint algorithm in the solve call inside `predict_neuralode(p)` which computes everything out-of-place when u0 is a StaticArray. Hence, we can move forward with the training of the NeuralODE
 
 ```@example sc_neuralode
-callback = function (state, l, pred; doplot = true)
+callback = function (state, l; doplot = true)
     display(l)
+    pred = predict_neuralode(state.u)
     plt = scatter(tsteps, data[1, :], label = "data")
     scatter!(plt, tsteps, pred[1, :], label = "prediction")
     if doplot

docs/src/examples/ode/second_order_adjoints.md

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

docs/src/examples/ode/second_order_neural.md

@@ -33,7 +33,7 @@ t = range(tspan[1], tspan[2], length = 20)
 model = Chain(Dense(2, 50, tanh), Dense(50, 2))
 ps, st = Lux.setup(Random.default_rng(), model)
 ps = ComponentArray(ps)
-model = StatefulLuxLayer{true}(model, ps, st)
+model = Lux.StatefulLuxLayer{true}(model, ps, st)
 
 ff(du, u, p, t) = model(u, p)
 prob = SecondOrderODEProblem{false}(ff, du0, u0, tspan, ps)
@@ -46,12 +46,12 @@ correct_pos = Float32.(transpose(hcat(collect(0:0.05:1)[2:end], collect(2:-0.05:
 
 function loss_n_ode(p)
     pred = predict(p)
-    sum(abs2, correct_pos .- pred[1:2, :]), pred
+    sum(abs2, correct_pos .- pred[1:2, :])
 end
 
 l1 = loss_n_ode(ps)
 
-callback = function (state, l, pred)
+callback = function (state, l)
     println(l)
     l < 0.01
 end

docs/src/examples/optimal_control/feedback_control.md

@@ -61,7 +61,7 @@ l = loss_univ(θ)
 ```@example udeneuralcontrol
 list_plots = []
 iter = 0
-cb = function (state, l)
+cb = function (state, l; makeplot = false)
     global list_plots, iter
 
     if iter == 0
@@ -71,9 +71,11 @@ cb = function (state, l)
 
     println(l)
 
-    plt = plot(predict_univ(state.u)', ylim = (0, 6))
-    push!(list_plots, plt)
-    display(plt)
+    if makeplot
+        plt = plot(predict_univ(state.u)', ylim = (0, 6))
+        push!(list_plots, plt)
+        display(plt)
+    end
     return false
 end
 ```
@@ -84,3 +86,7 @@ optf = Optimization.OptimizationFunction((x, p) -> loss_univ(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, θ)
 result_univ = Optimization.solve(optprob, PolyOpt(), callback = cb)
 ```
+
+```@example udeneuralcontrol
+cb(result_univ, result_univ.minimum; makeplot=true)
+```

docs/src/examples/optimal_control/optimal_control.md

@@ -37,7 +37,8 @@ of a local minimum. This looks like:
 
 ```@example neuraloptimalcontrol
 using Lux, ComponentArrays, OrdinaryDiffEq, Optimization, OptimizationOptimJL,
-      OptimizationOptimisers, SciMLSensitivity, Zygote, Plots, Statistics, Random
+      OptimizationOptimisers, SciMLSensitivity, Zygote, Plots, Statistics, Random,
+      ForwardDiff
 
 rng = Random.default_rng()
 tspan = (0.0f0, 8.0f0)
@@ -89,7 +90,7 @@ end
 # Setup and run the optimization
 
 loss1 = loss_adjoint(θ)
-adtype = Optimization.AutoZygote()
+adtype = Optimization.AutoForwardDiff()
 optf = Optimization.OptimizationFunction((x, p) -> loss_adjoint(x), adtype)
 
 optprob = Optimization.OptimizationProblem(optf, θ)