Add nested dual implementation of poisson control

examples/controlproblem.jl

@@ -49,6 +49,18 @@ using LinearAlgebra, DualArrays, Plots
 # If we propagate f as a DualVector above, we get the f.jacobian as the Hessian of J,
 # and this enables us to solve via newton's method.
 
+"""
+This calculates the scalar objective
+
+    (1/2) * ||u - d||² + (alpha / 2) * ||f||²
+
+where u solves K * u = f.
+"""
+function objective(K, f, d, alpha, h)
+    u = K \ f
+    return (h / 2) * sum((u - d) .^ 2) + (alpha * h / 2) * sum(f .^ 2)
+end
+
 """
 This calculates the gradient of the objective as detailed above.
 K: the discretised Laplacian
@@ -73,13 +85,20 @@ kappa: diffusivity constant
 alpha: regularisation parameter
 iters: number of iterations to run newton's method.
 """
-function solve_1D(f0, d, h, kappa, alpha, iters = 30)
+function solve_1D(f0, d, h, kappa, alpha, iters = 30; method = :adjoint)
     fvector = copy(f0)
     n = length(f0)
     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.data \ grads.value
+        gradient, hess = if method == :adjoint
+            grads = grad_objective(K, DualVector(fvector, I(length(fvector))), d, alpha)
+            (grads.value, grads.jacobian.data)
+        elseif method == :hessian
+            J = f -> objective(K, f, d, alpha, h)
+            gradient = J(DualVector(fvector, I(length(fvector)))).partials
+            (gradient, hessian(J, fvector))
+        end
+        step = hess \ gradient
         fvector -= step 
     end
 
@@ -97,16 +116,23 @@ kappa: diffusivity constant
 alpha: regularisation parameter
 iters: number of iterations to run newton's method.
 """
-function solve_2D(x0, d, h, kappa, alpha, iters = 30)
+function solve_2D(x0, d, h, kappa, alpha, iters = 30; method = :adjoint)
     fvector = copy(x0)
     n = isqrt(length(x0))
     T = Tridiagonal(-ones(n - 1), 2 * ones(n), -ones(n - 1))
     # The 2D Discretised Laplacian can be constructed as below
     # source: https://www.petercheng.me/blog/discrete-laplacian-matrix
     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.data \ grads.value
+        gradient, hess = if method == :adjoint
+            grads = grad_objective(K, DualVector(fvector, I(length(fvector))), d, alpha)
+            (grads.value, grads.jacobian.data)
+        elseif method == :hessian
+            J = f -> objective(K, f, d, alpha, h)
+            gradient = J(DualVector(fvector, I(length(fvector)))).partials
+            (gradient, hessian(J, fvector))
+        end
+        step = hess \ gradient
         fvector -= step 
     end
 
@@ -115,7 +141,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(save=undef)
+function plot_solution_1D(save=undef; method = :adjoint)
 
     # setup problem
     kappa = 1
@@ -131,17 +157,18 @@ function plot_solution_1D(save=undef)
     coeff = (kappa * pi^2) / (1 + alpha * kappa^2 * pi^4)
     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 Control", title="1D Poisson Control Solution")
-    plot!(x[2:end-1], fexact, label="Exact Control")
+    f = solve_1D(zeros(length(d)), d, h, kappa, alpha; method=method)
+    plt = plot(x[2:end-1], f, label="Computed Control ($(String(method)))", title="1D Poisson Control Solution")
+    plot!(plt, x[2:end-1], fexact, label="Exact Control")
+    display(plt)
     if save !== undef
-        savefig(save)
+        savefig(plt, 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(save=undef)
+function plot_solution_2D(save=undef; method = :adjoint)
 
     # setup problem
     kappa = 1
@@ -162,13 +189,14 @@ function plot_solution_2D(save=undef)
     # analytical optimal control as specified in dolfin-adjoint
     fexact = (1 / (1 + 4 * alpha * pi^4)) .* sin.(pi .* X) .* sin.(pi .* Y)
 
-    f = solve_2D(zeros(n * n), vec(d), h, kappa, alpha)
+    f = solve_2D(zeros(n * n), vec(d), h, kappa, alpha; method=method)
     F = reshape(f, n, n)
 
-    p1 = heatmap(x, x, F, title="Computed Control", aspect_ratio=1)
+    p1 = heatmap(x, x, F, title="Computed Control ($(String(method)))", aspect_ratio=1)
     p2 = heatmap(x, x, fexact, title="Exact Control", aspect_ratio=1)
-    plot(p1, p2, layout=(1, 2), size=(900, 400), plot_title="2D Poisson Control Solution", plot_titlevspan=0.12)
+    plt = plot(p1, p2, layout=(1, 2), size=(900, 400), plot_title="2D Poisson Control Solution", plot_titlevspan=0.12)
+    display(plt)
     if save !== undef
-        savefig(save)
+        savefig(plt, save)
     end
 end

src/DualArrays.jl

@@ -15,7 +15,7 @@ module DualArrays
 
 export DualVector, DualMatrix, Dual, jacobian, ArrayOperator, hessian
 
-import Base: +, -, ==, getindex, size, axes, broadcasted, show, sum, vcat, convert, *, isapprox, \, eltype, transpose, permutedims
+import Base: +, -, ==, getindex, setindex!, size, axes, broadcasted, show, sum, vcat, convert, *, isapprox, \, eltype, transpose, permutedims
 
 using LinearAlgebra, ArrayLayouts, FillArrays, DiffRules, TensorOperations, SparseArrayKit
 

src/arithmetic.jl

@@ -115,7 +115,16 @@ LinearAlgebra.dot(x::DualVector, y::AbstractVector) = Dual(dot(x.value, y), tran
 LinearAlgebra.dot(x::AbstractVector, y::DualVector) = Dual(dot(x, y.value), transpose(y.jacobian) * x)
 
 # solve
-\(x::AbstractMatrix, y::DualVector) = DualVector(x \ y.value, ArrayOperator{1}(x \ y.jacobian.data))
+function \(x::AbstractMatrix, y::DualVector)
+    DualVector(x \ y.value, ArrayOperator{1}(x \ y.jacobian.data))
+end
+
+function \(x::AbstractMatrix, y::DualMatrix)
+    jac = y.jacobian.data
+    reshaped = reshape(jac, size(jac, 1), :)
+    solved = x \ reshaped
+    DualMatrix(x \ y.value, ArrayOperator{2}(reshape(solved, size(jac))))
+end
 
 
 

src/indexing.jl

@@ -78,6 +78,12 @@ function getindex(x::DualMatrix, y::Vararg{Union{Colon, UnitRange}})
     DualMatrix(newval, newjac)
 end
 
+function setindex!(x::DualVector, y::Number, i::Int)
+    x.value[i] = y
+    x.jacobian.data[i, :] .= 0
+    return x
+end
+
 # DualArray array interface
 for op in (:size, :axes)
     @eval $op(a::DualArray) = $op(a.value)

test/broadcast_test.jl

@@ -71,4 +71,13 @@ using DualArrays, Test, SparseArrays, LinearAlgebra
     @test s.jacobian.data ≈ [0.0 -2.0 -3.0; -1.0 -1.0 -3.0; -1.0 -2.0 -2.0]
     @test m.jacobian.data ≈ [3.0 2.0 3.0; 2.0 6.0 6.0; 3.0 6.0 11.0]
     @test div.jacobian.data ≈ [0.25 -0.5 -0.75; -0.5 -0.5 -1.5; -0.75 -1.5 -1.75]
+end
+
+@testset "DualVector setindex!" begin
+    x = DualVector([1, 2, 3], [1 0 0; 0 1 0; 0 0 1])
+
+    x[2] = 4
+
+    @test x.value == [1, 4, 3]
+    @test x.jacobian.data == [1 0 0; 0 0 0; 0 0 1]
 end

test/runtests.jl

@@ -112,6 +112,16 @@ using DualArrays: ArrayOperator
         @test A \ b == DualVector([1, 1], [1.5 2; 1.5 2])
     end
 
+    @testset "Solve (DualMatrix)" begin
+        A = [1 0;0 2]
+        dm = DualMatrix([1.0 2.0; 3.0 4.0], zeros(2, 2, 2))
+
+        res = A \ dm
+
+        @test res.value == A \ dm.value
+        @test res.jacobian.data == zeros(2, 2, 2)
+    end
+
     @testset "Matrix multiplication" begin
         M = [1 1; 1 1]
         d = DualVector([2, 3], [4 5; 6 7])