Skip to content

Commit

Permalink
allocating dispatches for the Jacobian
Browse files Browse the repository at this point in the history
  • Loading branch information
@ChrisRackauckas
ChrisRackauckas committed Feb 15, 2018
1 parent 00a0cf5 commit 44892a2
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 42 deletions.
112 changes: 75 additions & 37 deletions src/jacobians.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,56 +5,91 @@ struct JacobianCache{CacheType1,CacheType2,CacheType3,fdtype,returntype,inplace}
end

function JacobianCache(
x :: AbstractArray{<:Number},
x1 :: Union{Void,AbstractArray{<:Number}} = nothing,
fx :: Union{Void,AbstractArray{<:Number}} = nothing,
fx1 :: Union{Void,AbstractArray{<:Number}} = nothing,
x,
fdtype :: Type{T1} = Val{:central},
returntype :: Type{T2} = eltype(x),
inplace :: Type{Val{T3}} = Val{true}) where {T1,T2,T3}
if eltype(x) <: Real && fdtype==Val{:complex}
x1 = zeros(Complex{eltype(x)}, size(x))
_fx = zeros(Complex{eltype(x)}, size(x))
else
x1 = similar(x)
_fx = similar(x)
end

if fdtype==Val{:complex}
if !(returntype<:Real)
fdtype_error(returntype)
end
if eltype(fx)!=Complex{eltype(x)}
_fx = zeros(Complex{eltype(x)}, size(x))
end
if typeof(fx1)!=Void
warn("fdtype==Val{:complex} doesn't benefit from caching fx1.")
end
_fx1 = nothing
if eltype(x1) != Complex{eltype(x)}
_x1 = zeros(Complex{eltype(x)}, size(x))
else
_x1 = x1
end
_fx1 = nothing
else
if eltype(x1) != eltype(x)
_x1 = similar(x)
else
_x1 = x1
end
if eltype(fx) != returntype
_fx = zeros(returntype, size(x))
_fx1 = similar(x)
end

JacobianCache(x1,_fx,_fx1,fdtype,returntype,inplace)
end

function JacobianCache(
x ,
fx,
fdtype :: Type{T1} = Val{:central},
returntype :: Type{T2} = eltype(x),
inplace :: Type{Val{T3}} = Val{true}) where {T1,T2,T3}

if eltype(x) <: Real && fdtype==Val{:complex}
x1 = zeros(Complex{eltype(x)}, size(x))
else
x1 = similar(x)
end

if eltype(fx) <: Real && fdtype==Val{:complex}
_fx = zeros(Complex{eltype(x)}, size(fx))
else
_fx = similar(fx)
end

if fdtype==Val{:complex}
_fx1 = nothing
else
_fx1 = similar(fx)
end

JacobianCache(x1,_fx,_fx1,fdtype,returntype,inplace)
end

function JacobianCache(
x1 ,
fx ,
fx1,
fdtype :: Type{T1} = Val{:central},
returntype :: Type{T2} = eltype(x),
inplace :: Type{Val{T3}} = Val{true}) where {T1,T2,T3}

if fdtype==Val{:complex}
!(returntype<:Real) && fdtype_error(returntype)

if eltype(fx) <: Real
_fx = zeros(Complex{eltype(x)}, size(fx))
else
_fx = fx
end
if eltype(fx1) != returntype
_fx1 = zeros(returntype, size(x))
if eltype(x1) <: Real
_x1 = zeros(Complex{eltype(x)}, size(x))
else
_fx1 = fx1
_x1 = x1
end
else
_x1 = x1
@assert eltype(fx) == T2
@assert eltype(fx1) == T2
_fx = fx
end
JacobianCache{typeof(_x1),typeof(_fx),typeof(_fx1),fdtype,returntype,inplace}(_x1,_fx,_fx1)
JacobianCache{typeof(_x1),typeof(_fx),typeof(fx1),fdtype,returntype,inplace}(_x1,_fx,fx1)
end

function finite_difference_jacobian(f, x::AbstractArray{<:Number},
fdtype :: Type{T1}=Val{:central},
returntype :: Type{T2}=eltype(x),
inplace :: Type{Val{T3}}=Val{true}) where {T1,T2,T3}

cache = JacobianCache(x,nothing,nothing,nothing,fdtype,returntype,inplace)
cache = JacobianCache(x,fdtype,returntype,inplace)
finite_difference_jacobian(f,x,cache)
end

Expand All @@ -81,11 +116,12 @@ function finite_difference_jacobian!(J::AbstractMatrix{<:Number}, f,x::AbstractA
if inplace == Val{true}
f(fx1, x1)
f(fx, x)
J[:,i] = (vfx1 - vfx) / epsilon
else
fx1 .= f(x1)
fx .= f(x)
J[:,i] = (vfx1 - vfx) / epsilon
end
@. J[:,i] = (vfx1 - vfx) / epsilon
x1[i] = x1_save
end
elseif fdtype == Val{:central}
Expand All @@ -100,11 +136,12 @@ function finite_difference_jacobian!(J::AbstractMatrix{<:Number}, f,x::AbstractA
if inplace == Val{true}
f(fx1, x1)
f(fx, x)
@. J[:,i] = (vfx1 - vfx) / (2*epsilon)
else
fx1 .= f(x1)
fx .= f(x)
fx1 = f(x1)
fx = f(x)
J[:,i] = (vfx1 - vfx) / (2*epsilon)
end
@. J[:,i] = (vfx1 - vfx) / (2*epsilon)
x1[i] = x1_save
x[i] = x_save
end
Expand All @@ -115,10 +152,11 @@ function finite_difference_jacobian!(J::AbstractMatrix{<:Number}, f,x::AbstractA
x1[i] += im * epsilon
if inplace == Val{true}
f(fx,x1)
@. J[:,i] = imag(vfx) / epsilon
else
fx .= f(x1)
fx = f(x1)
J[:,i] = imag(vfx) / epsilon
end
@. J[:,i] = imag(vfx) / epsilon
x1[i] = x1_save
end
else
Expand Down
10 changes: 5 additions & 5 deletions test/finitedifftests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -294,9 +294,9 @@ J = zeros(J_ref)
df = zeros(x)
df_ref = diag(J_ref)
epsilon = zeros(x)
forward_cache = DiffEqDiffTools.JacobianCache(x,similar(x),similar(x),similar(x),Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x,similar(x),similar(x),similar(x))
complex_cache = DiffEqDiffTools.JacobianCache(x,nothing,nothing,nothing,Val{:complex})
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x)
complex_cache = DiffEqDiffTools.JacobianCache(x,Val{:complex})

@time @testset "Jacobian StridedArray real-valued tests" begin
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache), J_ref) < 1e-4
Expand All @@ -317,8 +317,8 @@ J = zeros(J_ref)
df = zeros(x)
df_ref = diag(J_ref)
epsilon = zeros(real.(x))
forward_cache = DiffEqDiffTools.JacobianCache(x,similar(x),similar(x),similar(x),Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x,similar(x),similar(x),similar(x))
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x)

@time @testset "Jacobian StridedArray f : C^N -> C^N tests" begin
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache), J_ref) < 1e-4
Expand Down

0 comments on commit 44892a2

Please sign in to comment.