diff --git a/Project.toml b/Project.toml index 343f3cd..9dda8c8 100644 --- a/Project.toml +++ b/Project.toml @@ -9,7 +9,6 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" Requires = "ae029012-a4dd-5104-9daa-d747884805df" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" @@ -19,7 +18,6 @@ FiniteDiff = "2.8" ForwardDiff = "0.10" Ipopt = "1.0" Requires = "1.1" -ReverseDiff = "1.9" SparseDiffTools = "1.13" julia = "1.6" diff --git a/docs/src/quickstart.md b/docs/src/quickstart.md index e071ebc..0590aec 100644 --- a/docs/src/quickstart.md +++ b/docs/src/quickstart.md @@ -27,6 +27,9 @@ using SNOW If you want to use Snopt as the optimizer you need to build [Snopt.jl](https://github.com/byuflowlab/Snopt.jl) separately and load the package separate from this one ```using Snopt``` (either before or after ```using SNOW```). It can't be loaded by default because the code is not freely available. In this example we use Ipopt for the optimization as it is freely available. +!!! note "ReverseAD" + +Similar to "Snopt" `ReverseAD` is not available by default. You need to load the `ReverseDiff` package in order to use `ReverseAD` Next, we define the function we wish to optimize. diff --git a/src/ReverseAD.jl b/src/ReverseAD.jl new file mode 100644 index 0000000..1ef6c78 --- /dev/null +++ b/src/ReverseAD.jl @@ -0,0 +1,92 @@ +import .ReverseDiff +""" + createcache(sp::DensePattern, dtype::ReverseAD, func!, nx, ng) + +Cache for dense jacobian using reverse-mode AD. + +# Arguments +- `func!::function`: function of form: f = func!(g, x) +- `nx::Int`: number of design variables +- `ng::Int`: number of constraints +""" +function createcache(sp::DensePattern, dtype::ReverseAD, func!, nx, ng) + + function combine!(fg, x) + fg[1] = func!(@view(fg[2:end]), x) + end + + g = zeros(1 + ng) + x = zeros(nx) + + f_tape = ReverseDiff.JacobianTape(combine!, g, x) + cache = ReverseDiff.compile(f_tape) + J = DiffResults.JacobianResult(g, x) + + return DenseCache(combine!, g, J, cache, dtype) +end + +""" + evaluate!(g, df, dg, x, cache::DenseCache{T1,T2,T3,T4,T5} where {T1,T2,T3,T4,T5<:ReverseAD}) + +evaluate function and derivatives for a dense jacobian with reverse-mode AD + +# Arguments +- `g::Vector{Float}`: constraints, modified in place +- `df::Vector{Float}`: objective gradient, modified in place +- `dg::Vector{Float}`: constraint jacobian, modified in place (order specified by sparsity pattern) +- `x::Vector{Float}`: design variables, input +- `cache::DenseCache`: cache generated by `createcache` +""" +function evaluate!(g, df, dg, x, cache::DenseCache{T1,T2,T3,T4,T5} + where {T1,T2,T3,T4,T5<:ReverseAD}) + + ReverseDiff.jacobian!(cache.Jwork, cache.cache, x) + fg = DiffResults.value(cache.Jwork) # reference not copy + J = DiffResults.jacobian(cache.Jwork) # reference not copy + f = fg[1] + g[:] = fg[2:end] + df[:] = J[1, :] + dg[:] = J[2:end, :][:] + + return f +end + +""" + gradientcache(dtype::ReverseAD, func!, nx, ng) + +Cache for gradient using ReverseDiff + +# Arguments +- `func!::function`: function of form: f = func!(g, x) +- `nx::Int`: number of design variables +- `ng::Int`: number of constraints +""" +function gradientcache(dtype::ReverseAD, func!, nx, ng) + + function obj(x) + return func!(zeros(eltype(x[1]), ng), x) + end + + f_tape = ReverseDiff.GradientTape(obj, zeros(nx)) + cache = ReverseDiff.compile(f_tape) + + return GradOrJacCache(func!, nothing, cache, dtype) +end + +""" + gradient!(df, x, cache::GradOrJacCache{T1,T2,T3,T4} where {T1,T2,T3,T4<:ReverseAD}) + +evaluate gradient using ReverseDiff + +# Arguments +- `df::Vector{Float}`: objective gradient, modified in place +- `x::Vector{Float}`: design variables, input +- `cache::DenseCache`: cache generated by `gradientcache` +""" +function gradient!(df, x, cache::GradOrJacCache{T1,T2,T3,T4} + where {T1,T2,T3,T4<:ReverseAD}) + + ReverseDiff.gradient!(df, cache.cache, x) + + return nothing +end diff --git a/src/SNOW.jl b/src/SNOW.jl index 41b5fc1..71545fa 100644 --- a/src/SNOW.jl +++ b/src/SNOW.jl @@ -1,10 +1,9 @@ module SNOW -using ForwardDiff -using ReverseDiff +import ForwardDiff # using Zygote using DiffResults -using FiniteDiff +import FiniteDiff using SparseArrays using SparseDiffTools using Requires @@ -28,6 +27,7 @@ export IPOPT # conditionally load Snopt function __init__() @require Snopt="0e9dc826-d618-11e8-1f57-c34e87fde2c0" include("snopt.jl") + @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include("ReverseAD.jl") end end diff --git a/src/derivatives.jl b/src/derivatives.jl index 4991131..8fdf94e 100644 --- a/src/derivatives.jl +++ b/src/derivatives.jl @@ -6,6 +6,14 @@ abstract type AbstractDiffMethod end struct ForwardAD <: AbstractDiffMethod end + +""" + + ReverseAD + +In order to use ReverseAD, make sure the ReverseDiff package is loaded. +For instance by adding `import ReverseDiff` to your code. +""" struct ReverseAD <: AbstractDiffMethod end # struct RevZyg <: AbstractDiffMethod end # only used for gradients (not jacobians) struct ForwardFD <: AbstractDiffMethod end @@ -13,7 +21,7 @@ struct CentralFD <: AbstractDiffMethod end struct ComplexStep <: AbstractDiffMethod end struct UserDeriv <: AbstractDiffMethod end # user-specified derivatives -FD = Union{ForwardFD, CentralFD, ComplexStep} +const FD = Union{ForwardFD, CentralFD, ComplexStep} """ convert to type used in FiniteDiff package @@ -221,58 +229,7 @@ function evaluate!(g, df, dg, x, cache::DenseCache{T1,T2,T3,T4,T5} end -""" - createcache(sp::DensePattern, dtype::ReverseAD, func!, nx, ng) - -Cache for dense jacobian using reverse-mode AD. - -# Arguments -- `func!::function`: function of form: f = func!(g, x) -- `nx::Int`: number of design variables -- `ng::Int`: number of constraints -""" -function createcache(sp::DensePattern, dtype::ReverseAD, func!, nx, ng) - - function combine!(fg, x) - fg[1] = func!(@view(fg[2:end]), x) - end - - g = zeros(1 + ng) - x = zeros(nx) - - f_tape = ReverseDiff.JacobianTape(combine!, g, x) - cache = ReverseDiff.compile(f_tape) - J = DiffResults.JacobianResult(g, x) - - return DenseCache(combine!, g, J, cache, dtype) -end - - -""" - evaluate!(g, df, dg, x, cache::DenseCache{T1,T2,T3,T4,T5} where {T1,T2,T3,T4,T5<:ReverseAD}) - -evaluate function and derivatives for a dense jacobian with reverse-mode AD -# Arguments -- `g::Vector{Float}`: constraints, modified in place -- `df::Vector{Float}`: objective gradient, modified in place -- `dg::Vector{Float}`: constraint jacobian, modified in place (order specified by sparsity pattern) -- `x::Vector{Float}`: design variables, input -- `cache::DenseCache`: cache generated by `createcache` -""" -function evaluate!(g, df, dg, x, cache::DenseCache{T1,T2,T3,T4,T5} - where {T1,T2,T3,T4,T5<:ReverseAD}) - - ReverseDiff.jacobian!(cache.Jwork, cache.cache, x) - fg = DiffResults.value(cache.Jwork) # reference not copy - J = DiffResults.jacobian(cache.Jwork) # reference not copy - f = fg[1] - g[:] = fg[2:end] - df[:] = J[1, :] - dg[:] = J[2:end, :][:] - - return f -end """ @@ -374,46 +331,6 @@ struct GradOrJacCache{T1,T2,T3,T4} dtype::T4 end -""" - gradientcache(dtype::ReverseAD, func!, nx, ng) - -Cache for gradient using ReverseDiff - -# Arguments -- `func!::function`: function of form: f = func!(g, x) -- `nx::Int`: number of design variables -- `ng::Int`: number of constraints -""" -function gradientcache(dtype::ReverseAD, func!, nx, ng) - - function obj(x) - return func!(zeros(eltype(x[1]), ng), x) - end - - f_tape = ReverseDiff.GradientTape(obj, zeros(nx)) - cache = ReverseDiff.compile(f_tape) - - return GradOrJacCache(func!, nothing, cache, dtype) -end - - -""" - gradient!(df, x, cache::GradOrJacCache{T1,T2,T3,T4} where {T1,T2,T3,T4<:ReverseAD}) - -evaluate gradient using ReverseDiff - -# Arguments -- `df::Vector{Float}`: objective gradient, modified in place -- `x::Vector{Float}`: design variables, input -- `cache::DenseCache`: cache generated by `gradientcache` -""" -function gradient!(df, x, cache::GradOrJacCache{T1,T2,T3,T4} - where {T1,T2,T3,T4<:ReverseAD}) - - ReverseDiff.gradient!(df, cache.cache, x) - - return nothing -end # """ diff --git a/test/Project.toml b/test/Project.toml index 548a171..b6cb856 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,3 +1,3 @@ [deps] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" diff --git a/test/runtests.jl b/test/runtests.jl index e61e749..b4852b2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,6 @@ using SNOW using Test -using Zygote +import ReverseDiff checkallocations = false snopttest = false @@ -264,4 +264,4 @@ end # x = rand(5) # SNOW.get_dvs(prob, x) -# end \ No newline at end of file +# end