Polish examples (#32)

* initial commit

* fix keyword mismatch

* improve coverage

Author: max-vassili3v

Project.toml

@@ -9,16 +9,20 @@ BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SparseArrayKit = "a9a3c162-d163-4c15-8926-b8794fbefed2"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
+
 [compat]
 DiffRules = "1.15.1"
 ForwardDiff = "1.2.2"
 Lux = "1.29.1"
 Plots = "1.41.1"
+SparseArrayKit = "0.4.3"
 SparseArrays = "1.10"
-TensorOperations = "5.5.2"
+TensorOperations = "5.6.1"
+Zygote = "0.7.10"
 
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
@@ -27,7 +31,8 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-examples = ["Plots", "ForwardDiff", "BandedMatrices", "Lux"]
+examples = ["Plots", "ForwardDiff", "BandedMatrices", "Lux", "Zygote"]
 test = ["Test", "ForwardDiff", "BandedMatrices", "SparseArrays"]

examples/catenary.jl

@@ -1,4 +1,4 @@
-using DualArrays, LinearAlgebra, Plots, BenchmarkTools
+using DualArrays, LinearAlgebra, Plots, BenchmarkTools, BandedMatrices, ForwardDiff, Zygote
 
 """
 Solving the Catenary Problem using DualArrays.jl
@@ -11,6 +11,19 @@ The computation of finite differences keeps the Jacobian sparse,
 so we can use DualArrays.jl
 """
 
+const plot_args = (
+    xscale = :log10,
+    yscale = :log10,
+    lw = 2.5,
+    marker = :circle,
+    color = :steelblue4,
+    legend = :topleft,
+    framestyle = :box,
+    size = (960, 600),
+    left_margin = 10Plots.mm,
+    bottom_margin = 10Plots.mm,
+)
+
 function L(y, h, alpha, beta)
     """
     Evaluate functional with boundary conditions (alpha, beta)
@@ -24,33 +37,109 @@ function L(y, h, alpha, beta)
     return y_mid .* sqrt.(1 .+ dy.^2)
 end
 
-function learn_catenary(h = 0.1, alpha = cosh(-1), beta = cosh(1), epochs = 2000, lr = 0.01)
+function learn_catenary_dualarrays(h = 0.1, alpha = cosh(-1), beta = cosh(1), epochs = 2000, lr = 0.01)
+    n = Int(2 / h) - 1
+    y = ones(n) * (alpha + beta) / 2
+    for _ = 1:epochs
+        jac = DualArrays.jacobian(y -> L(y, h, alpha, beta), y, BandedMatrix)
+        grads = h * sum(jac, dims=1)
+        y -= lr * vec(grads)
+    end
+    return y
+end
+
+function learn_catenary_forwarddiff(h = 0.1, alpha = cosh(-1), beta = cosh(1), epochs = 2000, lr = 0.01)
+    n = Int(2 / h) - 1
+    y = ones(n) * (alpha + beta) / 2
+    for _ = 1:epochs
+        jac = ForwardDiff.jacobian(y -> L(y, h, alpha, beta), y)
+        grads = h * sum(jac, dims=1)
+        y -= lr * vec(grads)
+    end
+    return y
+end
+
+function learn_catenary_zygote(h = 0.1, alpha = cosh(-1), beta = cosh(1), epochs = 2000, lr = 0.01)
     n = Int(2 / h) - 1
     y = ones(n) * (alpha + beta) / 2
     for _ = 1:epochs
-        jac = jacobian(y -> L(y, h, alpha, beta), y, id="banded")
+        jac = only(Zygote.jacobian(y -> L(y, h, alpha, beta), y))
         grads = h * sum(jac, dims=1)
         y -= lr * vec(grads)
     end
     return y
 end
 
-function plot_solution(h)
+function plot_solution(save=undef;h = 0.1, alpha = cosh(-1), beta = cosh(1), epochs = 5000, lr = 0.02)
     x = collect(-1+h:h:1-h)
-    y = learn_catenary()
-    plot(x, y, label = "Approximation")
-    plot!(x, cosh.(x), label = "Exact Solution")
+    y = learn_catenary_dualarrays(h, alpha, beta, epochs, lr)
+    plot(x, y, label = "Approximate solution (DualArrays)", title="Catenary Solution", legend=:topleft)
+    plot!(x, cosh.(x), label = "Exact Solution (cosh(x))", ls = :dash)
+    if save !== undef
+        savefig(save)
+    end
 end
 
-function plot_times()
-    hs = [0.1, 0.05, 0.02, 0.01, 0.005, 0.002]
+function plot_times(save=undef; hs = [0.04, 0.02, 0.01, 0.005, 0.0025, 0.00125], epochs = 200, lr = 0.01)
     ns = Int.(2 ./ hs) .- 1
     dualvector_times = Float64[]
+    forwarddiff_times = Float64[]
+    zygote_times = Float64[]
 
     for (h, n) in zip(hs, ns)
         println("Computing solution with h = $h, n = $n")
-        push!(dualvector_times, @belapsed learn_catenary($h, cosh(-1), cosh(1), 2000, 0.01))
+        push!(dualvector_times, @belapsed learn_catenary_dualarrays($h, cosh(-1), cosh(1), $epochs, $lr))
+        push!(forwarddiff_times, @belapsed learn_catenary_forwarddiff($h, cosh(-1), cosh(1), $epochs, $lr))
+        push!(zygote_times, @belapsed learn_catenary_zygote($h, cosh(-1), cosh(1), $epochs, $lr))
+    end
+
+    plot(
+        ns,
+        dualvector_times,
+        ;
+        label = "DualArrays",
+        title = "Catenary Gradient Descent Runtime",
+        xlabel = "Number of points (n)",
+        ylabel = "Runtime (seconds)",
+        yticks = 10 .^ collect(-3:0.25:2),
+        xticks = 10 .^ collect(1.5:0.25:3.5),
+        plot_args...,
+    )
+    plot!(ns, forwarddiff_times, label = "ForwardDiff", lw = 2.5, marker = :square, color = :darkorange2)
+    plot!(ns, zygote_times, label = "Zygote", lw = 2.5, marker = :diamond, color = :forestgreen)
+    if save !== undef
+        savefig(save)
     end
+end
 
-    plot(ns, dualvector_times, label="DualArrays")
+function plot_memory(save=undef; hs = [0.04, 0.02, 0.01, 0.005, 0.0025, 0.00125], epochs = 200, lr = 0.01)
+    ns = Int.(2 ./ hs) .- 1
+    dualvector_memory = Float64[]
+    forwarddiff_memory = Float64[]
+    zygote_memory = Float64[]
+
+    for (h, n) in zip(hs, ns)
+        println("Computing solution with h = $h, n = $n")
+        push!(dualvector_memory, @ballocated learn_catenary_dualarrays($h, cosh(-1), cosh(1), $epochs, $lr))
+        push!(forwarddiff_memory, @ballocated learn_catenary_forwarddiff($h, cosh(-1), cosh(1), $epochs, $lr))
+        push!(zygote_memory, @ballocated learn_catenary_zygote($h, cosh(-1), cosh(1), $epochs, $lr))
+    end
+
+    plot(
+        ns,
+        dualvector_memory,
+        ;
+        label = "DualArrays",
+        title = "Catenary Gradient Descent Memory",
+        xlabel = "Number of points (n)",
+        ylabel = "Allocated bytes",
+        yticks = 10 .^ collect(6:0.5:11),
+        xticks = 10 .^ collect(1.5:0.25:3.5),
+        plot_args...,
+    )
+    plot!(ns, forwarddiff_memory, label = "ForwardDiff", lw = 2.5, marker = :square, color = :darkorange2)
+    plot!(ns, zygote_memory, label = "Zygote", lw = 2.5, marker = :diamond, color = :forestgreen)
+    if save !== undef
+        savefig(save)
+    end
 end

examples/controlproblem.jl

@@ -12,7 +12,7 @@ using LinearAlgebra, DualArrays, Plots
 # Provided f ∈ L²(Ω), there exists a unique solution u ∈ H₀¹(Ω) to the boundary value problem.
 # A problem arises that given a desired solution d, we want to find f such that the solution u
 # is approximately d. An intuitive way of asking this is: "given that we want a certain temperature
-# field in a room, how do we heat or cool the walls?"
+# field in a room, how do we position heat sources/sinks?"
 #
 # A way of doing this is to define an objective to be minimised given by:
 #
@@ -79,7 +79,7 @@ function solve_1D(f0, d, h, kappa, alpha, iters = 30)
     K = (kappa / h ^ 2) * Tridiagonal(-ones(n - 1), 2 * ones(n), -ones(n - 1))
     for _ = 1:iters
         grads = grad_objective(K, DualVector(fvector, I(length(fvector))), d, alpha)
-        step = grads.jacobian \ grads.value
+        step = grads.jacobian.data \ grads.value
         fvector -= step 
     end
 
@@ -106,7 +106,7 @@ function solve_2D(x0, d, h, kappa, alpha, iters = 30)
     K = (kappa / h ^ 2) * (kron(I(n), T) + kron(T, I(n)))
     for _ = 1:iters
         grads = grad_objective(K, DualVector(fvector, I(length(fvector))), d, alpha)
-        step = grads.jacobian \ grads.value
+        step = grads.jacobian.data \ grads.value
         fvector -= step 
     end
 
@@ -115,7 +115,7 @@ end
 
 # Solve the 1D Poisson control on the unit interval and plot the solution.
 # We compare it to a known analytical solution.
-function plot_solution_1D()
+function plot_solution_1D(save=undef)
 
     # setup problem
     kappa = 1
@@ -132,13 +132,16 @@ function plot_solution_1D()
     fexact = coeff .* sin.(pi .* x[2:end-1])
 
     f = solve_1D(zeros(length(d)), d, h, kappa, alpha)
-    plot(x[2:end-1], f, label="Computed Solution")
-    plot!(x[2:end-1], fexact, label="Exact Solution")
+    plot(x[2:end-1], f, label="Computed Control", title="1D Poisson Control Solution")
+    plot!(x[2:end-1], fexact, label="Exact Control")
+    if save !== undef
+        savefig(save)
+    end
 end
 
 # Solve the 2D Poisson control on the unit square and plot the solution.
 # We compare it to a known analytical solution given in the dolfin-adjoint example.
-function plot_solution_2D()
+function plot_solution_2D(save=undef)
 
     # setup problem
     kappa = 1
@@ -164,5 +167,8 @@ function plot_solution_2D()
 
     p1 = heatmap(x, x, F, title="Computed Control", aspect_ratio=1)
     p2 = heatmap(x, x, fexact, title="Exact Control", aspect_ratio=1)
-    plot(p1, p2, layout=(1, 2), size=(900, 400))
+    plot(p1, p2, layout=(1, 2), size=(900, 400), plot_title="2D Poisson Control Solution", plot_titlevspan=0.12)
+    if save !== undef
+        savefig(save)
+    end
 end

examples/jacobians.jl

@@ -9,7 +9,7 @@
 #
 ###
 
-using DualArrays, ForwardDiff, Plots, BenchmarkTools
+using DualArrays, ForwardDiff, Plots, BenchmarkTools, BandedMatrices
 
 f(x) = sin.(x)
 
@@ -20,7 +20,7 @@ function plot_times()
 
     for n in ns
         println("Computing jacobian for n = $n")
-        push!(dualarray_times, @belapsed DualArrays.jacobian(f, rand($n)))
+        push!(dualarray_times, @belapsed DualArrays.jacobian(f, rand($n), BandedMatrix))
         push!(forwarddiff_times, @belapsed ForwardDiff.jacobian(f, rand($n)))
     end
 

examples/nestedduals.jl

@@ -42,14 +42,14 @@
 #
 # We now implement an example using f(x) = x[1] * x[2] below.
 
-using DualArrays
+using DualArrays, SparseArrayKit
 f(x) = x[1] * x[2]
 
 # a + Bϵ
 x = DualVector([2, 3], [1 0; 0 1])
 
 # C + Dϵ
-y = DualMatrix([1 0;0 1], zeros(2,2,2))
+y = DualMatrix([1 0;0 1], SparseArray{Float64, 3}(undef, (2, 2, 2)))
 
 # Setup the nested dual vector (a + Bϵ) + (C + Dϵ)ϵ'
 z = DualVector(x, y)
@@ -65,7 +65,7 @@ gradient = result.value.partials
 # equivalent to c
 gradient2 = result.partials.value
 # equivalent to d
-hessian = result.partials.jacobian.data
+_hessian = result.partials.jacobian.data
 
 # We expect this to be 6
 print("Value: ", value, "\n")
@@ -74,5 +74,5 @@ print("Gradient: ", gradient, "\n")
 # We expect this to be the same
 print("Gradient: ", gradient2, "\n")
 # We expect this to be [0 1; 1 0]
-print("Hessian: ", hessian, "\n")
+print("Hessian: ", _hessian, "\n")