Merge branch 'master' of https://github.com/invenia/FDM.jl

src/FDM.jl

@@ -2,6 +2,8 @@ module FDM
 
     using Printf, LinearAlgebra
 
+    const AV = AbstractVector
+
     include("methods.jl")
     include("numerics.jl")
     include("grad.jl")

src/grad.jl

@@ -3,7 +3,7 @@
 
 Approximate the gradient of `f` at `x` using `fdm`. Assumes that `f(x)` is scalar.
 """
-function grad(fdm, f, x::AbstractArray{T}) where T
+function grad(fdm, f, x::AbstractArray{T}) where T<:Real
     v, dx, tmp = fill(zero(T), size(x)), similar(x), similar(x)
     for n in eachindex(x)
         v[n] = one(T)
@@ -17,3 +17,33 @@ function grad(fdm, f, x::AbstractArray{T}) where T
     end
     return dx
 end
+
+"""
+    jacobian(fdm, f, x::AbstractVector{<:Real}, D::Int)
+    jacobian(fdm, f, x::AbstractVector{<:Real})
+
+Approximate the Jacobian of `f` at `x` using `fdm`. `f(x)` must be a length `D` vector. If
+`D` is not provided, then `f(x)` is computed once to determine the output size.
+"""
+function jacobian(fdm, f, x::AV{T}, D::Int) where {T<:Real}
+    J = Matrix{T}(undef, D, length(x))
+    for d in 1:D
+        J[d, :] = grad(fdm, x->f(x)[d], x)
+    end
+    return J
+end
+jacobian(fdm, f, x::AV{<:Real}) = jacobian(fdm, f, x, length(f(x)))
+
+"""
+    jvp(fdm, f, x::AbstractVector{<:Real}, ẋ::AbstractVector{<:Real})
+
+Convenience function to compute `jacobian(f, x) * ẋ`.
+"""
+jvp(fdm, f, x::AV{<:Real}, ẋ::AV{<:Real}) = jacobian(fdm, f, x) * ẋ
+
+"""
+    j′vp(fdm, f, ȳ::AbstractVector{<:Real}, x::AbstractVector{<:Real})
+
+Convenience function to compute `jacobian(f, x)' * ȳ`.
+"""
+j′vp(fdm, f, ȳ::AV{<:Real}, x::AV{<:Real}) = jacobian(fdm, f, x, length(ȳ))' * ȳ

test/grad.jl

@@ -1,8 +1,28 @@
-using FDM: grad
+using FDM: grad, jacobian, jvp, j′vp
 
 @testset "grad" begin
     x = randn(MersenneTwister(123456), 2)
     xc = copy(x)
     @test grad(central_fdm(5, 1), x->sin(x[1]) + cos(x[2]), x) ≈ [cos(x[1]), -sin(x[2])]
     @test xc == x
 end
+
+function check_jac_and_jvp_and_j′vp(fdm, f, ȳ, x, ẋ, J_exact)
+    xc = copy(x)
+    @test jacobian(fdm, f, x, length(ȳ)) ≈ J_exact
+    @test jacobian(fdm, f, x) == jacobian(fdm, f, x, length(ȳ))
+    @test jvp(fdm, f, x, ẋ) ≈ J_exact * ẋ
+    @test j′vp(fdm, f, ȳ, x) ≈ J_exact' * ȳ
+    @test xc == x
+end
+
+@testset "jacobian / jvp / j′vp" begin
+    rng, P, Q, fdm = MersenneTwister(123456), 3, 2, central_fdm(5, 1)
+    ȳ, A, x, ẋ = randn(rng, P), randn(rng, P, Q), randn(rng, Q), randn(rng, Q)
+    Ac = copy(A)
+
+    check_jac_and_jvp_and_j′vp(fdm, x->A * x, ȳ, x, ẋ, A)
+    @test Ac == A
+    check_jac_and_jvp_and_j′vp(fdm, x->sin.(A * x), ȳ, x, ẋ, cos.(A * x) .* A)
+    @test Ac == A
+end

test/runtests.jl

@@ -1,4 +1,4 @@
-using FDM, Test, Random, Printf
+using FDM, Test, Random, Printf, LinearAlgebra
 
 @testset "FDM" begin
     include("methods.jl")