Handle omitted event-time terms via reduced-system IW fit

Author: xiaobaaaa

src/avar.jl

@@ -19,6 +19,28 @@ end
 
 _ginvsym(X::AbstractMatrix) = Matrix(-FixedEffectModels.invsym!(Symmetric(copy(X))))
 
+function _share_vcov_expand(reduced::AbstractMatrix, keep_cols::AbstractVector{Bool}, nresid::Integer)
+    full_cols = length(keep_cols)
+    keep_idx = findall(keep_cols)
+    expanded_idx = Int[]
+
+    for resid_idx in 1:nresid
+        append!(expanded_idx, (resid_idx - 1) * full_cols .+ keep_idx)
+    end
+
+    full = zeros(Float64, full_cols * nresid, full_cols * nresid)
+    full[expanded_idx, expanded_idx] .= Matrix(reduced)
+    return Symmetric(full)
+end
+
+function _share_vcov_dispatch(sample::DataFrame, emat::Matrix{Float64}, Xmat::Matrix{Float64}, vcov::CovarianceEstimator)
+    vcov isa Vcov.SimpleCovariance && return _share_vcov_simple(emat, Xmat)
+    vcov isa Vcov.RobustCovariance && return _share_vcov_robust(emat, Xmat)
+    vcov isa Vcov.ClusterCovariance && return _share_vcov_cluster(sample, emat, Xmat, vcov)
+
+    throw(ArgumentError("Unsupported covariance estimator $(typeof(vcov)) for the cohort-share step."))
+end
+
 function _share_vcov(table, e::AbstractMatrix, X::AbstractMatrix, vcov::CovarianceEstimator)
     sample = DataFrame(table; copycols = false)
     Xmat = Matrix{Float64}(X)
@@ -31,11 +53,13 @@ function _share_vcov(table, e::AbstractMatrix, X::AbstractMatrix, vcov::Covarian
         emat = emat[keep, :]
     end
 
-    vcov isa Vcov.SimpleCovariance && return _share_vcov_simple(emat, Xmat)
-    vcov isa Vcov.RobustCovariance && return _share_vcov_robust(emat, Xmat)
-    vcov isa Vcov.ClusterCovariance && return _share_vcov_cluster(sample, emat, Xmat, vcov)
+    keep_cols = vec(any(.!iszero.(Xmat), dims = 1))
+    if all(keep_cols)
+        return _share_vcov_dispatch(sample, emat, Xmat, vcov)
+    end
 
-    throw(ArgumentError("Unsupported covariance estimator $(typeof(vcov)) for the cohort-share step."))
+    reduced = _share_vcov_dispatch(sample, emat, Xmat[:, keep_cols], vcov)
+    return _share_vcov_expand(reduced, keep_cols, size(emat, 2))
 end
 
 function _share_vcov_simple(e::Matrix{Float64}, X::Matrix{Float64})

src/eventfit.jl

@@ -1,6 +1,52 @@
-function eventreg(@nospecialize(df),     
-    @nospecialize(formula::FormulaTerm),
-    @nospecialize(rel_varlist::Vector{Symbol}),
+function _nonzero_relmask(df::DataFrame, nvarlist::Vector{Symbol})
+    isempty(nvarlist) && return Bool[]
+    X = Matrix{Float64}(df[:, nvarlist])
+    return vec(any(.!iszero.(X), dims = 1))
+end
+
+function _estimable_relmask(vcov_mat::AbstractMatrix, nrel_times::Integer, ncohort::Integer)
+    nrel_times == 0 && return Bool[]
+    interaction_dim = nrel_times * ncohort
+    diag_v = diag(Matrix(vcov_mat)[1:interaction_dim, 1:interaction_dim])
+    keep_rel = falses(nrel_times)
+    for i in 1:nrel_times
+        block = ((i - 1) * ncohort + 1):(i * ncohort)
+        keep_rel[i] = any(.!isnan.(diag_v[block]))
+    end
+    return keep_rel
+end
+
+function _zero_omitted_rowscols(vcov_mat::AbstractMatrix)
+    VV = Matrix{Float64}(vcov_mat)
+    omitted = isnan.(diag(VV))
+    kept = .!omitted
+
+    if any(kept) && any(isnan.(VV[kept, kept]))
+        throw(ArgumentError("Non-omitted entries of the interacted-regression covariance matrix contain NaN values."))
+    end
+
+    if any(omitted)
+        VV[omitted, :] .= 0.0
+        VV[:, omitted] .= 0.0
+    end
+
+    return VV
+end
+
+function _reinsert_rel_outputs(coef_reduced::AbstractVector, vcov_reduced::AbstractMatrix, keep_rel::AbstractVector{Bool}, full_size::Integer)
+    coef_full = zeros(Float64, full_size)
+    vcov_full = fill(NaN, full_size, full_size)
+
+    keep_idx = findall(keep_rel)
+    coef_full[keep_idx] .= coef_reduced
+    vcov_full[keep_idx, keep_idx] .= Matrix(vcov_reduced)
+
+    return coef_full, Symmetric(vcov_full)
+end
+
+function eventreg(@nospecialize(df),     
+    @nospecialize(formula::FormulaTerm),
+    @nospecialize(rel_varlist::Vector{Symbol}),
     @nospecialize(control_cohort::Symbol),
     @nospecialize(cohort::Symbol),
     @nospecialize(vcov::CovarianceEstimator = Vcov.simple());
@@ -43,21 +89,29 @@ function StatsAPI.fit(::Type{EventStudyInteract},
     if dof_add != 0
         @warn "The dof_add keyword is no longer forwarded to FixedEffectModels.reg on version 1.13 and is currently ignored."
     end    
-    # Prepare the varlists
-    nvarlist = Symbol[]
-    dvarlist = Symbol[]
-    for l in rel_varlist
+    # Prepare the varlists
+    nvarlist = Symbol[]
+    dvarlist = Symbol[]
+    for l in rel_varlist
         push!(dvarlist,l)
         n_l = Symbol("n",l)
         df[!,n_l] = df[!,l]
-        df[!,n_l][df[!,control_cohort].==1] .= 0
-        push!(nvarlist,n_l)
-    end
+        df[!,n_l][df[!,control_cohort].==1] .= 0
+        push!(nvarlist,n_l)
+    end
+    full_dvarlist = copy(dvarlist)
 
-    cohort_list = unique(df[!,cohort][df[!,control_cohort].==0])
-    cohort_list = sort(cohort_list)
-    cohort_list = Int.(cohort_list)
-    nrel_times = length(nvarlist)
+    prekeep_rel = _nonzero_relmask(df, nvarlist)
+    if !all(prekeep_rel)
+        nvarlist = nvarlist[prekeep_rel]
+        dvarlist = dvarlist[prekeep_rel]
+    end
+    isempty(nvarlist) && error("No estimable relative-time indicators remain after removing unsupported columns.")
+
+    cohort_list = unique(df[!,cohort][df[!,control_cohort].==0])
+    cohort_list = sort(cohort_list)
+    cohort_list = Int.(cohort_list)
+    nrel_times = length(nvarlist)
     ncohort = length(cohort_list)
 
     # Prepare interaction terms for the interacted regression
@@ -97,11 +151,21 @@ function StatsAPI.fit(::Type{EventStudyInteract},
     progress_bar = progress_bar,
     subset = subset)
 
-    #Caculate the weights.
-
-    df = df[result.esample .== 1, :]
-
-
+    postkeep_rel = _estimable_relmask(result.vcov, nrel_times, ncohort)
+    if !all(postkeep_rel)
+        nvarlist = nvarlist[postkeep_rel]
+        dvarlist = dvarlist[postkeep_rel]
+    end
+    isempty(nvarlist) && error("No estimable relative-time indicators remain after fitting the interacted regression.")
+
+    nrel_times = length(nvarlist)
+    rel_keep_mask = falses(length(rel_varlist))
+    rel_keep_mask[findall(prekeep_rel)] .= postkeep_rel
+
+    #Caculate the weights.
+
+    df = df[result.esample .== 1, :]
+
     if ncohort == 1
         ff_w = ones(1, nrel_times)
         Sigma_ff = zeros(nrel_times, nrel_times)
@@ -127,40 +191,25 @@ function StatsAPI.fit(::Type{EventStudyInteract},
         Sigma_ff = _share_vcov(stage2df, e, X, vcov)
     end
 
-    b = result.coef
-
-    V = diag(result.vcov)
-
-    replace!(V, NaN => 0)
-
-    # Convert the delta estimate vector to a matrix where each column is a relative time
-    end_ = 0
-    evt_bb = zeros(ncohort,0)
-    evt_VV = zeros(ncohort,0)
-
-    for i in 1:nrel_times
-        start = end_ + 1
-        end_ = start + ncohort - 1
-        b_i = b[start:end_]
-        evt_bb = hcat(evt_bb, b_i)
-        V_i = V[start:end_]
-        evt_VV = hcat(evt_VV, V_i)
-    end
-
-    # Take weighted average for IW estimators
-    w = ff_w
-    delta = evt_bb
-    b_iw = sum(w .* delta, dims=1)
-    nc, nr = size(w)
-
-    # Ptwise variance from cohort share estimation and interacted regression
-    VV = result.vcov # VCV from the interacted regression
-    replace!(VV, NaN => 0)
-
-    VV = VV[1:nr*nc, 1:nr*nc] # in case reports _cons
-
-    wlong = w'.*repeat(ee(1, nr)', 1, nc)
-
+    interaction_keep = repeat(postkeep_rel, inner = ncohort)
+    interaction_idx = findall(interaction_keep)
+    b = result.coef[interaction_idx]
+
+    # Convert the interacted coefficients to a matrix where each column is a relative time
+    evt_bb = reshape(b, ncohort, nrel_times)
+
+    # Take weighted average for IW estimators
+    w = ff_w
+    delta = evt_bb
+    b_iw = sum(w .* delta, dims=1)
+    nc, nr = size(w)
+
+    # Ptwise variance from cohort share estimation and interacted regression
+    VV_full = _zero_omitted_rowscols(result.vcov) # VCV from the interacted regression
+    VV = VV_full[interaction_idx, interaction_idx]
+
+    wlong = w'.*repeat(ee(1, nr)', 1, nc)
+
     for i in 2:nrel_times      # loop from 2 to nrel_times
         wlong = [wlong (w' .* repeat(ee(i, nr)', 1, nc))]  # concatenate 
     end
@@ -179,10 +228,9 @@ function StatsAPI.fit(::Type{EventStudyInteract},
         end
     end
 
-    coef_iw = vec(b_iw')
-
-    vcov_iw = Symmetric(V_iw)
-    vcov_iw_diag = diag(vcov_iw)
+    coef_iw_reduced = vec(b_iw')
+    vcov_iw_reduced = Symmetric(V_iw)
+    coef_iw, vcov_iw = _reinsert_rel_outputs(coef_iw_reduced, vcov_iw_reduced, rel_keep_mask, length(rel_varlist))
 
 
 
@@ -213,7 +261,7 @@ function StatsAPI.fit(::Type{EventStudyInteract},
     nclusters = result.nclusters
     esample = result.esample
     fekeys = result.fekeys
-    coef_names = [string(x) for x in dvarlist]
+    coef_names = [string(x) for x in full_dvarlist]
     # coef_names = coefnames
     response_name = responsename(result)
     # response_name, coef_names = coefnames(formula_schema)

test/runtests.jl

@@ -351,6 +351,39 @@ try
             @test all(isfinite, diag(vcov(model)))
         end
 
+        @testset "Zero relative-time columns do not poison clustered share vcov" begin
+            df = make_test_panel()
+            df.g_3 = zeros(Int, nrow(df))
+
+            formula = term(:Y) ~ sum(fe.(term.([:id, :t])))
+            baseline = eventreg(
+                df,
+                formula,
+                [:g_2, :g0, :g1, :g2],
+                :never_treat,
+                :first_treat,
+                Vcov.cluster(:id);
+                progress_bar = false,
+            )
+            with_zero = eventreg(
+                df,
+                formula,
+                [:g_3, :g_2, :g0, :g1, :g2],
+                :never_treat,
+                :first_treat,
+                Vcov.cluster(:id);
+                progress_bar = false,
+            )
+
+            @test coefnames(with_zero) == string.([:g_3, :g_2, :g0, :g1, :g2])
+            @test coef(with_zero)[1] == 0.0
+            @test isnan(diag(vcov(with_zero))[1])
+            @test all(isnan, vcov(with_zero)[1, :])
+            @test all(isnan, vcov(with_zero)[:, 1])
+            @test coef(with_zero)[2:end] ≈ coef(baseline) atol = 1.0e-10 rtol = 1.0e-10
+            @test diag(vcov(with_zero))[2:end] ≈ diag(vcov(baseline)) atol = 1.0e-10 rtol = 1.0e-10
+        end
+
         @testset "Single-cohort equivalence to direct FE regression" begin
             report = run_single_cohort_equivalence()