Merge pull request #26 from Roger-luo/roger/multiargs

add multi arg jacobian

Author: willtebbutt

src/FiniteDifferences.jl

@@ -7,4 +7,7 @@ module FiniteDifferences
     include("methods.jl")
     include("numerics.jl")
     include("grad.jl")
+
+
+    @deprecate jacobian(fdm, f, x::Vector, D::Int) jacobian(fdm, f, x; len=D)
 end

src/grad.jl

@@ -1,40 +1,75 @@
 export grad, jacobian, jvp, j′vp, to_vec
+replace_arg(x, xs::Tuple, k::Int) = ntuple(p -> p == k ? x : xs[p], length(xs))
 
 """
-    grad(fdm, f, x::Vector{<:Number})
+    grad(fdm, f, xs...)
 
-Approximate the gradient of `f` at `x` using `fdm`. Assumes that `f(x)` is scalar.
+Approximate the gradient of `f` at `xs...` using `fdm`. Assumes that `f(xs...)` is scalar.
 """
-function grad(fdm, f, x::Vector{T}) where T<:Number
-    v, dx, tmp = fill(zero(T), size(x)), similar(x), similar(x)
-    for n in eachindex(x)
-        v[n] = one(T)
-        dx[n] = fdm(function(ϵ)
-                tmp .= x .+ ϵ .* v
-                return f(tmp)
-            end,
-            zero(T),
-        )
-        v[n] = zero(T)
+function grad end
+
+function grad(fdm, f, x::AbstractArray{T}) where T <: Number
+    dx = similar(x)
+    tmp = similar(x)
+    for k in eachindex(x)
+        dx[k] = fdm(zero(T)) do ϵ
+            tmp .= x
+            tmp[k] += ϵ
+            return f(tmp)
+        end
     end
     return dx
 end
 
+grad(fdm, f, x::Real) = fdm(f, x)
+grad(fdm, f, x::Tuple) = grad(fdm, (xs...)->f(xs), x...)
+
+function grad(fdm, f, d::Dict{K, V}) where {K, V}
+    dd = Dict{K, V}()
+    for (k, v) in d
+        function f′(x)
+            tmp = copy(d)
+            tmp[k] = x
+            return f(tmp)
+        end
+        dd[k] = grad(fdm, f′, v)
+    end
+    return dd
+end
+
+function grad(fdm, f, x)
+    v, back = to_vec(x)
+    return back(grad(fdm, x->f(back(v)), v))
+end
+
+function grad(fdm, f, xs...)
+    return ntuple(length(xs)) do k
+        grad(fdm, x->f(replace_arg(x, xs, k)...), xs[k])
+    end
+end
+
 """
-    jacobian(fdm, f, x::Vector{<:Number}, D::Int)
-    jacobian(fdm, f, x::Vector{<:Number})
+    jacobian(fdm, f, xs::Union{Real, AbstractArray{<:Real}}; len::Int=length(f(x)))
 
 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::Vector{T}, D::Int) where {T<:Number}
-    J = Matrix{T}(undef, D, length(x))
-    for d in 1:D
-        J[d, :] = grad(fdm, x->f(x)[d], x)
+function jacobian(fdm, f, x::Union{T, AbstractArray{T}}; len::Int=length(f(x))) where {T <: Number}
+    J = Matrix{float(T)}(undef, len, length(x))
+    for d in 1:len
+        gs = grad(fdm, x->f(x)[d], x)
+        for k in 1:length(x)
+            J[d, k] = gs[k]
+        end
     end
     return J
 end
-jacobian(fdm, f, x::Vector{<:Number}) = jacobian(fdm, f, x, length(f(x)))
+
+function jacobian(fdm, f, xs...; len::Int=length(f(xs...)))
+    return ntuple(length(xs)) do k
+        jacobian(fdm, x->f(replace_arg(x, xs, k)...), xs[k]; len=len)
+    end
+end
 
 """
     _jvp(fdm, f, x::Vector{<:Number}, ẋ::AbstractVector{<:Number})
@@ -48,7 +83,7 @@ _jvp(fdm, f, x::Vector{<:Number}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* x
 
 Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
 """
-_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x, length(ȳ))) * ȳ
+_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))) * ȳ
 
 """
     jvp(fdm, f, x, ẋ)

test/grad.jl

@@ -26,8 +26,8 @@ Base.length(x::DummyType) = size(x.X, 1)
 
     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 jacobian(fdm, f, x; len=length(ȳ)) ≈ J_exact
+        @test jacobian(fdm, f, x) == jacobian(fdm, f, x; len=length(ȳ))
         @test _jvp(fdm, f, x, ẋ) ≈ J_exact * ẋ
         @test _j′vp(fdm, f, ȳ, x) ≈ transpose(J_exact) * ȳ
         @test xc == x
@@ -44,6 +44,62 @@ Base.length(x::DummyType) = size(x.X, 1)
         @test Ac == A
     end
 
+    @testset "multi vars jacobian/grad" begin
+        rng, fdm = MersenneTwister(123456), central_fdm(5, 1)
+        
+        f1(x, y) = x * y + x
+        f2(x, y) = sum(x * y + x)
+        f3(x::Tuple) = sum(x[1]) + x[2]
+        f4(d::Dict) = sum(d[:x]) + d[:y]
+
+        @testset "jacobian" begin
+            @testset "check multiple matrices" begin
+                x, y = rand(rng, 3, 3), rand(rng, 3, 3)
+                jac_xs = jacobian(fdm, f1, x, y)
+                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)
+                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)
+            end
+
+            @testset "check mixed scalar and matrices" begin
+                x, y = rand(3, 3), 2
+                jac_xs = jacobian(fdm, f1, x, y)
+                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)
+                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)
+            end
+        end
+
+        @testset "grad" begin
+            @testset "check multiple matrices" begin
+                x, y = rand(rng, 3, 3), rand(rng, 3, 3)
+                dxs = grad(fdm, f2, x, y)
+                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)
+                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)
+            end
+
+            @testset "check mixed scalar & matrices" begin
+                x, y = rand(rng, 3, 3), 2
+                dxs = grad(fdm, f2, x, y)
+                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)
+                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)        
+            end
+
+            @testset "check tuple" begin
+                x, y = rand(rng, 3, 3), 2
+                dxs = grad(fdm, f3, (x, y))
+                @test dxs[1] ≈ grad(fdm, x->f3((x, y)), x)
+                @test dxs[2] ≈ grad(fdm, y->f3((x, y)), y)    
+            end
+
+            @testset "check dict" begin
+                x, y = rand(rng, 3, 3), 2
+                d = Dict(:x=>x, :y=>y)
+                dxs = grad(fdm, f4, d)
+                @test dxs[:x] ≈ grad(fdm, x->f3((x, y)), x)
+                @test dxs[:y] ≈ grad(fdm, y->f3((x, y)), y)
+            end
+        end
+    end
+
     function test_to_vec(x)
         x_vec, back = to_vec(x)
         @test x_vec isa Vector