Autodiff across FE result

[just-cameron]

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.

FixedEffectModels.jl/src/FixedEffectModel.jl

Lines 8 to 45 in 2a2661e

	struct FixedEffectModel <: RegressionModel
	coef::Vector{Float64} # Vector of coefficients
	vcov::Matrix{Float64} # Covariance matrix
	vcov_type::CovarianceEstimator
	nclusters::Union{NamedTuple, Nothing}
	esample::BitVector # Is the row of the original dataframe part of the estimation sample?
	residuals::Union{AbstractVector, Nothing}
	fe::DataFrame
	fekeys::Vector{Symbol}
	coefnames::Vector # Name of coefficients
	yname::Union{String, Symbol} # Name of dependent variable
	formula::FormulaTerm # Original formula
	formula_predict::FormulaTerm
	contrasts::Dict
	nobs::Int64 # Number of observations
	dof_residual::Int64 # nobs - degrees of freedoms
	rss::Float64 # Sum of squared residuals
	tss::Float64 # Total sum of squares
	r2::Float64 # R squared
	adjr2::Float64 # R squared adjusted
	F::Float64 # F statistics
	p::Float64 # p value for the F statistics
	# for FE
	iterations::Union{Int, Nothing} # Number of iterations
	converged::Union{Bool, Nothing} # Has the demeaning algorithm converged?
	r2_within::Union{Float64, Nothing} # within r2 (with fixed effect
	# for IV
	F_kp::Union{Float64, Nothing} # First Stage F statistics KP
	p_kp::Union{Float64, Nothing} # First Stage p value KP
	end

Is there an appetite for this? I think it'd be lovely to be able to AD through high-dimensional fixed effect estimates.

[schrimpf]

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.