diff --git a/docs/src/release_notes.md b/docs/src/release_notes.md index 3a6b0b9..90e788b 100644 --- a/docs/src/release_notes.md +++ b/docs/src/release_notes.md @@ -4,11 +4,12 @@ These release notes adhere to the [keep a changelog](https://keepachangelog.com/ ## Version [0.8.0](https://github.com/dscolby/CausalELM.jl/releases/tag/v0.8.0) - 2024-10-31 ### Added * Implemented randomization inference-based confidence intervals [#78](https://github.com/dscolby/CausalELM.jl/issues/78) +* Added marginal effects to model summaries [#78](https://github.com/dscolby/CausalELM.jl/issues/78) ### Fixed * Removed unnecessary include and using statements * Slightly sped up the randomization inference implementation and clarified it in the docs [#77](https://github.com/dscolby/CausalELM.jl/issues/77) * Fixed the randomization inference index selection procedure for interrupted time series estimators -* Inlined certain methods to slightly improve performance [#79](https://github.com/dscolby/CausalELM.jl/issues/79) +* Inlined certain methods to slightly improve performance [#76](https://github.com/dscolby/CausalELM.jl/issues/76) ## Version [v0.7.0](https://github.com/dscolby/CausalELM.jl/releases/tag/v0.7.0) - 2024-06-22 ### Added diff --git a/src/estimators.jl b/src/estimators.jl index 44550e7..9cec74f 100644 --- a/src/estimators.jl +++ b/src/estimators.jl @@ -48,6 +48,7 @@ mutable struct InterruptedTimeSeries Y₀::Array{Float64} X₁::Array{Float64} Y₁::Array{Float64} + marginal_effect::Float64 @model_config individual_effect end @@ -77,6 +78,7 @@ function InterruptedTimeSeries( float(Y₀), X₁, float(Y₁), + NaN, "difference", true, task, @@ -137,6 +139,7 @@ julia> m5 = GComputation(x_df, t_df, y_df) mutable struct GComputation <: CausalEstimator @standard_input_data @model_config average_effect + marginal_effect::Float64 ensemble::ELMEnsemble function GComputation( @@ -173,6 +176,7 @@ mutable struct GComputation <: CausalEstimator num_feats, num_neurons, NaN, + NaN, ) end end @@ -220,6 +224,7 @@ julia> m2 = DoubleMachineLearning(x_df, t_df, y_df) mutable struct DoubleMachineLearning <: CausalEstimator @standard_input_data @model_config average_effect + marginal_effect::Float64 folds::Integer end @@ -256,6 +261,7 @@ function DoubleMachineLearning( num_feats, num_neurons, NaN, + NaN, folds, ) end @@ -285,6 +291,7 @@ julia> estimate_causal_effect!(m1) fit!(learner) its.causal_effect = predict(learner, its.X₁) .- its.Y₁ + its.marginal_effect = mean(its.causal_effect) return its.causal_effect end @@ -309,7 +316,9 @@ julia> estimate_causal_effect!(m1) ``` """ @inline function estimate_causal_effect!(g::GComputation) - g.causal_effect = mean(g_formula!(g)) + causal_effect, marginal_effect = g_formula!(g) + g.causal_effect, g.marginal_effect = mean(causal_effect), mean(marginal_effect) + return g.causal_effect end @@ -330,6 +339,7 @@ julia> g_formula!(m2) """ @inline function g_formula!(g) # Keeping this separate for reuse with S-Learning covariates, y = hcat(g.X, g.T), g.Y + x₁, x₀ = hcat(g.X, ones(size(g.X, 1))), hcat(g.X, zeros(size(g.X, 1))) if g.quantity_of_interest ∈ ("ITT", "ATE", "CATE") Xₜ = hcat(covariates[:, 1:(end - 1)], ones(size(covariates, 1))) @@ -350,10 +360,9 @@ julia> g_formula!(m2) ) fit!(g.ensemble) - yₜ, yᵤ = predict(g.ensemble, Xₜ), predict(g.ensemble, Xᵤ) - return vec(yₜ) - vec(yᵤ) + return vec(yₜ) - vec(yᵤ), predict(g.ensemble, x₁) - predict(g.ensemble, x₀) end """ @@ -374,27 +383,35 @@ julia> estimate_causal_effect!(m2) """ @inline function estimate_causal_effect!(DML::DoubleMachineLearning) X, T, Y = generate_folds(DML.X, DML.T, DML.Y, DML.folds) - DML.causal_effect = 0 + DML.causal_effect, DML.marginal_effect = 0, 0 + Δ = var_type(DML.T) isa Binary ? 1.0 : 1.5e-8mean(DML.T) # Cross fitting by training on the main folds and predicting residuals on the auxillary - for fld in 1:(DML.folds) - X_train, X_test = reduce(vcat, X[1:end .!== fld]), X[fld] - Y_train, Y_test = reduce(vcat, Y[1:end .!== fld]), Y[fld] - T_train, T_test = reduce(vcat, T[1:end .!== fld]), T[fld] - - Ỹ, T̃ = predict_residuals(DML, X_train, X_test, Y_train, Y_test, T_train, T_test) + for fold in 1:(DML.folds) + X_train, X_test = reduce(vcat, X[1:end .!== fold]), X[fold] + Y_train, Y_test = reduce(vcat, Y[1:end .!== fold]), Y[fold] + T_train, T_test = reduce(vcat, T[1:end .!== fold]), T[fold] + T_train₊ = var_type(DML.T) isa Binary ? T_train .* 0 : T_train .+ Δ + + Ỹ, T̃, T̃₊ = predict_residuals( + DML, X_train, X_test, Y_train, Y_test, T_train, T_test, T_train₊ + ) DML.causal_effect += T̃\Ỹ + DML.marginal_effect += (T̃₊\Ỹ - DML.causal_effect) / Δ end + DML.causal_effect /= DML.folds + DML.marginal_effect /= DML.folds return DML.causal_effect end """ - predict_residuals(D, x_train, x_test, y_train, y_test, t_train, t_test) + predict_residuals(D, x_train, x_test, y_train, y_test, t_train, t_test, t_train₊) -Predict treatment and outcome residuals for double machine learning or R-learning. +Predict treatment, outcome, and marginal effect residuals for double machine learning or +R-learning. # Notes This method should not be called directly. @@ -406,7 +423,7 @@ julia> x_train, x_test = X[1:80, :], X[81:end, :] julia> y_train, y_test = Y[1:80], Y[81:end] julia> t_train, t_test = T[1:80], T[81:100] julia> m1 = DoubleMachineLearning(X, T, Y) -julia> predict_residuals(m1, x_train, x_test, y_train, y_test, t_train, t_test) +julia> predict_residuals(m1, x_train, x_test, y_train, y_test, t_train, t_test, zeros(100)) ``` """ @inline function predict_residuals( @@ -417,6 +434,7 @@ julia> predict_residuals(m1, x_train, x_test, y_train, y_test, t_train, t_test) yₜₑ::Vector{Float64}, tₜᵣ::Vector{Float64}, tₜₑ::Vector{Float64}, + tₜᵣ₊::Vector{Float64} ) y = ELMEnsemble( xₜᵣ, yₜᵣ, D.sample_size, D.num_machines, D.num_feats, D.num_neurons, D.activation @@ -426,12 +444,17 @@ julia> predict_residuals(m1, x_train, x_test, y_train, y_test, t_train, t_test) xₜᵣ, tₜᵣ, D.sample_size, D.num_machines, D.num_feats, D.num_neurons, D.activation ) + t₊ = ELMEnsemble( + xₜᵣ, tₜᵣ₊, D.sample_size, D.num_machines, D.num_feats, D.num_neurons, D.activation + ) + fit!(y) fit!(t) + fit!(t₊) # Estimate a model with T + a finite difference - yₚᵣ, tₚᵣ = predict(y, xₜₑ), predict(t, xₜₑ) + yₚᵣ, tₚᵣ, tₚᵣ₊ = predict(y, xₜₑ), predict(t, xₜₑ), predict(t₊, xₜₑ) - return yₜₑ - yₚᵣ, tₜₑ - tₚᵣ + return yₜₑ - yₚᵣ, tₜₑ - tₚᵣ, tₜₑ - tₚᵣ₊ end """ diff --git a/src/inference.jl b/src/inference.jl index 28da1e3..71bbf5a 100644 --- a/src/inference.jl +++ b/src/inference.jl @@ -59,7 +59,8 @@ function summarize(mod; kwargs...) "Standard Error", "p-value", "Lower 2.5% CI", - "Upper 97.5% CI" + "Upper 97.5% CI", + "Marginal Effect" ] if haskey(kwargs, :inference) && kwargs[:inference] == true @@ -82,7 +83,8 @@ function summarize(mod; kwargs...) stderr, p, lower_ci, - upper_ci + upper_ci, + mod.marginal_effect ] for (nicename, value) in zip(nicenames, values) @@ -124,7 +126,8 @@ function summarize(its::InterruptedTimeSeries; kwargs...) "Standard Error", "p-value", "Lower 2.5% CI", - "Upper 97.5% CI" + "Upper 97.5% CI", + "Marginal Effect" ] values = [ @@ -140,7 +143,8 @@ function summarize(its::InterruptedTimeSeries; kwargs...) stderr, p, l, - u + u, + its.marginal_effect ] for (nicename, value) in zip(nicenames, values) diff --git a/src/metalearners.jl b/src/metalearners.jl index 7f2f355..559b15d 100644 --- a/src/metalearners.jl +++ b/src/metalearners.jl @@ -45,6 +45,7 @@ julia> m4 = SLearner(x_df, t_df, y_df) mutable struct SLearner <: Metalearner @standard_input_data @model_config individual_effect + marginal_effect::Vector{Float64} ensemble::ELMEnsemble function SLearner( @@ -76,6 +77,7 @@ mutable struct SLearner <: Metalearner num_feats, num_neurons, fill(NaN, size(T, 1)), + fill(NaN, size(T, 1)), ) end end @@ -123,6 +125,7 @@ julia> m3 = TLearner(x_df, t_df, y_df) mutable struct TLearner <: Metalearner @standard_input_data @model_config individual_effect + marginal_effect::Vector{Float64} μ₀::ELMEnsemble μ₁::ELMEnsemble @@ -154,6 +157,7 @@ mutable struct TLearner <: Metalearner num_feats, num_neurons, fill(NaN, size(T, 1)), + fill(NaN, size(T, 1)), ) end end @@ -201,6 +205,7 @@ julia> m3 = XLearner(x_df, t_df, y_df) mutable struct XLearner <: Metalearner @standard_input_data @model_config individual_effect + marginal_effect::Vector{Float64} μ₀::ELMEnsemble μ₁::ELMEnsemble ps::Array{Float64} @@ -233,6 +238,7 @@ mutable struct XLearner <: Metalearner num_feats, num_neurons, fill(NaN, size(T, 1)), + fill(NaN, size(T, 1)), ) end end @@ -278,6 +284,7 @@ julia> m2 = RLearner(x_df, t_df, y_df) mutable struct RLearner <: Metalearner @standard_input_data @model_config individual_effect + marginal_effect::Vector{Float64} folds::Integer end @@ -315,6 +322,7 @@ function RLearner( num_feats, num_neurons, fill(NaN, size(T, 1)), + fill(NaN, size(T, 1)), folds, ) end @@ -363,6 +371,7 @@ julia> m3 = DoublyRobustLearner(X, T, Y, W=w) mutable struct DoublyRobustLearner <: Metalearner @standard_input_data @model_config individual_effect + marginal_effect::Vector{Float64} folds::Integer end @@ -398,6 +407,7 @@ function DoublyRobustLearner( num_feats, num_neurons, fill(NaN, size(T, 1)), + fill(NaN, size(T, 1)), 2, ) end @@ -421,7 +431,7 @@ julia> estimate_causal_effect!(m4) ``` """ @inline function estimate_causal_effect!(s::SLearner) - s.causal_effect = g_formula!(s) + s.causal_effect, s.marginal_effect = g_formula!(s) return s.causal_effect end @@ -458,6 +468,7 @@ julia> estimate_causal_effect!(m5) fit!(t.μ₁) predictionsₜ, predictionsᵪ = predict(t.μ₁, t.X), predict(t.μ₀, t.X) t.causal_effect = @fastmath vec(predictionsₜ - predictionsᵪ) + t.marginal_effect = t.causal_effect return t.causal_effect end @@ -488,6 +499,8 @@ julia> estimate_causal_effect!(m1) (x.ps .* predict(μχ₀, x.X)) .+ ((1 .- x.ps) .* predict(μχ₁, x.X)) )) + x.marginal_effect = x.causal_effect # Works since T is binary + return x.causal_effect end @@ -510,6 +523,7 @@ julia> estimate_causal_effect!(m1) """ @inline function estimate_causal_effect!(R::RLearner) X, T̃, Ỹ = generate_folds(R.X, R.T, R.Y, R.folds) + T̃₊, Δ = similar(T̃), var_type(R.T) isa Binary ? 1.0 : 1.5e-8mean(R.T) R.X, R.T, R.Y = reduce(vcat, X), reduce(vcat, T̃), reduce(vcat, Ỹ) # Get residuals from out-of-fold predictions @@ -517,19 +531,26 @@ julia> estimate_causal_effect!(m1) X_train, X_test = reduce(vcat, X[1:end .!== f]), X[f] Y_train, Y_test = reduce(vcat, Ỹ[1:end .!== f]), Ỹ[f] T_train, T_test = reduce(vcat, T̃[1:end .!== f]), T̃[f] - Ỹ[f], T̃[f] = predict_residuals(R, X_train, X_test, Y_train, Y_test, T_train, T_test) + T_train₊ = var_type(R.T) isa Binary ? T_train .* 0 : T_train .+ Δ + Ỹ[f], T̃[f], T̃₊[f] = predict_residuals( + R, X_train, X_test, Y_train, Y_test, T_train, T_test, T_train₊ + ) end # Using target transformation and the weight trick to minimize the causal loss T̃², target = reduce(vcat, T̃).^2, reduce(vcat, Ỹ) ./ reduce(vcat, T̃) Xʷ, Yʷ = R.X .* T̃², target .* T̃² - - # Fit a weighted residual-on-residual model + T̃²₊, target₊ = reduce(vcat, T̃₊).^2, reduce(vcat, Ỹ) ./ reduce(vcat, T̃₊) final_model = ELMEnsemble( Xʷ, Yʷ, R.sample_size, R.num_machines, R.num_feats, R.num_neurons, R.activation ) - fit!(final_model) + + # Using finite differences to calculate marginal effects + final_model₊ = deepcopy(final_model) + final_model₊.X, final_model₊.Y = R.X .* T̃²₊, target₊ .* T̃²₊ + fit!(final_model); fit!(final_model₊) R.causal_effect = predict(final_model, R.X) + R.marginal_effect = (predict(final_model₊, final_model.X) - R.causal_effect) ./ Δ return R.causal_effect end @@ -563,6 +584,7 @@ julia> estimate_causal_effect!(m1) causal_effect ./= 2 DRE.causal_effect = causal_effect + DRE.marginal_effect = causal_effect return DRE.causal_effect end diff --git a/test/test_estimators.jl b/test/test_estimators.jl index 91e6efb..dddaa07 100644 --- a/test/test_estimators.jl +++ b/test/test_estimators.jl @@ -54,7 +54,7 @@ t_train, t_test = float(t[1:80]), float(t[81:end]) y_train, y_test = float(y[1:80]), float(y[81:end]) residual_predictor = DoubleMachineLearning(x, t, y, num_neurons=5) residuals = CausalELM.predict_residuals( - residual_predictor, x_train, x_test, y_train, y_test, t_train, t_test + residual_predictor, x_train, x_test, y_train, y_test, t_train, t_test, zeros(80) ) @testset "Interrupted Time Series Estimation" begin @@ -79,9 +79,11 @@ residuals = CausalELM.predict_residuals( @testset "Interrupted Time Series Estimation" begin @test isa(its.causal_effect, Array) + @test its.marginal_effect == CausalELM.mean(its.causal_effect) # Without autocorrelation @test isa(its_no_ar.causal_effect, Array) + @test its_no_ar.marginal_effect == CausalELM.mean(its_no_ar.causal_effect) end end @@ -105,7 +107,9 @@ end @testset "G-Computation Estimation" begin @test isa(g_computer.causal_effect, Float64) + @test isnan(g_computer.marginal_effect) == false @test isa(g_computer_binary_out.causal_effect, Float64) + @test isnan(g_computer.marginal_effect) == false # Check that the estimats for ATE and ATT are different @test g_computer.causal_effect !== gcomputer_att.causal_effect @@ -131,6 +135,7 @@ end @testset "Double Machine Learning Post-estimation Structure" begin @test dm.causal_effect isa Float64 + @test isnan(dm.marginal_effect) == false end end diff --git a/test/test_metalearners.jl b/test/test_metalearners.jl index b0cfcc0..95fd77e 100644 --- a/test/test_metalearners.jl +++ b/test/test_metalearners.jl @@ -82,7 +82,9 @@ estimate_causal_effect!(dr_learner_df) @testset "S-Learner Estimation" begin @test isa(slearner1.causal_effect, Array{Float64}) + @test all(isnan, slearner1.marginal_effect) == false @test isa(s_learner_binary.causal_effect, Array{Float64}) + @test all(isnan, s_learner_binary.marginal_effect) == false end end @@ -98,7 +100,9 @@ end @testset "T-Learner Estimation" begin @test isa(tlearner1.causal_effect, Array{Float64}) + @test all(isnan, tlearner1.marginal_effect) == false @test isa(t_learner_binary.causal_effect, Array{Float64}) + @test all(isnan, t_learner_binary.marginal_effect) == false end end @@ -133,7 +137,9 @@ end @test typeof(xlearner3.μ₁) <: CausalELM.ELMEnsemble @test xlearner3.ps isa Array{Float64} @test xlearner3.causal_effect isa Array{Float64} + @test all(isnan, xlearner3.marginal_effect) == false @test x_learner_binary.causal_effect isa Array{Float64} + @test all(isnan, x_learner_binary.marginal_effect) == false end end @@ -152,6 +158,7 @@ end @test length(rlearner.causal_effect) == length(y) @test eltype(rlearner.causal_effect) == Float64 @test all(isnan, rlearner.causal_effect) == false + @test all(isnan, rlearner.marginal_effect) == false end end @@ -175,6 +182,7 @@ end @test length(dr_learner.causal_effect) === length(y) @test eltype(dr_learner.causal_effect) == Float64 @test all(isnan, dr_learner.causal_effect) == false + @test all(isnan, dr_learner.marginal_effect) == false @test dr_learner_df.causal_effect isa Vector @test length(dr_learner_df.causal_effect) === length(y) @test eltype(dr_learner_df.causal_effect) == Float64 diff --git a/test/test_utilities.jl b/test/test_utilities.jl index 7129114..ed6137f 100644 --- a/test/test_utilities.jl +++ b/test/test_utilities.jl @@ -16,6 +16,7 @@ model_config_avg_ground_truth = quote num_feats::Integer num_neurons::Int64 causal_effect::Float64 + average_marginal_effect::Float64 end model_config_ind_ground_truth = quote @@ -29,6 +30,7 @@ model_config_ind_ground_truth = quote num_feats::Integer num_neurons::Int64 causal_effect::Array{Float64} + average_marginal_effect::Float64 end # Fields for the user supplied data