Wrong JVP and VJP ?

Fixes #60

These two functions used the wrong dimensions to determine the size of the result vector. This should fix this problem.

Wrote tests to verify the the functions  jacobian_transpose_v numerically matches transpose(J)*v and jacobian_times_v numerically matches J*v.

src/Jacobian.jl

@@ -145,7 +145,7 @@ export sparse_jacobian
 Returns a vector of Node, where each element in the vector is the symbolic form of `Jv`. Also returns `v_vector` a vector of the `v` variables. This is useful if you want to generate a function to evaluate `Jv` and you want to separate the inputs to the function and the `v` variables."""
 function jacobian_times_v(terms::AbstractVector{T}, partial_variables::AbstractVector{S}) where {T<:Node,S<:Node}
     graph = DerivativeGraph(terms)
-    v_vector = make_variables(gensym(), domain_dimension(graph))
+    v_vector = make_variables(gensym(), length(partial_variables))
     factor!(graph)
     for (variable, one_v) in zip(partial_variables, v_vector)
         new_edges = PathEdge[]
@@ -245,9 +245,7 @@ function jacobian_transpose_v(terms::AbstractVector{T}, partial_variables::Abstr
         end
     end
 
-    outdim = domain_dimension(graph)
-
-    result = Vector{Node}(undef, outdim)
+    result = Vector{Node}(undef, length(partial_variables))
 
     for i in eachindex(result)
         result[i] = Node(0.0)
@@ -259,14 +257,14 @@ function jacobian_transpose_v(terms::AbstractVector{T}, partial_variables::Abstr
         for root_index in 1:codomain_dimension(graph)
             var_index = variable_node_to_index(graph, var)
             if var_index !== nothing
-                result[var_index] += evaluate_path(graph, root_index, var_index)
+                result[i] += evaluate_path(graph, root_index, var_index)
             else
-                result[var_index] = zero(Node) #TODO fix this so get more generic zero value
+                result[i] = zero(Node) #TODO fix this so get more generic zero value
             end
         end
     end
 
-    return result, r_vector #need v_vector values if want to make executable after making symbolic form. Need to differentiate between variables that were in original graph and variables introduced by v_vector
+    return result, r_vector #need r_vector values if want to make executable after making symbolic form. Need to differentiate between variables that were in original graph and variables introduced by r_vector
 end
 export jacobian_transpose_v
 

test/FDTests.jl

@@ -1399,8 +1399,9 @@ end
     @test DA == derivative(A, nq1)
 end
 
-@testitem "FD.jacobian_times_v" begin
+@testitem "jacobian_times_v" begin
     import FastDifferentiation as FD
+    using Random
     include("ShareTestCode.jl")
 
     order = 10
@@ -1427,10 +1428,43 @@ end
         @test isapprox(slow_val, fast_val, rtol=1e-9)
     end
 
+    #test for https://github.com/brianguenter/FastDifferentiation.jl/issues/60
+    function heat(u, p, t)
+        gamma, dx = p
+        dxxu = (circshift(u, 1) .- 2 * u .+ circshift(u, -1)) / dx^2
+        return gamma * dxxu
+    end
+
+
+    Ns = 5
+    xmin, xmax = 0.0, 1.0
+    dx = (xmax - xmin) / (Ns - 1)
+
+    xs = collect(LinRange(xmin, xmax, Ns))
+    using FastDifferentiation
+    FastDifferentiation.@variables uv pv tv
+    uvs = make_variables(:uv, Ns)
+    pvs = make_variables(:pv, 2)
+    tvs = make_variables(:tv, 1)
+    fv = heat(uvs, pvs, tvs)
+    js = FastDifferentiation.jacobian(fv, uvs) # The Jacobian js is of shape Ns x Ns
+
+    vjps, rvs = FastDifferentiation.jacobian_times_v(fv, uvs) # The variable vjps is of shape Ns+2 and rvs is of shape Ns
+
+    input_vec = [uvs; pvs; tvs; rvs]
+    jv = make_function(FastDifferentiation.Node.(js * rvs), input_vec)
+    jtv2 = make_function(vjps, input_vec)
+
+    for i in 1:100
+        rng = Random.Xoshiro(123)
+        tvec = rand(rng, length(input_vec))
+        @test isapprox(jv(tvec), jtv2(tvec))
+    end
 end
 
 @testitem "jacobian_transpose_v" begin
     import FastDifferentiation as FD
+    using Random
     include("ShareTestCode.jl")
 
     order = 10
@@ -1455,6 +1489,39 @@ end
 
         @test isapprox(slow_val, fast_val, rtol=1e-8)
     end
+
+    #test for https://github.com/brianguenter/FastDifferentiation.jl/issues/60
+    function heat(u, p, t)
+        gamma, dx = p
+        dxxu = (circshift(u, 1) .- 2 * u .+ circshift(u, -1)) / dx^2
+        return gamma * dxxu
+    end
+
+
+    Ns = 5
+    xmin, xmax = 0.0, 1.0
+    dx = (xmax - xmin) / (Ns - 1)
+
+    xs = collect(LinRange(xmin, xmax, Ns))
+    using FastDifferentiation
+    FastDifferentiation.@variables uv pv tv
+    uvs = make_variables(:uv, Ns)
+    pvs = make_variables(:pv, 2)
+    tvs = make_variables(:tv, 1)
+    fv = heat(uvs, pvs, tvs)
+    js = FastDifferentiation.jacobian(fv, uvs) # The Jacobian js is of shape Ns x Ns
+
+    jvps, vvs = FastDifferentiation.jacobian_transpose_v(fv, uvs) # The variable jvps is of shape Ns and rvs is of shape Ns+2
+
+    input_vec = [uvs; pvs; tvs; vvs]
+    jtv = make_function(FastDifferentiation.Node.(js'vvs), input_vec)
+    jtv2 = make_function(jvps, input_vec)
+
+    for i in 1:100
+        rng = Random.Xoshiro(123)
+        tvec = rand(rng, length(input_vec))
+        @test isapprox(jtv(tvec), jtv2(tvec))
+    end
 end
 
 @testitem "hessian" begin