diff --git a/src/fdconversion.jl b/src/fdconversion.jl new file mode 100644 index 000000000..bf5bb3a6f --- /dev/null +++ b/src/fdconversion.jl @@ -0,0 +1,179 @@ +module FDConversion + +import Symbolics +import SymbolicUtils +import FastDifferentiation as FD +import Random + +#= +try implementing these + https://github.com/JuliaSymbolics/SymbolicUtils.jl/blob/master/src/interface.jl +=# + +"""converts from Node to Symbolics expression""" +function _to_symbolics!(a::T, cache::IdDict, variable_map::IdDict) where {T<:FD.Node} + tmp = get(cache, a, nothing) + if tmp !== nothing + return tmp + else + if FD.arity(a) === 0 + if FD.is_constant(a) + cache[a] = Symbolics.Num(FD.value(a)) + elseif FD.is_variable(a) + tmp = Symbolics.variable(FD.value(a)) + cache[a] = tmp + variable_map[a] = tmp + else + throw(ErrorException("Node with 0 children was neither constant nor variable. This should never happen.")) + end + else + if FD.arity(a) === 1 + cache[a] = a.node_value(_to_symbolics!(a.children[1], cache, variable_map)) + else + cache[a] = foldl(a.node_value, _to_symbolics!.(a.children, Ref(cache), Ref(variable_map))) + end + end + end +end + +function to_symbolics(a::T, cache::IdDict=IdDict(), variable_map::IdDict=IdDict()) where {T<:FD.Node} + _to_symbolics!(a, cache, variable_map) + return cache[a], variable_map +end +export to_symbolics + +function to_symbolics(a::AbstractArray{T}, cache::IdDict=IdDict(), variable_map::IdDict=IdDict()) where {T<:FD.Node} + _to_symbolics!.(a, Ref(cache), Ref(variable_map)) + return map(x -> cache[x], a), variable_map +end + +"""Converts expression x from Symbolics to FastDifferentiation form. Returns a tuple of the **FD** form of the expression and a Dict that maps Symbolics variables to **FD** variables""" +function to_fd(x::Real, cache=IdDict(), substitutions=IdDict()) + result = _to_FD!(x, cache, substitutions) + syms = collect(filter(x -> SymbolicUtils.issym(x), keys(cache))) + sym_map = Dict(zip(syms, map(x -> cache[x], syms))) + + return result, sym_map +end +export to_fd + +function to_fd(x::AbstractArray{<:Real}) + cache = IdDict() + substitutions = IdDict() + result = _to_FD!.(x, Ref(cache), Ref(substitutions)) + syms = collect(filter(x -> SymbolicUtils.issym(x), keys(cache))) + sym_map = Dict(zip(syms, map(x -> cache[x], syms))) #map from symbolics variables to FD variables + return result, sym_map +end + +function _to_FD!(sym_node, cache::IdDict, visited::IdDict) + # Substitutions are done on a Node graph, not a SymbolicsUtils.Sym graph. As a consequence the values + # in substitutions Dict are Node not Sym type. cache has keys (op,args...) where op is generally a function type but sometimes a Sym, + # and args are all Node types. + + @assert !(typeof(sym_node) <: Symbolics.Arr) "Differentiation of expressions involving arrays and array variables is not yet supported." + if SymbolicUtils.istree(Symbolics.unwrap(sym_node)) + @assert !SymbolicUtils.issym(SymbolicUtils.operation(Symbolics.unwrap(sym_node))) "expressions of the form `q(t)` not yet supported." + end + + symx = isa(sym_node, Symbolics.Num) ? sym_node.val : sym_node + @assert typeof(symx) != Symbolics.Num + + + tmpsub = get(visited, symx, nothing) + if tmpsub !== nothing + return visited[symx] #substitute Node object for symbolic object + end + + tmp = get(cache, symx, nothing) + + if tmp !== nothing + return tmp + elseif !SymbolicUtils.istree(symx) + if SymbolicUtils.issym(symx) + tmpnode = FD.Node(Symbol(symx)) + else #must be a number of some kind + tmpnode = FD.Node(symx) + end + + cache[symx] = tmpnode + + return tmpnode + else + numargs = length(SymbolicUtils.arguments(symx)) + symargs = SymbolicUtils.arguments(symx) + + args = _to_FD!.(symargs, Ref(cache), Ref(visited)) + + key = (SymbolicUtils.operation(symx), args...) + tmp = get(cache, key, nothing) + + if tmp !== nothing + return tmp + else + tmpnode = FD.Node(SymbolicUtils.operation(symx), args...) + cache[key] = tmpnode + + return tmpnode + end + end +end + + +function remap(fd_function_to_call, symbolics_function, differentiation_variables::AbstractVector{Symbolics.Num}, fast_differentiation::Bool) + fd_func, variable_map = to_fd(symbolics_function) + partial_vars = map(x -> variable_map[x], differentiation_variables) + tmp = fd_function_to_call(fd_func, partial_vars) + if fast_differentiation + return fd_func, variable_map #return FastDifferentiation expression to be passed to make_function for efficient evaluation + else + reverse_map = IdDict{Any,Any}(map(x -> variable_map[x] => x, differentiation_variables)) + reverse_variable_map = deepcopy(reverse_map) + return to_symbolics(tmp, reverse_map, reverse_variable_map) #return Symbolics expression for further evaluation + end +end + + +""" +Converts from `Symbolics` form to `FastDifferentiation` form and computes Jacobian with respect to `diff_variables`. +If `fast_differentiation=false` returns result in Symbolics form. If `fast_differentiation=true` the result will be a 2-tuple. The first tuple entry will be the jacobian of `symbolics_function` converted to `FastDifferentiation` form. +The second tuple term will be `differentiation_variables` converted to `FastDifferentiation` form. +This tuple can be passed to `fd_make_function` to create an efficient executable.""" +fd_jacobian(symbolics_function::AbstractArray{Symbolics.Num}, differentiation_variables::AbstractVector{Symbolics.Num}; fast_differentiation=false) = remap(FD.jacobian, symbolics_function, differentiation_variables, fast_differentiation) +export fd_jacobian + +""" +Converts from `Symbolics` form to `FastDifferentiation` form and computes sparse Jacobian with respect to `diff_variables`. +If `fast_differentiation=false` returns result in Symbolics form. If `fast_differentiation=true` the result will be a 2-tuple. The first tuple entry will be the sparse jacobian of `symbolics_function` converted to `FastDifferentiation` form. +The second tuple term will be `differentiation_variables` converted to `FastDifferentiation` form. +This tuple can be passed to `fd_make_function` to create an efficient executable.""" +fd_sparse_jacobian(symbolics_function::AbstractArray{Symbolics.Num}, differentiation_variables::AbstractVector{Symbolics.Num}; fast_differentiation=false) = remap(FD.sparse_jacobian, symbolics_function, differentiation_variables, fast_differentiation) +export fd_sparse_jacobian + +""" +Converts from `Symbolics` form to `FastDifferentiation` form and computes Hessian with respect to `diff_variables`. +If `fast_differentiation=false` returns result in Symbolics form. If `fast_differentiation=true` the result will be a 2-tuple. The first tuple entry will be the hessian of `symbolics_function` converted to `FastDifferentiation` form. +The second tuple term will be `differentiation_variables` converted to `FastDifferentiation` form. +This tuple can be passed to `fd_make_function` to create an efficient executable.""" +fd_hessian(symbolics_function::Symbolics.Num, differentiation_variables::AbstractVector{Symbolics.Num}; fast_differentiation=false) = remap(FD.hessian, symbolics_function, differentiation_variables, fast_differentiation) +export fd_sparse_hessian + +""" +Converts from `Symbolics` form to `FastDifferentiation` form and computes sparse Hessian with respect to `diff_variables`. +If `fast_differentiation=false` returns result in Symbolics form. If `fast_differentiation=true` the result will be a 2-tuple. The first tuple entry will be the sparse Hessian of `symbolics_function` converted to `FastDifferentiation` form. +The second tuple term will be `differentiation_variables` converted to `FastDifferentiation` form. +This tuple can be passed to `fd_make_function` to create an efficient executable.""" +fd_sparse_hessian(symbolics_function::Symbolics.Num, differentiation_variables::AbstractVector{Symbolics.Num}; fast_differentiation=false) = remap(FD.sparse_hessian, symbolics_function, differentiation_variables, fast_differentiation) +export fd_sparse_hessian + +""" +Creates an efficient runtime generated function to evaluate the function defined in `a[1]` with arguments given by `a[2]`. Used in conjuction with fd_jacobian,fd_sparse_jacobian,fd_hessian,fd_sparse_hessian. +""" +fd_make_function(a::Tuple{T,S}; in_place=false) where {T<:AbstractArray{<:FD.Node},S<:AbstractVector{<:FD.Node}} = FD.make_function(a[1], a[2], in_place=in_place) +export fd_make_function +#etc. for Jv Jᵀv Hv + +# export FastDifferentiation.make_function + +end # module FSDConvert + diff --git a/test/fdconversion.jl b/test/fdconversion.jl new file mode 100644 index 000000000..16714a544 --- /dev/null +++ b/test/fdconversion.jl @@ -0,0 +1,143 @@ + +import Symbolics +import FastDifferentiation as FD +using FDConversion +using Random +using Test + +#Function used in tests for FDConversion +function P(::Type{T}, l, m, z::T) where {T} + if l == 0 && m == 0 + return T(1) + elseif l == m + return (1 - 2m) * P(T, m - 1, m - 1, z) + elseif l == m + 1 + return (2m + 1) * z * P(T, m, m, z) + else + return ((2l - 1) / (l - m) * z * P(T, l - 1, m, z) - (l + m - 1) / (l - m) * P(T, l - 2, m, z)) + end +end + + +function S(::Type{T}, m, x::T, y::T) where {T} + if m == 0 + return T(0) + else + return x * C(T, m - 1, x, y) - y * S(T, m - 1, x, y) + end +end + + +function C(::Type{T}, m, x::T, y::T) where {T} + if m == 0 + return T(1) + else + return x * S(T, m - 1, x, y) + y * C(T, m - 1, x, y) + end +end + +function factorial_approximation(::Type{T}, x) where {T} + local n1 = x + sqrt(2 * T(π) * n1) * (n1 / T(ℯ) * sqrt(n1 * sinh(1 / T(n1)) + 1 / (810 * T(n1)^6)))^n1 +end + + +function compare_factorial_approximation() + for n in 1:30 + println("n $n relative error $((factorial(big(n))-factorial_approximation(BigFloat,n))/factorial(big(n)))") + end +end + + +function N(::Type{T}, l, m) where {T} + @assert m >= 0 + if m == 0 + return sqrt((2l + 1 / (4 * T(π)))) + else + # return sqrt((2l+1)/2π * factorial(big(l-m))/factorial(big(l+m))) + #use factorial_approximation instead of factorial because the latter does not use Stirlings approximation for large n. Get error for n > 2 unless using BigInt but if use BigInt get lots of rational numbers in symbolic result. + return sqrt((2l + 1) / 2 * T(π) * factorial_approximation(T, l - m) / factorial_approximation(T, l + m)) + end +end + + +"""l is the order of the spherical harmonic""" +function Y(::Type{T}, l, m, x::T, y::T, z::T) where {T} + @assert l >= 0 + @assert abs(m) <= l + if m < 0 + return N(T, l, abs(m)) * P(T, l, abs(m), z) * S(T, abs(m), x, y) + else + return N(T, l, m) * P(T, l, m, z) * C(T, m, x, y) + end +end + +Y(l, m, x::T, y::T, z::T) where {T<:FD.Node} = Y(FD.Node, l, m, x, y, z) +Y(l, m, x::T, y::T, z::T) where {T<:Number} = Y(T, l, m, x, y, z) + + +function SHFunctions(max_l, x, y, z) + @assert typeof(x) == typeof(y) == typeof(z) + result = Vector(undef, 0) + + for l in 0:max_l-1 + for m in -l:l + push!(result, Y(l, m, x, y, z)) + end + end + + return result +end + + + +@testset "conversion from Symbolics to FD" begin + Symbolics.@variables x y + + symbolics_expr = x^2 + y * (x^2) + dag, tmp = to_fd(symbolics_expr) + vars = collect(values(tmp)) + fdx, fdy = FD.value(vars[1]) == :x ? (vars[1], vars[2]) : (vars[2], vars[1]) #need to find the variables since they can be in any order + + correct_fun = FD.make_function([dag], [fdx, fdy]) + + + #verify that all the node expressions exist in the dag. Can't rely on them being in a particular order because Symbolics can + #arbitrarily choose how to reorder trees. + num_tests = 100 + rng = Random.Xoshiro(8392) + for _ in 1:num_tests + (xval, yval) = rand(rng, 2) + FDval = correct_fun([xval, yval])[1] + Syval = Symbolics.substitute(symbolics_expr, Dict([(x, xval), (y, yval)])) + + @test isapprox(FDval, Syval.val) + + end +end + + +@testset "conversion from FD to Symbolics" begin + order = 8 + FD.@variables x y z + + FD_funcs = FD.Node.(SHFunctions(order, x, y, z)) + Sym_funcs, variables = to_symbolics(FD_funcs) + + sx, sy, sz = map(p -> variables[p], [x, y, z]) + + FD_eval = FD.make_function(FD_funcs, [x, y, z]) + rng = Random.Xoshiro(8392) + num_tests = 100 + for _ in 1:num_tests + tx, ty, tz = rand(rng, BigFloat, 3) + subs = Dict([sx => tx, sy => ty, sz => tz]) + res = Symbolics.substitute.(Sym_funcs, Ref(subs)) + + FD_res = FD_eval([tx, ty, tz]) + + @test isapprox(FD_res, res, atol=1e-12) + end +end + +