Some progress on FFJORD

Author: avik-pal

README.md

@@ -72,3 +72,4 @@ explore various ways to integrate the two methodologies:
   - `Flux` is no longer re-exported from `DiffEqFlux`. Instead we reexport `Lux`.
   - `NeuralDAE` now allows an optional `du0` as input.
   - `TensorLayer` is now a Lux Neural Network.
+  - APIs for quite a few layer constructions have changed. Please refer to the updated documentation for more details.

docs/make.jl

@@ -1,22 +1,22 @@
 using Documenter, DiffEqFlux
 
-cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml", force = true)
-cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)
+cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml"; force = true)
+cp("./docs/Project.toml", "./docs/src/assets/Project.toml"; force = true)
 
 ENV["GKSwstype"] = "100"
 using Plots
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 
 include("pages.jl")
 
-makedocs(sitename = "DiffEqFlux.jl",
+makedocs(; sitename = "DiffEqFlux.jl",
     authors = "Chris Rackauckas et al.",
     clean = true, doctest = false, linkcheck = true,
     warnonly = [:docs_block, :missing_docs],
     modules = [DiffEqFlux],
-    format = Documenter.HTML(assets = ["assets/favicon.ico"],
+    format = Documenter.HTML(; assets = ["assets/favicon.ico"],
         canonical = "https://docs.sciml.ai/DiffEqFlux/stable/"),
     pages = pages)
 
-deploydocs(repo = "github.com/SciML/DiffEqFlux.jl.git";
+deploydocs(; repo = "github.com/SciML/DiffEqFlux.jl.git",
     push_preview = true)

docs/src/examples/GPUs.md

@@ -23,14 +23,14 @@ u0 = Float32[2.0; 0.0] |> gpu
 prob_gpu = ODEProblem(dudt, u0, tspan, ps)
 
 # Runs on a GPU
-sol_gpu = solve(prob_gpu, Tsit5(), saveat = tsteps)
+sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)
 ```
 
 Or we could directly use the neural ODE layer function, like:
 
 ```julia
 using DiffEqFlux: NeuralODE
-prob_neuralode_gpu = NeuralODE(model, tspan, Tsit5(), saveat = tsteps)
+prob_neuralode_gpu = NeuralODE(model, tspan, Tsit5(); saveat = tsteps)
 ```
 
 If one is using `Lux.Chain`, then the computation takes place on the GPU with
@@ -55,13 +55,13 @@ 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)
+sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)
 ```
 
 or via the NeuralODE struct:
 
 ```julia
-prob_neuralode_gpu = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
+prob_neuralode_gpu = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
 prob_neuralode_gpu(u0, p, st)
 ```
 
@@ -83,22 +83,22 @@ rng = Random.default_rng()
 u0 = Float32[2.0; 0.0]
 datasize = 30
 tspan = (0.0f0, 1.5f0)
-tsteps = range(tspan[1], tspan[2], length = datasize)
+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 = gpu(solve(prob_trueode, Tsit5(), saveat = tsteps))
+ode_data = gpu(solve(prob_trueode, Tsit5(); saveat = tsteps))
 
 dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
 u0 = Float32[2.0; 0.0] |> gpu
 p, st = Lux.setup(rng, dudt2)
 p = p |> ComponentArray |> gpu
 st = st |> gpu
 
-prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
+prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
 
 function predict_neuralode(p)
     gpu(first(prob_neuralode(u0, p, st)))
@@ -119,8 +119,8 @@ callback = function (p, l, pred; doplot = false)
     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")
+    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))
@@ -132,7 +132,7 @@ adtype = Optimization.AutoZygote()
 optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, p)
 result_neuralode = Optimization.solve(optprob,
-    Adam(0.05),
+    Adam(0.05);
     callback = callback,
     maxiters = 300)
 ```

docs/src/examples/augmented_neural_ode.md

@@ -27,8 +27,8 @@ function concentric_sphere(dim, inner_radius_range, outer_radius_range,
         push!(data, reshape(random_point_in_sphere(dim, outer_radius_range...), :, 1))
         push!(labels, -ones(1, 1))
     end
-    data = cat(data..., dims = 2)
-    labels = cat(labels..., dims = 2)
+    data = cat(data...; dims = 2)
+    labels = cat(labels...; dims = 2)
     DataLoader((data |> Flux.gpu, labels |> Flux.gpu); batchsize = batch_size,
         shuffle = true,
         partial = false)
@@ -41,7 +41,7 @@ function construct_model(out_dim, input_dim, hidden_dim, augment_dim)
     node = NeuralODE(Flux.Chain(Flux.Dense(input_dim, hidden_dim, relu),
             Flux.Dense(hidden_dim, hidden_dim, relu),
             Flux.Dense(hidden_dim, input_dim)) |> Flux.gpu,
-        (0.0f0, 1.0f0), Tsit5(), save_everystep = false,
+        (0.0f0, 1.0f0), Tsit5(); save_everystep = false,
         reltol = 1.0f-3, abstol = 1.0f-3, save_start = false) |> Flux.gpu
     node = augment_dim == 0 ? node : AugmentedNDELayer(node, augment_dim)
     return Flux.Chain((x, p = node.p) -> node(x, p),
@@ -54,15 +54,15 @@ end
 function plot_contour(model, npoints = 300)
     grid_points = zeros(Float32, 2, npoints^2)
     idx = 1
-    x = range(-4.0f0, 4.0f0, length = npoints)
-    y = range(-4.0f0, 4.0f0, length = npoints)
+    x = range(-4.0f0, 4.0f0; length = npoints)
+    y = range(-4.0f0, 4.0f0; length = npoints)
     for x1 in x, x2 in y
         grid_points[:, idx] .= [x1, x2]
         idx += 1
     end
     sol = reshape(model(grid_points |> Flux.gpu), npoints, npoints) |> Flux.cpu
 
-    return contour(x, y, sol, fill = true, linewidth = 0.0)
+    return contour(x, y, sol; fill = true, linewidth = 0.0)
 end
 
 loss_node(x, y) = mean((model(x) .- y) .^ 2)
@@ -91,7 +91,7 @@ opt = Adam(0.005)
 println("Training Neural ODE")
 
 for _ in 1:10
-    Flux.train!(loss_node, Flux.params(parameters, model), dataloader, opt, cb = cb)
+    Flux.train!(loss_node, Flux.params(parameters, model), dataloader, opt; cb = cb)
 end
 
 plt_node = plot_contour(model)
@@ -103,7 +103,7 @@ println()
 println("Training Augmented Neural ODE")
 
 for _ in 1:10
-    Flux.train!(loss_node, Flux.params(parameters, model), dataloader, opt, cb = cb)
+    Flux.train!(loss_node, Flux.params(parameters, model), dataloader, opt; cb = cb)
 end
 
 plt_anode = plot_contour(model)
@@ -153,8 +153,8 @@ function concentric_sphere(dim, inner_radius_range, outer_radius_range,
         push!(data, reshape(random_point_in_sphere(dim, outer_radius_range...), :, 1))
         push!(labels, -ones(1, 1))
     end
-    data = cat(data..., dims = 2)
-    labels = cat(labels..., dims = 2)
+    data = cat(data...; dims = 2)
+    labels = cat(labels...; dims = 2)
     return DataLoader((data |> Flux.gpu, labels |> Flux.gpu); batchsize = batch_size,
         shuffle = true,
         partial = false)
@@ -181,7 +181,7 @@ function construct_model(out_dim, input_dim, hidden_dim, augment_dim)
     node = NeuralODE(Flux.Chain(Flux.Dense(input_dim, hidden_dim, relu),
             Flux.Dense(hidden_dim, hidden_dim, relu),
             Flux.Dense(hidden_dim, input_dim)) |> Flux.gpu,
-        (0.0f0, 1.0f0), Tsit5(), save_everystep = false,
+        (0.0f0, 1.0f0), Tsit5(); save_everystep = false,
         reltol = 1.0f-3, abstol = 1.0f-3, save_start = false) |> Flux.gpu
     node = augment_dim == 0 ? node : (AugmentedNDELayer(node, augment_dim) |> Flux.gpu)
     return Flux.Chain((x, p = node.p) -> node(x, p),
@@ -200,15 +200,15 @@ Here, we define a utility to plot our model regression results as a heatmap.
 function plot_contour(model, npoints = 300)
     grid_points = zeros(2, npoints^2)
     idx = 1
-    x = range(-4.0f0, 4.0f0, length = npoints)
-    y = range(-4.0f0, 4.0f0, length = npoints)
+    x = range(-4.0f0, 4.0f0; length = npoints)
+    y = range(-4.0f0, 4.0f0; length = npoints)
     for x1 in x, x2 in y
         grid_points[:, idx] .= [x1, x2]
         idx += 1
     end
     sol = reshape(model(grid_points |> Flux.gpu), npoints, npoints) |> Flux.cpu
 
-    return contour(x, y, sol, fill = true, linewidth = 0.0)
+    return contour(x, y, sol; fill = true, linewidth = 0.0)
 end
 ```
 
@@ -269,7 +269,7 @@ for `20` epochs.
 model, parameters = construct_model(1, 2, 64, 0)
 
 for _ in 1:10
-    Flux.train!(loss_node, Flux.params(model, parameters), dataloader, opt, cb = cb)
+    Flux.train!(loss_node, Flux.params(model, parameters), dataloader, opt; cb = cb)
 end
 ```
 
@@ -288,7 +288,7 @@ a function which can be expressed by the neural ode. For more details and proofs
 model, parameters = construct_model(1, 2, 64, 1)
 
 for _ in 1:10
-    Flux.train!(loss_node, Flux.params(model, parameters), dataloader, opt, cb = cb)
+    Flux.train!(loss_node, Flux.params(model, parameters), dataloader, opt; cb = cb)
 end
 ```
 

docs/src/examples/collocation.md

@@ -20,20 +20,20 @@ rng = Random.default_rng()
 u0 = Float32[2.0; 0.0]
 datasize = 300
 tspan = (0.0f0, 1.5f0)
-tsteps = range(tspan[1], tspan[2], length = datasize)
+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)
+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)
+plot!(tsteps, u'; lw = 5)
 savefig("colloc.png")
 plot(tsteps, du')
 savefig("colloc_du.png")
@@ -63,11 +63,11 @@ optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
 
 result_neuralode = Optimization.solve(optprob,
-    Adam(0.05),
+    Adam(0.05);
     callback = callback,
     maxiters = 10000)
 
-prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
+prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
 nn_sol, st = prob_neuralode(u0, result_neuralode.u, st)
 scatter(tsteps, data')
 plot!(nn_sol)
@@ -88,13 +88,13 @@ optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
 
 numerical_neuralode = Optimization.solve(optprob,
-    Adam(0.05),
+    Adam(0.05);
     callback = callback,
     maxiters = 300)
 
 nn_sol, st = prob_neuralode(u0, numerical_neuralode.u, st)
 scatter(tsteps, data')
-plot!(nn_sol, lw = 5)
+plot!(nn_sol; lw = 5)
 ```
 
 ## Generating the Collocation
@@ -112,20 +112,20 @@ rng = Random.default_rng()
 u0 = Float32[2.0; 0.0]
 datasize = 300
 tspan = (0.0f0, 1.5f0)
-tsteps = range(tspan[1], tspan[2], length = datasize)
+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)
+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)
+plot!(tsteps, u'; lw = 5)
 ```
 
 We can then differentiate the smoothed function to get estimates of the
@@ -165,11 +165,11 @@ optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
 
 result_neuralode = Optimization.solve(optprob,
-    Adam(0.05),
+    Adam(0.05);
     callback = callback,
     maxiters = 10000)
 
-prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
+prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
 nn_sol, st = prob_neuralode(u0, result_neuralode.u, st)
 scatter(tsteps, data')
 plot!(nn_sol)
@@ -196,13 +196,13 @@ optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
 optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))
 
 numerical_neuralode = Optimization.solve(optprob,
-    Adam(0.05),
+    Adam(0.05);
     callback = callback,
     maxiters = 300)
 
 nn_sol, st = prob_neuralode(u0, numerical_neuralode.u, st)
 scatter(tsteps, data')
-plot!(nn_sol, lw = 5)
+plot!(nn_sol; lw = 5)
 ```
 
 This method then has a good global starting position, making it less

docs/src/examples/hamiltonian_nn.md

@@ -12,7 +12,7 @@ Now we make some simplifying assumptions, and assign ``m = 1`` and ``k = 1``. An
 using Lux, DiffEqFlux, OrdinaryDiffEq, Statistics, Plots, Zygote, ForwardDiff, Random,
     ComponentArrays, Optimization, OptimizationOptimisers, IterTools
 
-t = range(0.0f0, 1.0f0, length = 1024)
+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, :)
@@ -46,12 +46,12 @@ res = Optimization.solve(opt_prob, opt, dataloader)
 
 ps_trained = res.u
 
-model = NeuralHamiltonianDE(hnn, (0.0f0, 1.0f0), Tsit5(), save_everystep = false,
+model = NeuralHamiltonianDE(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")
+plot(data[1, :], data[2, :]; lw = 4, label = "Original")
+plot!(pred[1, :], pred[2, :]; lw = 4, label = "Predicted")
 xlabel!("Position (q)")
 ylabel!("Momentum (p)")
 ```
@@ -66,15 +66,15 @@ The HNN predicts the gradients ``(\dot q, \dot p)`` given ``(q, p)``. Hence, we
 using Flux, DiffEqFlux, DifferentialEquations, Statistics, Plots, ReverseDiff, Random,
     IterTools, Lux, ComponentArrays, Optimization, OptimizationOptimisers
 
-t = range(0.0f0, 1.0f0, length = 1024)
+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)
+data = cat(q_t, p_t; dims = 1)
+target = cat(dqdt, dpdt; dims = 1)
 B = 256
 NEPOCHS = 500
 dataloader = ncycle(((selectdim(data, 2, ((i - 1) * B + 1):(min(i * B, size(data, 2)))),
@@ -117,12 +117,12 @@ ps_trained = res.u
 In order to visualize the learned trajectories, we need to solve the ODE. We will use the `NeuralHamiltonianDE` layer, which is essentially a wrapper over `HamiltonianNN` layer, and solves the ODE.
 
 ```@example hamiltonian
-model = NeuralHamiltonianDE(hnn, (0.0f0, 1.0f0), Tsit5(), save_everystep = false,
+model = NeuralHamiltonianDE(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")
+plot(data[1, :], data[2, :]; lw = 4, label = "Original")
+plot!(pred[1, :], pred[2, :]; lw = 4, label = "Predicted")
 xlabel!("Position (q)")
 ylabel!("Momentum (p)")
 ```