diff --git a/examples/controlproblem.jl b/examples/controlproblem.jl index 861703f..f1afbac 100644 --- a/examples/controlproblem.jl +++ b/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,7 +116,7 @@ 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)) @@ -105,8 +124,15 @@ function solve_2D(x0, d, h, kappa, alpha, iters = 30) # 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 \ No newline at end of file diff --git a/src/DualArrays.jl b/src/DualArrays.jl index 360c597..2161772 100644 --- a/src/DualArrays.jl +++ b/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 diff --git a/src/arithmetic.jl b/src/arithmetic.jl index d7f8142..f835ffb 100644 --- a/src/arithmetic.jl +++ b/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 diff --git a/src/indexing.jl b/src/indexing.jl index b357811..2cbf31e 100644 --- a/src/indexing.jl +++ b/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) diff --git a/test/broadcast_test.jl b/test/broadcast_test.jl index e368e08..1d66fd3 100644 --- a/test/broadcast_test.jl +++ b/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 \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 510c265..ad16e8c 100644 --- a/test/runtests.jl +++ b/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])