Skip to content

Commit

Permalink
Merge c1c9ce7 into fb14851
Browse files Browse the repository at this point in the history
  • Loading branch information
@huanglangwen
huanglangwen committed Oct 15, 2019
2 parents fb14851 + c1c9ce7 commit 4b56524
Show file tree
Hide file tree
Showing 4 changed files with 180 additions and 65 deletions.
159 changes: 136 additions & 23 deletions src/jacobians.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,35 +93,40 @@ function JacobianCache(
JacobianCache{typeof(_x1),typeof(_fx),typeof(fx1),typeof(colorvec),typeof(sparsity),fdtype,returntype,inplace}(_x1,_fx,fx1,colorvec,sparsity)
end

function finite_difference_jacobian!(J::AbstractMatrix,
f,
x::AbstractArray{<:Number},
fdtype :: Type{T1}=Val{:forward},
returntype :: Type{T2}=eltype(x),
inplace :: Type{Val{T3}}=Val{true},
f_in :: Union{T2,Nothing}=nothing;
relstep=default_relstep(fdtype, eltype(x)),
absstep=relstep,
colorvec = eachindex(x),
sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) where {T1,T2,T3}
function _make_Ji(rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
pick_inds = [i for i in 1:length(rows_index) if colorvec[cols_index[i]] == color_i]
rows_index_c = rows_index[pick_inds]
cols_index_c = cols_index[pick_inds]
len_rows = length(pick_inds)
unused_rows = setdiff(1:nrows,rows_index_c)
perm_rows = sortperm(vcat(rows_index_c,unused_rows))
cols_index_c = vcat(cols_index_c,zeros(Int,nrows-len_rows))[perm_rows]
Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:nrows, j in 1:ncols]
Ji
end

cache = JacobianCache(x, fdtype, returntype, inplace)
finite_difference_jacobian!(J, f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity)
function Base.vec(x::Number)
x
end

function finite_difference_jacobian(f, x::AbstractArray{<:Number},
fdtype :: Type{T1}=Val{:forward},
returntype :: Type{T2}=eltype(x),
inplace :: Type{Val{T3}}=Val{true},
f_in :: Union{T2,Nothing}=nothing;
relstep=default_relstep(fdtype, eltype(x)),
absstep=relstep,
colorvec = eachindex(x),
sparsity = nothing,
jac_prototype = nothing,
dir=true) where {T1,T2,T3}

cache = JacobianCache(x, fdtype, returntype, inplace)
finite_difference_jacobian(f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity, dir=dir)
if f_in isa Nothing
fx = f(x)
else
fx = f_in
end
cache = JacobianCache(x, fx, fdtype, returntype, Val{false})
finite_difference_jacobian(f, x, cache, fx; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity, jac_prototype=jac_prototype, dir=dir)
end

function finite_difference_jacobian(
Expand All @@ -133,11 +138,119 @@ function finite_difference_jacobian(
absstep=relstep,
colorvec = cache.colorvec,
sparsity = cache.sparsity,
jac_prototype = nothing,
dir=true) where {T1,T2,T3,cType,sType,fdtype,returntype,inplace}
_J = false .* x .* x'
_J isa SMatrix ? J = MArray(_J) : J = _J
finite_difference_jacobian!(J, f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity, dir=dir)
_J isa SMatrix ? SArray(J) : J

x1, fx, fx1 = cache.x1, cache.fx, cache.fx1

if !(f_in isa Nothing)
vecfx = vec(f_in)
elseif fdtype == Val{:forward}
vecfx = vec(f(x))
elseif fdtype == Val{:complex} && returntype <: Real
vecfx = real(fx)
else
vecfx = vec(fx)
end
vecx = vec(x)
vecx1 = vec(x1)
J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false.*vecfx.*x' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
nrows, ncols = size(J)

if !(sparsity isa Nothing)
rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
rows_index = [rows_index[i] for i in 1:length(rows_index)]
cols_index = [cols_index[i] for i in 1:length(cols_index)]
end

if fdtype == Val{:forward}
@inbounds for color_i ∈ 1:maximum(colorvec)
if sparsity isa Nothing
x_save = vecx[color_i]
epsilon = compute_epsilon(Val{:forward}, x_save, relstep, absstep, dir)
_vecx1 = Base.setindex(vecx,x_save+epsilon,color_i)
_x1 = reshape(_vecx1,size(x))
vecfx1 = vec(f(_x1))
dx = (vecfx1-vecfx)/epsilon
J = J + mapreduce(i -> i==color_i ? dx : zeros(eltype(x),nrows), hcat, 1:ncols)
else
tmp = norm(vecx .* (colorvec .== color_i))
epsilon = compute_epsilon(Val{:forward}, sqrt(tmp), relstep, absstep, dir)
_vecx = @. vecx + epsilon * (colorvec == color_i)
_x = reshape(_vecx,size(x))
vecfx1 = vec(f(_x))
dx = (vecfx1-vecfx)/epsilon
Ji = _make_Ji(rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
J = J + Ji
end
end
elseif fdtype == Val{:central}
@inbounds for color_i ∈ 1:maximum(colorvec)
if sparsity isa Nothing
x1_save = vecx1[color_i]
x_save = vecx[color_i]
epsilon = compute_epsilon(Val{:forward}, x1_save, relstep, absstep, dir)
_vecx1 = Base.setindex(vecx1,x1_save+epsilon,color_i)
_vecx = Base.setindex(vecx,x_save-epsilon,color_i)
_x1 = reshape(_vecx1,size(x))
_x = reshape(_vecx,size(x))
vecfx1 = vec(f(_x1))
vecfx = vec(f(_x))
dx = (vecfx1-vecfx)/(2epsilon)
J = J + mapreduce(i -> i==color_i ? dx : zeros(eltype(x),nrows), hcat, 1:ncols)
else
tmp = norm(vecx1 .* (colorvec .== color_i))
epsilon = compute_epsilon(Val{:forward}, sqrt(tmp), relstep, absstep, dir)
_vecx1 = @. vecx1 + epsilon * (colorvec == color_i)
_vecx = @. vecx - epsilon * (colorvec == color_i)
_x1 = reshape(_vecx1,size(x))
_x = reshape(_vecx, size(x))
vecfx1 = vec(f(_x1))
vecfx = vec(f(_x))
dx = (vecfx1-vecfx)/(2epsilon)
Ji = _make_Ji(rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
J = J + Ji
end
end
elseif fdtype == Val{:complex} && returntype <: Real
epsilon = eps(eltype(x))
@inbounds for color_i ∈ 1:maximum(colorvec)
if sparsity isa Nothing
x_save = vecx[color_i]
_vecx = Base.setindex(complex.(vecx),x_save+im*epsilon,color_i)
_x = reshape(_vecx,size(x))
vecfx = vec(f(_x))
dx = imag(vecfx)/epsilon
J = J + mapreduce(i -> i==color_i ? dx : zeros(eltype(x),nrows), hcat, 1:ncols)
else
_vecx = @. vecx + im * epsilon * (colorvec == color_i)
_x = reshape(_vecx,size(x))
vecfx = vec(f(_x))
dx = imag(vecfx)/epsilon
Ji = _make_Ji(rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
J = J + Ji
end
end
else
fdtype_error(returntype)
end
J
end

function finite_difference_jacobian!(J::AbstractMatrix,
f,
x::AbstractArray{<:Number},
fdtype :: Type{T1}=Val{:forward},
returntype :: Type{T2}=eltype(x),
inplace :: Type{Val{T3}}=Val{true},
f_in :: Union{T2,Nothing}=nothing;
relstep=default_relstep(fdtype, eltype(x)),
absstep=relstep,
colorvec = eachindex(x),
sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) where {T1,T2,T3}

cache = JacobianCache(x, fdtype, returntype, inplace)
finite_difference_jacobian!(J, f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity)
end

function finite_difference_jacobian!(
Expand Down Expand Up @@ -203,7 +316,7 @@ function finite_difference_jacobian!(
if inplace == Val{true}
@. x1 = x1 + epsilon * (_color == color_i)
else
_x1 = @. _x1 + epsilon * (_color == color_i)
_x1 = @. x1 + epsilon * (_color == color_i)
end
end

Expand Down Expand Up @@ -297,8 +410,8 @@ function finite_difference_jacobian!(
@. x1 = x1 + epsilon * (_color == color_i)
@. x = x - epsilon * (_color == color_i)
else
_x1 = @. _x1 + epsilon * (_color == color_i)
_x = @. _x - epsilon * (_color == color_i)
_x1 = @. x1 + epsilon * (_color == color_i)
_x = @. x - epsilon * (_color == color_i)
end
end

Expand Down
12 changes: 8 additions & 4 deletions test/coloring_tests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ function f(dx,x)
dx[end] = x[end-1] - 2x[end]
nothing
end
function oopf(x)
global fcalls += 1
vcat([-2x[1]+x[2]],[x[i-1]-2x[i]+x[i+1] for i in 2:length(x)-1],[x[end-1]-2x[end]])
end

function second_derivative_stencil(N)
A = zeros(N,N)
Expand All @@ -20,7 +24,7 @@ function second_derivative_stencil(N)
A
end

J = DiffEqDiffTools.finite_difference_jacobian(f,rand(30))
J = DiffEqDiffTools.finite_difference_jacobian(oopf,rand(30))
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
@test fcalls == 31
Expand Down Expand Up @@ -83,11 +87,11 @@ DiffEqDiffTools.finite_difference_jacobian!(_J2,f,rand(30),Val{:complex},colorve
@test _J2 ≈ _J

_Jb = BandedMatrices.BandedMatrix(similar(_J2),(1,1))
DiffEqDiffTools.finite_difference_jacobian!(_Jb, f, rand(30), colorvec=colorvec=repeat(1:3,10))
DiffEqDiffTools.finite_difference_jacobian!(_Jb, f, rand(30), colorvec=repeat(1:3,10))
@test _Jb ≈ _J

_Jtri = Tridiagonal(similar(_J2))
DiffEqDiffTools.finite_difference_jacobian!(_Jtri, f, rand(30), colorvec=colorvec=repeat(1:3,10))
DiffEqDiffTools.finite_difference_jacobian!(_Jtri, f, rand(30), colorvec=repeat(1:3,10))
@test _Jtri ≈ _J

#https://github.com/JuliaDiffEq/DiffEqDiffTools.jl/issues/67#issuecomment-516871956
Expand All @@ -111,4 +115,4 @@ DiffEqDiffTools.finite_difference_jacobian!(Jsparse, f, x, colorvec=colorsbbb)
Jbb = BlockBandedMatrix(similar(Jsparse),(fill(100, 100), fill(100, 100)),(1,1));
colorsbb = ArrayInterface.matrix_colors(Jbb)
DiffEqDiffTools.finite_difference_jacobian!(Jbb, f, x, colorvec=colorsbb)
@test Jbb ≈ Jsparse
@test Jbb ≈ Jsparse
55 changes: 32 additions & 23 deletions test/finitedifftests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -307,57 +307,66 @@ central_cache = DiffEqDiffTools.GradientCache(df,x,Val{:central})
end

# Jacobian tests
function f(fvec,x)
function iipf(fvec,x)
fvec[1] = (x[1]+3)*(x[2]^3-7)+18
fvec[2] = sin(x[2]*exp(x[1])-1)
end
function oopf(x)
[(x[1]+3)*(x[2]^3-7)+18,
sin(x[2]*exp(x[1])-1)]
end
x = rand(2); y = rand(2)
z = copy(x)
ff(df, x) = !all(x .<= z) ? error() : f(df, x)
f(y,x)
iipff(df, x) = !all(x .<= z) ? error() : iipf(df, x)
iipf(y,x)
oopff(x) = !all(x .<= z) ? error() : oopf(x)
J_ref = [[-7+x[2]^3 3*(3+x[1])*x[2]^2]; [exp(x[1])*x[2]*cos(1-exp(x[1])*x[2]) exp(x[1])*cos(1-exp(x[1])*x[2])]]
J = zero(J_ref)
df = zero(x)
df_ref = diag(J_ref)
epsilon = zero(x)
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x,Val{:central})
complex_cache = DiffEqDiffTools.JacobianCache(x,Val{:complex})
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward},eltype(x),Val{false})
central_cache = DiffEqDiffTools.JacobianCache(x,Val{:central},eltype(x),Val{false})
complex_cache = DiffEqDiffTools.JacobianCache(x,Val{:complex},eltype(x),Val{false})
f_in = copy(y)

@time @testset "Jacobian StridedArray real-valued tests" begin
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(ff, x, forward_cache, dir=-1), J_ref) < 1e-4
@test_throws Any err_func(DiffEqDiffTools.finite_difference_jacobian(ff, x, forward_cache), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache, relstep=sqrt(eps())), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache, f_in), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, central_cache), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, Val{:central}), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, complex_cache), J_ref) < 1e-14
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopff, x, forward_cache, dir=-1), J_ref) < 1e-4
@test_throws Any err_func(DiffEqDiffTools.finite_difference_jacobian(oopff, x, forward_cache), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache, relstep=sqrt(eps())), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache, f_in), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, central_cache), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, Val{:central}), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, complex_cache), J_ref) < 1e-14
end

function f(fvec,x)
function iipf(fvec,x)
fvec[1] = (im*x[1]+3)*(x[2]^3-7)+18
fvec[2] = sin(x[2]*exp(x[1])-1)
end
function oopf(x)
[(im*x[1]+3)*(x[2]^3-7)+18,
sin(x[2]*exp(x[1])-1)]
end
x = rand(2) + im*rand(2)
y = similar(x)
f(y,x)
iipf(y,x)
J_ref = [[im*(-7+x[2]^3) 3*(3+im*x[1])*x[2]^2]; [exp(x[1])*x[2]*cos(1-exp(x[1])*x[2]) exp(x[1])*cos(1-exp(x[1])*x[2])]]
J = zero(J_ref)
df = zero(x)
df_ref = diag(J_ref)
epsilon = zero(real.(x))
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward})
central_cache = DiffEqDiffTools.JacobianCache(x,Val{:central})
forward_cache = DiffEqDiffTools.JacobianCache(x,Val{:forward},eltype(x),Val{false})
central_cache = DiffEqDiffTools.JacobianCache(x,Val{:central},eltype(x),Val{false})
f_in = copy(y)

@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
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache, relstep=sqrt(eps())), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, forward_cache, f_in), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, central_cache), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(f, x, Val{:central}), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache, relstep=sqrt(eps())), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, forward_cache, f_in), J_ref) < 1e-4
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, central_cache), J_ref) < 1e-8
@test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, Val{:central}), J_ref) < 1e-8
end

# Hessian tests
Expand Down
19 changes: 4 additions & 15 deletions test/out_of_place_tests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,23 +16,12 @@ function second_derivative_stencil(N)
end

x = @SVector ones(30)
J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:forward}, eltype(x), Val{false})
J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:forward}, eltype(x))
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
DiffEqDiffTools.finite_difference_jacobian!(_J, f,x, Val{:forward}, eltype(x), Val{false},
colorvec=repeat(1:3,10))
@test _J ≈ second_derivative_stencil(30)

J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:central}, eltype(x), Val{false})
J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:central}, eltype(x))
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
DiffEqDiffTools.finite_difference_jacobian!(_J, f,x, Val{:central}, eltype(x), Val{false},
colorvec=repeat(1:3,10))
@test _J ≈ second_derivative_stencil(30)

J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:complex}, eltype(x), Val{false})
J = DiffEqDiffTools.finite_difference_jacobian(f,x, Val{:complex}, eltype(x))
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
DiffEqDiffTools.finite_difference_jacobian!(_J, f,x, Val{:complex}, eltype(x), Val{false},
colorvec=repeat(1:3,10))
@test _J ≈ second_derivative_stencil(30)

0 comments on commit 4b56524

Please sign in to comment.