Open cpfiffer opened 2 years ago
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.
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 ofFloat64
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.