FixedEffectModels.jl
FixedEffectModels.jl copied to clipboard
Autodiff across FE result
Someone noted to me that FixedEffectModels.jl is tricky to use AD on because there are so many explicit Float64
type constraints -- does anyone have a good sense of how much effort it would take to remove/parameterize/reduce the explicit type constraints here?
As an example, the FixedEffectModel
struct has a lot of Float64
explicit types that could be parametric instead.
https://github.com/FixedEffects/FixedEffectModels.jl/blob/2a2661e508c579b6c9c8a82f1388debe0c8229e8/src/FixedEffectModel.jl#L8-L45
Is there an appetite for this? I think it'd be lovely to be able to AD through high-dimensional fixed effect estimates.
What do you have in mind? coef
inherits its element type from y
and X
. These are constructed from the dataframe and formula, and converted to Float64
. I could imagine getting the type of y
and X
from the types of the columns of the dataframe to allow automatic differentiation with respect to the data, but it seems somewhat unusual to try to autodiff through a dataframe.
If you use the lower level code in FixedEffects.jl, there is some compatibility with AD.
using ForwardDiff, FiniteDiff, FixedEffects, LinearAlgebra
N = 10
K = 2
d1=repeat(1:(N÷2), inner = 2)
d2=repeat(1:2, inner = N÷2)
p1 = FixedEffect(d1)
p2 = FixedEffect(d2)
x = rand(N,K)
y = rand(N)
function regfe(y,x,fes)
x = copy(x)
y = copy(y)
w = FixedEffects.uweights(eltype(x), size(x, 1))
feM = AbstractFixedEffectSolver{eltype(x)}(fes, w, Val{:cpu}, Threads.nthreads())
solve_residuals!(x, feM)
return((x'*x) \ (x'*y))
end
regfe(y, x, [p1, p2])
# derivative wrt x
J1 = ForwardDiff.jacobian(x->regfe(y, reshape(x,N,K), [p1, p2]) , vec(x))
J2 = FiniteDiff.finite_difference_jacobian(x->regfe(y, reshape(x,N,K), [p1, p2]) , vec(x))
norm(J1- J2)
# derivative wrt y
J1 = ForwardDiff.jacobian(y->regfe(y, reshape(x,N,K), [p1, p2]) , y)
J2 = FiniteDiff.finite_difference_jacobian(y->regfe(y, reshape(x,N,K), [p1, p2]) , y)
norm(J1- J2)
It's dreadfully slow, but appears to give correct results. Reverse mode packages are probably harder. The in-place mutation of solve_residuals!
and iteration in lsmr
probably mean it'd require defining custom methods. Doing so could also greatly improve the ForwardDiff performance.