From c3a6365a8dd4e04ed53498970512cc279f734ae4 Mon Sep 17 00:00:00 2001 From: brianguenter <1brianguenter@gmail.com> Date: Mon, 11 Mar 2024 16:48:16 -0700 Subject: [PATCH] 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 | 12 ++++----- test/FDTests.jl | 69 ++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 73 insertions(+), 8 deletions(-) diff --git a/src/Jacobian.jl b/src/Jacobian.jl index e693f0e9..6679cb43 100644 --- a/src/Jacobian.jl +++ b/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 diff --git a/test/FDTests.jl b/test/FDTests.jl index ff6de565..7d98ab87 100644 --- a/test/FDTests.jl +++ b/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