ForwardDiff.jl
ForwardDiff.jl copied to clipboard
Unitful support
Unitful and ForwardDiff don't play well together. But it would be good if they did! stripping the units out of a model to run ForwardDiff for a jacobian etc. isn't always feasible, in my case it would mean adding a lot of manual conversions and tests.
Loosening some Real types to Number would get some of the way, but there are probably deeper problems.
Edit: to make more sense.
There seems to be a relatively easy solution of stripping the units from the input vector and reapplying the units to the Dual number vector passed back from ForwardDiff, and stripping them back off before returning the results vector.
The problem with this is when using ForwardDiff inside other packages - say DiffEqSensitivity, which will expect units in the output based on the input, and it all gets very messy to work with.
+1. To provide a concrete example of the failure:
julia> using Unitful, ForwardDiff
julia> f(x::Vector) = sum(sin, x) + prod(tan, x) * sum(sqrt, x);
julia> x = [1.0u"m", 2.0u"m^2"]
2-element Array{Quantity{Float64,D,U} where U where D,1}:
1.0 m
2.0 m^2
julia> g = x -> ForwardDiff.gradient(f, x); # g = ∇f
julia> g(x)
ERROR: MethodError: no method matching ForwardDiff.Partials(::Tuple{Float64,Quantity{Float64,Unitful.Dimensions{()},Unitful.FreeUnits{(),Unitful.Dimensions{()}}}})
Closest candidates are:
ForwardDiff.Partials(::Tuple{Vararg{V,N}}) where {N, V} at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/partials.jl:2
Stacktrace:
[1] macro expansion at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/partials.jl:0 [inlined]
[2] single_seed(::Type{ForwardDiff.Partials{2,Quantity{Float64,D,U} where U where D}}, ::Val{1}) at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/partials.jl:10
[3] macro expansion at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/apiutils.jl:0 [inlined]
[4] construct_seeds(::Type{ForwardDiff.Partials{2,Quantity{Float64,D,U} where U where D}}) at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/apiutils.jl:51
[5] ForwardDiff.GradientConfig(::typeof(f), ::Array{Quantity{Float64,D,U} where U where D,1}, ::ForwardDiff.Chunk{2}, ::ForwardDiff.Tag{typeof(f),Quantity{Float64,D,U} where U where D}) at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/config.jl:121 (repeats 3 times)
[6] gradient(::Function, ::Array{Quantity{Float64,D,U} where U where D,1}) at /home/cst-jean/.julia/packages/ForwardDiff/hnKaN/src/gradient.jl:15
Loosening some Real types to Number would get some of the way, but there are probably deeper problems.
It also assumes uniformly-typed inputs:
struct Partials{N,V} <: AbstractVector{V}
values::NTuple{N,V}
end
There seems to be a relatively easy solution of stripping the units from the input vector and reapplying the units to the Dual number vector passed back from ForwardDiff, and stripping them back off before returning the results vector.
If there was a generic way to do that, that would be very useful.
Is it possible that this is essentially fixed by https://github.com/JuliaDiff/ForwardDiff.jl/pull/369?
Possibly yes. At least the Dual number types. The differentiation methods might need more work.
I remember finding a few other problems besides loosening Real types, but I can't remember details. It would be good to have tests for units working here and in packages like DiffEqSensitivity, but I don't have any time to put that together currently.
I haven't tested it, so it might work strictly speaking. But as long as it assumes an input of the form AbstractVector{T}
, won't performance be very poor when T is not a concrete type? Each unit has a different type.
I just tried to play with taking the derivative of a simple Unitful-embedded function:
julia> using Unitful, ForwardDiff
julia> foo(x) = x / 1.0u"s"
foo (generic function with 1 method)
julia> ForwardDiff.derivative(foo, 2.0)
ERROR: MethodError: no method matching extract_derivative(::Type{ForwardDiff.Tag{typeof(foo),Float64}}, ::Quantity{ForwardDiff.Dual{ForwardDiff.Tag{typeof(foo),Float64},Float64,1},𝐓^-1,Unitful.FreeUnits{(s^-1,),𝐓^-1,nothing}})
so I am not sure how https://github.com/JuliaDiff/ForwardDiff.jl/pull/369 fixes this issue, but maybe that's not how to do it? Is there an example of ForwardDiff and Unitful working together?
Is there any plan to make this work?
FYI for anybody else looking, the code above works if you strip the unit again before computing the derivative.
I.e. something like
ForwardDiff.derivative(ustrip∘foo, 1.0)
works. Generally, the input and output need to be not-unitful, but inside the function you can use units.
More generally, I found that sometimes autodiff was working as I expected, and sometimes not.
For instance, I've found that generally dealing with Unitful quantities works /inside functions/, as long as the input and output types are unit stripped. I've boiled it down to the following observation:
We're dealing with two type wrappers around, say, Float64
. One is the Dual
wrapper, and one is the Quantity
wrapper.
Generally, for both many operations are overloaded, so that we can compute (Base.:+(lhs::Dual, rhs::Dual)
and (Base.:+(lhs::Quantity, rhs::Quantity)
and so on.
However, when we want both type wrappers at the same time, it's not obvious whether we should use Quantity{Dual{T}}
or Dual{Quantity{T}}
-- in some sense they live in an orthogonal world. However, practically speaking, it seems towork well to work with Quantity{Dual{T}}
.
Notice for example:
using ForwardDiff: Dual
using Unitful, Unitful.DefaultSymbols
Dual(1.0)m # has type Quantity{Dual{T}} -> works
Dual(1.0m) # has type Dual{Quantity{T}} -> doesn't work
For me, that means two things:
- I write all my Unitful function interfaces such that they accept an arbitrary numeric type
<:Real
. See my PR https://github.com/PainterQubits/Unitful.jl/pull/698 that allows writing interfaces like
circumference_of_circle(r::WithDims(u"m")) = pi*r^2
which essentially allows Quantity{Float64, Meters}
and Quantity{Dual{Float64}, Meters}
, etc. (it fixes the dimension template parameter but leaves open the numeric type parameter).
- When using autodiff, I wrap the function from "both sides" with
ustrip
(at the "end") and by adding the unit (at the "beginning"), i.e.
area_of_circle(r::WithDims(u"m")) = pi*r^2
let f = (x->ustrip(m^2, x))∘area_of_circle∘(x->x*m)
ForwardDiff.derivative(f, 1.0)
end
Generally, this works also for gradient
, jacobian
, and hessian
:
import ForwardDiff: derivative, gradient, jacobian, hessian
import Underscores: @_
derivative((@_ ustrip(m^2,__) ∘ area_of_circle ∘ (__*m)), 1.0)
gradient((@_ ustrip.(m,__) ∘ sum ∘ (__.*m)), [1.;2;3])
jacobian((@_ ustrip.(m^2,__) ∘ [sum(x->x^2, __); 2*sum(__)^2] ∘ (__.*m)), [1.;2;3])
hessian((@_ ustrip(m^5,__) ∘ (x->x[1]^2*x[2]^3) ∘ (__.*m)), [1.;1.;1.])