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

Author: Vaibhavdixit02

docs/src/getting_started/fit_simulation.md

@@ -99,12 +99,14 @@ function loss(newp)
     newprob = remake(prob, p = newp)
     sol = solve(newprob, saveat = 1)
     loss = sum(abs2, sol .- xy_data)
-    return loss, sol
+    return loss
 end
 
 # Define a callback function to monitor optimization progress
-function callback(p, l, sol)
+function callback(state, l)
     display(l)
+    newprob = remake(prob, p = state.u)
+    sol = solve(newprob, saveat = 1)
     plt = plot(sol, ylim = (0, 6), label = ["Current x Prediction" "Current y Prediction"])
     scatter!(plt, t_data, xy_data', label = ["x Data" "y Data"])
     display(plt)
@@ -278,37 +280,28 @@ function loss(newp)
     newprob = remake(prob, p = newp)
     sol = solve(newprob, saveat = 1)
     l = sum(abs2, sol .- xy_data)
-    return l, sol
+    return l
 end
 ```
 
-Notice that our loss function returns the loss value as the first return,
-but returns extra information (the ODE solution with the new parameters)
-as an extra return argument.
-We will explain why this extra return information is helpful in the next section.
-
 ### Step 5: Solve the Optimization Problem
 
 This step will look very similar to [the first optimization tutorial](@ref first_opt),
-except now we have a new loss function `loss` which returns both the loss value
-and the associated ODE solution.
-(In the previous tutorial, `L` only returned the loss value.)
 The `Optimization.solve` function can accept an optional callback function
 to monitor the optimization process using extra arguments returned from `loss`.
 
 The callback syntax is always:
 
 ```
 callback(
-    optimization variables,
+    state,
     the current loss value,
-    other arguments returned from the loss function, ...
 )
 ```
 
-In this case, we will provide the callback the arguments `(p, l, sol)`,
-since it always takes the current state of the optimization first (`p`)
-then the returns from the loss function (`l, sol`).
+In this case, we will provide the callback the arguments `(state, l)`,
+since it always takes the current state of the optimization first (`state`)
+then the current loss value (`l`).
 The return value of the callback function should default to `false`.
 `Optimization.solve` will halt if/when the callback function returns `true` instead.
 Typically the `return` statement would monitor the loss value
@@ -318,8 +311,10 @@ More details about callbacks in Optimization.jl can be found
 [here](https://docs.sciml.ai/Optimization/stable/API/solve/).
 
 ```@example odefit
-function callback(p, l, sol)
+function callback(p, l)
     display(l)
+    newprob = remake(prob, p = p)
+    sol = solve(newprob, saveat = 1)
     plt = plot(sol, ylim = (0, 6), label = ["Current x Prediction" "Current y Prediction"])
     scatter!(plt, t_data, xy_data', label = ["x Data" "y Data"])
     display(plt)

docs/src/showcase/blackhole.md

@@ -490,10 +490,8 @@ function loss(NN_params)
         prob_nn, RK4(), u0 = u0, p = NN_params, saveat = tsteps, dt = dt, adaptive = false))
     pred_waveform = compute_waveform(dt_data, pred, mass_ratio, model_params)[1]
 
-    loss = (sum(abs2,
-        view(waveform, obs_to_use_for_training) .-
-        view(pred_waveform, obs_to_use_for_training)))
-    return loss, pred_waveform
+    loss = ( sum(abs2, view(waveform,obs_to_use_for_training) .- view(pred_waveform,obs_to_use_for_training) ) )
+    return loss
 end
 ```
 
@@ -508,10 +506,11 @@ We'll use the following callback to save the history of the loss values.
 ```@example ude
 losses = []
 
-callback(θ, l, pred_waveform; doplot = true) = begin
+callback(state, l; doplot = true) = begin
     push!(losses, l)
     #=  Disable plotting as it trains since in docs
     display(l)
+    waveform = compute_waveform(dt_data, soln, mass_ratio, model_params)[1]
     # plot current prediction against data
     plt = plot(tsteps, waveform,
         markershape=:circle, markeralpha = 0.25,

docs/src/showcase/missing_physics.md

@@ -222,7 +222,7 @@ current loss:
 ```@example ude
 losses = Float64[]
 
-callback = function (p, l)
+callback = function (state, l)
     push!(losses, l)
     if length(losses) % 50 == 0
         println("Current loss after $(length(losses)) iterations: $(losses[end])")

docs/src/showcase/pinngpu.md

@@ -148,7 +148,7 @@ prob = discretize(pde_system, discretization)
 ## Step 6: Solve the Optimization Problem
 
 ```@example pinn
-callback = function (p, l)
+callback = function (state, l)
     println("Current loss is: $l")
     return false
 end