Skip to content

Commit

Permalink
Improve and extend autodiff interface. (#98)
Browse files Browse the repository at this point in the history
  • Loading branch information
@pkofod
pkofod authored Feb 7, 2019
1 parent 12ceb15 commit 46fd262
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 32 deletions.
14 changes: 14 additions & 0 deletions src/NLSolversBase.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,20 @@ export AbstractObjective,
export AbstractConstraints, OnceDifferentiableConstraints,
TwiceDifferentiableConstraints, ConstraintBounds

function diffeqdiff_fdtype(autodiff)
if autodiff == :finiteforward
fdtype = Val{:forward}
elseif autodiff == :finitecomplex
fdtype = Val{:complex}
elseif any(autodiff .== (:finite, :central, :finitecentral))
fdtype = Val{:central}
end
fdtype
end

is_finitediff(autodiff) = autodiff ∈ (:central, :finite, :finiteforward, :finitecomplex)
is_forwarddiff(autodiff) = autodiff ∈ (:forward, :forwarddiff, true)

function x_of_nans(x)
x_out = fill!(similar(x), NaN)
x_out
Expand Down
42 changes: 28 additions & 14 deletions src/objective_types/oncedifferentiable.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,20 @@ function OnceDifferentiable(f, x_seed::AbstractArray{T},
autodiff) where T

if typeof(f) <: Union{InplaceObjective, NotInplaceObjective}

fF = make_f(f, x_seed, F)
dfF = make_df(f, x_seed, F)
fdfF = make_fdf(f, x_seed, F)

return OnceDifferentiable(fF, dfF, fdfF, x_seed, F, DF)
else
if any(autodiff .== (:finite, :central))
# TODO: Allow user to specify Val{:central}, Val{:forward}, :Val{:complex} (requires care when using :forward I think)
gcache = DiffEqDiffTools.GradientCache(x_seed, x_seed, Val{:central})
if is_finitediff(autodiff)

# Figure out which Val-type to use for DiffEqDiffTools based on our
# symbol interface.
fdtype = diffeqdiff_fdtype(autodiff)
gcache = DiffEqDiffTools.GradientCache(x_seed, x_seed, fdtype)

function g!(storage, x)
DiffEqDiffTools.finite_difference_gradient!(storage, f, x, gcache)
return
Expand All @@ -50,7 +56,7 @@ function OnceDifferentiable(f, x_seed::AbstractArray{T},
g!(storage, x)
return f(x)
end
elseif autodiff == :forward
elseif is_forwarddiff(autodiff)
gcfg = ForwardDiff.GradientConfig(f, x_seed)
g! = (out, x) -> ForwardDiff.gradient!(out, f, x, gcfg)

Expand All @@ -62,6 +68,7 @@ function OnceDifferentiable(f, x_seed::AbstractArray{T},
else
error("The autodiff value $autodiff is not support. Use :finite or :forward.")
end

return OnceDifferentiable(f, g!, fg!, x_seed, F, DF)
end
end
Expand All @@ -87,33 +94,40 @@ function OnceDifferentiable(f, x::AbstractArray, F::AbstractArray, DF::AbstractA
fdfF = make_fdf(f, x, F)
return OnceDifferentiable(fF, dfF, fdfF, x, F, DF)
else
if autodiff == :central || autodiff == :finite
central_cache = DiffEqDiffTools.JacobianCache(similar(x), similar(F), similar(F))
function fj!(F, J, x)
if is_finitediff(autodiff)

# Figure out which Val-type to use for DiffEqDiffTools based on our
# symbol interface.
fdtype = diffeqdiff_fdtype(autodiff)

central_cache = DiffEqDiffTools.JacobianCache(similar(x), similar(F), similar(F), fdtype)
function fj_finitediff!(F, J, x)
f(F, x)
DiffEqDiffTools.finite_difference_jacobian!(J, f, x, central_cache)
F
end
function j!(J, x)
function j_finitediff!(J, x)
F_cache = similar(F)
fj!(F_cache, J, x)
fj_finitediff!(F_cache, J, x)
end
return OnceDifferentiable(f, j!, fj!, x, F, DF)
elseif autodiff == :forward || autodiff == true
return OnceDifferentiable(f, j_finitediff!, fj_finitediff!, x, F, DF)

elseif is_forwarddiff(autodiff)

jac_cfg = ForwardDiff.JacobianConfig(f, F, x, chunk)
ForwardDiff.checktag(jac_cfg, f, x)

F2 = copy(F)
function g!(J, x)
function j_forwarddiff!(J, x)
ForwardDiff.jacobian!(J, f, F2, x, jac_cfg, Val{false}())
end
function fg!(F, J, x)
function fj_forwarddiff!(F, J, x)
jac_res = DiffResults.DiffResult(F, J)
ForwardDiff.jacobian!(jac_res, f, F2, x, jac_cfg, Val{false}())
DiffResults.value(jac_res)
end

return OnceDifferentiable(f, g!, fg!, x, F, DF)
return OnceDifferentiable(f, j_forwarddiff!, fj_forwarddiff!, x, F, DF)
else
error("The autodiff value $(autodiff) is not supported. Use :finite or :forward.")
end
Expand Down
45 changes: 27 additions & 18 deletions src/objective_types/twicedifferentiable.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,19 +50,21 @@ function TwiceDifferentiable(f, g,
n_x = length(x_seed)

g! = df!_from_df(g, F, inplace)

fg! = make_fdf(x_seed, F, f, g!)

if autodiff == :finite
# TODO: Create / request Hessian functionality in DiffEqDiffTools?
# (Or is it better to use the finite difference Jacobian of a gradient?)
# TODO: Allow user to specify Val{:central}, Val{:forward}, :Val{:complex}
jcache = DiffEqDiffTools.JacobianCache(x_seed, Val{:central})
if is_finitediff(autodiff)

# Figure out which Val-type to use for DiffEqDiffTools based on our
# symbol interface.
fdtype = diffeqdiff_fdtype(autodiff)

jcache = DiffEqDiffTools.JacobianCache(x_seed, fdtype)
function h!(storage, x)
DiffEqDiffTools.finite_difference_jacobian!(storage, g!, x, jcache)
return
end
elseif autodiff == :forward

elseif is_forwarddiff(autodiff)
hcfg = ForwardDiff.HessianConfig(f, similar(x_seed))
h! = (out, x) -> ForwardDiff.hessian!(out, f, x, hcfg)
else
Expand All @@ -76,16 +78,18 @@ TwiceDifferentiable(d::NonDifferentiable, x_seed::AbstractVector{T} = d.x_f, F::

function TwiceDifferentiable(d::OnceDifferentiable, x_seed::AbstractVector{T} = d.x_f,
F::Real = real(zero(T)); autodiff = :finite) where T<:Real
if autodiff == :finite
# TODO: Create / request Hessian functionality in DiffEqDiffTools?
# (Or is it better to use the finite difference Jacobian of a gradient?)
# TODO: Allow user to specify Val{:central}, Val{:forward}, :Val{:complex}
jcache = DiffEqDiffTools.JacobianCache(x_seed, Val{:central})
if is_finitediff(autodiff)

# Figure out which Val-type to use for DiffEqDiffTools based on our
# symbol interface.
fdtype = diffeqdiff_fdtype(autodiff)

jcache = DiffEqDiffTools.JacobianCache(x_seed, fdtype)
function h!(storage, x)
DiffEqDiffTools.finite_difference_jacobian!(storage, d.df, x, jcache)
return
end
elseif autodiff == :forward
elseif is_forwarddiff(autodiff)
hcfg = ForwardDiff.HessianConfig(d.f, similar(gradient(d)))
h! = (out, x) -> ForwardDiff.hessian!(out, d.f, x, hcfg)
else
Expand All @@ -96,9 +100,13 @@ end

function TwiceDifferentiable(f, x::AbstractVector, F::Real = real(zero(eltype(x)));
autodiff = :finite, inplace = true)
if autodiff == :finite
# TODO: Allow user to specify Val{:central}, Val{:forward}, Val{:complex}
gcache = DiffEqDiffTools.GradientCache(x, x, Val{:central})
if is_finitediff(autodiff)

# Figure out which Val-type to use for DiffEqDiffTools based on our
# symbol interface.
fdtype = diffeqdiff_fdtype(autodiff)
gcache = DiffEqDiffTools.GradientCache(x, x, fdtype)

function g!(storage, x)
DiffEqDiffTools.finite_difference_gradient!(storage, f, x, gcache)
return
Expand All @@ -107,13 +115,14 @@ function TwiceDifferentiable(f, x::AbstractVector, F::Real = real(zero(eltype(x)
g!(storage, x)
return f(x)
end
# TODO: Allow user to specify Val{:central}, Val{:forward}, :Val{:complex}

function h!(storage::Matrix, x::Vector)
# TODO: Wait to use DiffEqDiffTools until they introduce the Hessian feature
Calculus.finite_difference_hessian!(f, x, storage)
return
end
elseif autodiff == :forward
elseif is_forwarddiff(autodiff)

gcfg = ForwardDiff.GradientConfig(f, x)
g! = (out, x) -> ForwardDiff.gradient!(out, f, x, gcfg)

Expand Down

0 comments on commit 46fd262

Please sign in to comment.