Refactor jacobian and grad (#66)

* Implements Jacobian via jvp

* Improves perf of jacobian

* Refactors grad

* Refactors jacobian

* Tidies a few things up

* Update src/grad.jl

Co-Authored-By: Lyndon White <[email protected]>

* Changes composition for application.

Co-Authored-By: Lyndon White <[email protected]>

* Clarifies jacobian algorithm

Co-Authored-By: Lyndon White <[email protected]>

* Clarifies grad function

Co-Authored-By: Lyndon White <[email protected]>

* Fixes to_vec docstring.

Co-Authored-By: Lyndon White <[email protected]>

* Update src/to_vec.jl

Co-Authored-By: Lyndon White <[email protected]>

* Moves exports to main file

* Fixes jpvp docstring

* Tidies up to_vec(::Vector)

* Tidies up to_vec(::Dict)

* Tidies up to_vec(::Dict) further

* Tidies up to_vec(::Tuple) further

* Deprecates len parameter

* Tweaks jacobian comment

* Tidies up docstring

* Tweaks depwarn

* Tweaks jacobian deprecation

Co-Authored-By: Nick Robinson <[email protected]>

* Fixes typo in docs

Co-Authored-By: Wessel <[email protected]>

Co-authored-by: Lyndon White <[email protected]>
Co-authored-by: Nick Robinson <[email protected]>
Co-authored-by: Wessel <[email protected]>

Author: 4 people

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDifferences"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.9.2"
+version = "0.9.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/FiniteDifferences.jl

@@ -2,12 +2,10 @@ module FiniteDifferences
 
     using Printf, LinearAlgebra
 
-    const AV = AbstractVector
+    export to_vec, grad, jacobian, jvp, j′vp
 
     include("methods.jl")
     include("numerics.jl")
+    include("to_vec.jl")
     include("grad.jl")
-
-
-    @deprecate jacobian(fdm, f, x::Vector, D::Int) jacobian(fdm, f, x; len=D)
 end

src/grad.jl

@@ -1,96 +1,50 @@
-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, xs...)
+    jacobian(fdm, f, x...)
 
-Approximate the gradient of `f` at `xs...` using `fdm`. Assumes that `f(xs...)` is scalar.
+Approximate the Jacobian of `f` at `x` using `fdm`. Results will be returned as a
+`Matrix{<:Number}` of `size(length(y_vec), length(x_vec))` where `x_vec` is the flattened
+version of `x`, and `y_vec` the flattened version of `f(x...)`. Flattening performed by
+[`to_vec`](@ref).
 """
-function grad end
-
-function _grad(fdm, f, x::AbstractArray{T}) where T <: Number
-    # x must be mutable, we will mutate it and then mutate it back.
-    dx = similar(x)
-    for k in eachindex(x)
-        dx[k] = fdm(zero(T)) do ϵ
-            xk = x[k]
-            x[k] = xk +  ϵ
-            ret = f(x)
-            x[k] = xk  # Can't do `x[k] -= ϵ` as floating-point math is not associative
+function jacobian(fdm, f, x::Vector{<:Number}; len=nothing)
+    len !== nothing && Base.depwarn(
+        "`len` keyword argument to `jacobian` is no longer required " *
+        "and will not be permitted in the future.",
+         :jacobian
+    )
+    ẏs = map(eachindex(x)) do n
+        return fdm(zero(eltype(x))) do ε
+            xn = x[n]
+            x[n] = xn + ε
+            ret = first(to_vec(f(x)))
+            x[n] = xn  # Can't do `x[n] -= ϵ` as floating-point math is not associative
             return ret
         end
     end
-    return (dx, )
+    return (hcat(ẏs...), )
 end
 
-grad(fdm, f, x::Array{<:Number}) = _grad(fdm, f, x)
-# Fallback for when we don't know `x` will be mutable:
-grad(fdm, f, x::AbstractArray{<:Number}) = _grad(fdm, f, similar(x).=x)
-
-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}
-    ∇d = Dict{K, V}()
-    for (k, v) in d
-        dk = d[k]
-        function f′(x)
-            d[k] = x
-            return f(d)
-        end
-        ∇d[k] = grad(fdm, f′, v)[1]
-        d[k] = dk
-    end
-    return (∇d, )
+function jacobian(fdm, f, x; len=nothing)
+    x_vec, from_vec = to_vec(x)
+    return jacobian(fdm, f ∘ from_vec, x_vec; len=len)
 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])[1]
-    end
-end
-
-"""
-    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 `len` vector. If
-`len` is not provided, then `f(x)` is computed once to determine the output size.
-"""
-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)[1]
-        for k in 1:length(x)
-            J[d, k] = gs[k]
-        end
-    end
-    return (J, )
-end
-
-function jacobian(fdm, f, xs...; len::Int=length(f(xs...)))
+function jacobian(fdm, f, xs...; len=nothing)
     return ntuple(length(xs)) do k
         jacobian(fdm, x->f(replace_arg(x, xs, k)...), xs[k]; len=len)[1]
     end
 end
 
+replace_arg(x, xs::Tuple, k::Int) = ntuple(p -> p == k ? x : xs[p], length(xs))
+
 """
     _jvp(fdm, f, x::Vector{<:Number}, ẋ::AbstractVector{<:Number})
 
 Convenience function to compute `jacobian(f, x) * ẋ`.
 """
-_jvp(fdm, f, x::Vector{<:Number}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* ẋ), zero(eltype(x)))
-
-"""
-    _j′vp(fdm, f, ȳ::AbstractVector{<:Number}, x::Vector{<:Number})
-
-Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
-"""
-_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))[1]) * ȳ
+function _jvp(fdm, f, x::Vector{<:Number}, ẋ::Vector{<:Number})
+    return fdm(ε -> f(x .+ ε .* ẋ), zero(eltype(x)))
+end
 
 """
     jvp(fdm, f, x, ẋ)
@@ -111,110 +65,23 @@ end
 """
     j′vp(fdm, f, ȳ, x...)
 
-Compute an adjoint with any types of arguments for which [`to_vec`](@ref) is defined.
+Compute an adjoint with any types of arguments `x` for which [`to_vec`](@ref) is defined.
 """
 function j′vp(fdm, f, ȳ, x)
     x_vec, vec_to_x = to_vec(x)
     ȳ_vec, _ = to_vec(ȳ)
-    return (vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec)), )
+    return (vec_to_x(_j′vp(fdm, first ∘ to_vec ∘ f ∘ vec_to_x, ȳ_vec, x_vec)), )
 end
-j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]
 
-"""
-    to_vec(x) -> Tuple{<:AbstractVector, <:Function}
-
-Transform `x` into a `Vector`, and return a closure which inverts the transformation.
-"""
-function to_vec(x::Number)
-    function Number_from_vec(x_vec)
-        return first(x_vec)
-    end
-    return [x], Number_from_vec
-end
-
-# Vectors
-to_vec(x::Vector{<:Number}) = (x, identity)
-function to_vec(x::Vector)
-    x_vecs_and_backs = map(to_vec, x)
-    x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
-    function Vector_from_vec(x_vec)
-        sz = cumsum([map(length, x_vecs)...])
-        return [backs[n](x_vec[sz[n]-length(x_vecs[n])+1:sz[n]]) for n in eachindex(x)]
-    end
-    return vcat(x_vecs...), Vector_from_vec
-end
-
-# Arrays
-function to_vec(x::Array{<:Number})
-    function Array_from_vec(x_vec)
-        return reshape(x_vec, size(x))
-    end
-    return vec(x), Array_from_vec
-end
-
-function to_vec(x::Array)
-    x_vec, back = to_vec(reshape(x, :))
-    function Array_from_vec(x_vec)
-        return reshape(back(x_vec), size(x))
-    end
-    return x_vec, Array_from_vec
-end
-
-# AbstractArrays
-function to_vec(x::T) where {T<:LinearAlgebra.AbstractTriangular}
-    x_vec, back = to_vec(Matrix(x))
-    function AbstractTriangular_from_vec(x_vec)
-        return T(reshape(back(x_vec), size(x)))
-    end
-    return x_vec, AbstractTriangular_from_vec
-end
-
-function to_vec(x::Symmetric)
-    function Symmetric_from_vec(x_vec)
-        return Symmetric(reshape(x_vec, size(x)))
-    end
-    return vec(Matrix(x)), Symmetric_from_vec
-end
-
-function to_vec(X::Diagonal)
-    function Diagonal_from_vec(x_vec)
-        return Diagonal(reshape(x_vec, size(X)...))
-    end
-    return vec(Matrix(X)), Diagonal_from_vec
-end
-
-function to_vec(X::Transpose)
-    function Transpose_from_vec(x_vec)
-        return Transpose(permutedims(reshape(x_vec, size(X))))
-    end
-    return vec(Matrix(X)), Transpose_from_vec
-end
+j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]
 
-function to_vec(X::Adjoint)
-    function Adjoint_from_vec(x_vec)
-        return Adjoint(conj!(permutedims(reshape(x_vec, size(X)))))
-    end
-    return vec(Matrix(X)), Adjoint_from_vec
+function _j′vp(fdm, f, ȳ::Vector{<:Number}, x::Vector{<:Number})
+    return transpose(first(jacobian(fdm, f, x))) * ȳ
 end
 
-# Non-array data structures
-
-function to_vec(x::Tuple)
-    x_vecs, x_backs = zip(map(to_vec, x)...)
-    sz = cumsum([map(length, x_vecs)...])
-    function Tuple_from_vec(v)
-        return ntuple(n->x_backs[n](v[sz[n]-length(x_vecs[n])+1:sz[n]]), length(x))
-    end
-    return vcat(x_vecs...), Tuple_from_vec
-end
+"""
+    grad(fdm, f, xs...)
 
-# Convert to a vector-of-vectors to make use of existing functionality.
-function to_vec(d::Dict)
-    d_vec_vec = [val for val in values(d)]
-    d_vec, back = to_vec(d_vec_vec)
-    function Dict_from_vec(v)
-        v_vec_vec = back(v)
-        return Dict([(key, v_vec_vec[n]) for (n, key) in enumerate(keys(d))])
-    end
-    return d_vec, Dict_from_vec
-end
+Compute the gradient of `f` for any `xs` for which [`to_vec`](@ref) is defined.
+"""
+grad(fdm, f, xs...) = j′vp(fdm, f, 1, xs...)  # `j′vp` with seed of 1

src/to_vec.jl

@@ -0,0 +1,100 @@
+"""
+    to_vec(x)
+
+Transform `x` into a `Vector`, and return the vector, and a closure which inverts the
+transformation.
+"""
+function to_vec(x::Number)
+    function Number_from_vec(x_vec)
+        return first(x_vec)
+    end
+    return [x], Number_from_vec
+end
+
+# Vectors
+to_vec(x::Vector{<:Number}) = (x, identity)
+function to_vec(x::Vector)
+    x_vecs_and_backs = map(to_vec, x)
+    x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
+    function Vector_from_vec(x_vec)
+        sz = cumsum(map(length, x_vecs))
+        return [backs[n](x_vec[sz[n] - length(x_vecs[n]) + 1:sz[n]]) for n in eachindex(x)]
+    end
+    return vcat(x_vecs...), Vector_from_vec
+end
+
+# Arrays
+function to_vec(x::Array{<:Number})
+    function Array_from_vec(x_vec)
+        return reshape(x_vec, size(x))
+    end
+    return vec(x), Array_from_vec
+end
+
+function to_vec(x::Array)
+    x_vec, back = to_vec(reshape(x, :))
+    function Array_from_vec(x_vec)
+        return reshape(back(x_vec), size(x))
+    end
+    return x_vec, Array_from_vec
+end
+
+# AbstractArrays
+function to_vec(x::T) where {T<:LinearAlgebra.AbstractTriangular}
+    x_vec, back = to_vec(Matrix(x))
+    function AbstractTriangular_from_vec(x_vec)
+        return T(reshape(back(x_vec), size(x)))
+    end
+    return x_vec, AbstractTriangular_from_vec
+end
+
+function to_vec(x::Symmetric)
+    function Symmetric_from_vec(x_vec)
+        return Symmetric(reshape(x_vec, size(x)))
+    end
+    return vec(Matrix(x)), Symmetric_from_vec
+end
+
+function to_vec(X::Diagonal)
+    function Diagonal_from_vec(x_vec)
+        return Diagonal(reshape(x_vec, size(X)...))
+    end
+    return vec(Matrix(X)), Diagonal_from_vec
+end
+
+function to_vec(X::Transpose)
+    function Transpose_from_vec(x_vec)
+        return Transpose(permutedims(reshape(x_vec, size(X))))
+    end
+    return vec(Matrix(X)), Transpose_from_vec
+end
+
+function to_vec(X::Adjoint)
+    function Adjoint_from_vec(x_vec)
+        return Adjoint(conj!(permutedims(reshape(x_vec, size(X)))))
+    end
+    return vec(Matrix(X)), Adjoint_from_vec
+end
+
+# Non-array data structures
+
+function to_vec(x::Tuple)
+    x_vecs, x_backs = zip(map(to_vec, x)...)
+    sz = cumsum(collect(map(length, x_vecs)))
+    function Tuple_from_vec(v)
+        return ntuple(length(x)) do n
+            return x_backs[n](v[sz[n] - length(x_vecs[n]) + 1:sz[n]])
+        end
+    end
+    return reduce(vcat, x_vecs), Tuple_from_vec
+end
+
+# Convert to a vector-of-vectors to make use of existing functionality.
+function to_vec(d::Dict)
+    d_vec, back = to_vec(collect(values(d)))
+    function Dict_from_vec(v)
+        v_vec_vec = back(v)
+        return Dict(key => v_vec_vec[n] for (n, key) in enumerate(keys(d)))
+    end
+    return d_vec, Dict_from_vec
+end