Partially Linear Regression (PLR) Tutorial

This tutorial demonstrates how to use the DoubleMLPLR model for estimating treatment effects in a partially linear regression framework.

Overview

The Partially Linear Regression model assumes:

$$Y = \theta D + g_0(X) + \epsilon$$

$$D = m_0(X) + v$$

Where:

$Y$ is the outcome variable
$D$ is the treatment variable (can be continuous or binary)
$X$ are control variables (covariates)
$\theta$ is the treatment effect we want to estimate
$g_0(X)$ and $m_0(X)$ are nuisance functions estimated via ML

Load packages and set up ML models

using DoubleML; using StableRNGs; using MLJ; using TreeParzen; using MLJDecisionTreeInterface; using EvoTrees

begin
    EvoTreeRegressor = @load EvoTreeRegressor pkg = EvoTrees verbosity = 0
    RandomForestRegressor = @load RandomForestRegressor pkg = DecisionTree verbosity = 0
end

MLJDecisionTreeInterface.RandomForestRegressor

Generate PLR data

# PLR Data
data_plr = DoubleML.make_plr_CCDDHNR2018(1000, alpha = 0.5, dim_x = 20, rng = StableRNG(42))

DoubleMLData{Float32, Vector{Float32}}(Float32[0.21417502, 0.9692105, -0.30895242, -0.048992738, 2.6029072, 1.4883567, 3.2605982, 1.0506742, -1.6848451, -0.40627423  …  -0.27870387, 2.7265549, 1.6948698, 2.7781668, -0.66583717, 0.78912425, 1.5714866, 0.21167372, 0.35720244, 1.2874476], Float32[-1.2644778, 1.2479526, -0.07056438, -1.4938519, 0.06656515, 1.2609595, 1.7690849, -0.95087236, -2.1373367, 0.64379275  …  -0.8578431, 0.8335334, 0.060642224, 1.8599323, -1.3742881, -1.7428911, 0.71109384, -0.05331813, -0.5296481, 1.9346514], Float32[-0.67025167 -0.14986733 … 0.41640848 -0.30865937; 2.085484 0.17391905 … -0.9364084 0.844609; … ; -0.5976384 -0.30607623 … 0.2812654 -0.3663869; 0.32168102 0.5612614 … 0.33550787 0.44790605], 1000, 20, :y, :d, [:X1, :X2, :X3, :X4, :X5, :X6, :X7, :X8, :X9, :X10, :X11, :X12, :X13, :X14, :X15, :X16, :X17, :X18, :X19, :X20])

We can check what models are available for predicting the outcome variable:

# Find matching models
models() do model
    matching(model, data_plr.x, data_plr.y)
end

12-element Vector{NamedTuple{(:name, :package_name, :is_supervised, :abstract_type, :constructor, :deep_properties, :docstring, :fit_data_scitype, :human_name, :hyperparameter_ranges, :hyperparameter_types, :hyperparameters, :implemented_methods, :inverse_transform_scitype, :is_pure_julia, :is_wrapper, :iteration_parameter, :load_path, :package_license, :package_url, :package_uuid, :predict_scitype, :prediction_type, :reporting_operations, :reports_feature_importances, :supports_class_weights, :supports_online, :supports_training_losses, :supports_weights, :tags, :target_in_fit, :transform_scitype, :input_scitype, :target_scitype, :output_scitype)}}:
 (name = CatBoostRegressor, package_name = CatBoost, ... )
 (name = DecisionTreeRegressor, package_name = BetaML, ... )
 (name = EvoTreeGaussian, package_name = EvoTrees, ... )
 (name = EvoTreeMLE, package_name = EvoTrees, ... )
 (name = EvoTreeRegressor, package_name = EvoTrees, ... )
 (name = GaussianMixtureRegressor, package_name = BetaML, ... )
 (name = NeuralNetworkRegressor, package_name = BetaML, ... )
 (name = NeuralNetworkRegressor, package_name = MLJFlux, ... )
 (name = PartLS, package_name = PartitionedLS, ... )
 (name = RandomForestRegressor, package_name = BetaML, ... )
 (name = SRRegressor, package_name = SymbolicRegression, ... )
 (name = SRTestRegressor, package_name = SymbolicRegression, ... )

Run a simple model

begin
    # Simple PLR with RandomForest
    ml_m = RandomForestRegressor(rng = StableRNG(42))
    ml_g = RandomForestRegressor(rng = StableRNG(42))

    dml_plr_simple = DoubleML.DoubleMLPLR(data_plr, ml_g, ml_m, n_folds = 4, n_rep = 1)

    fit!(dml_plr_simple)
end

DoubleMLPLR{Float32, MLJDecisionTreeInterface.RandomForestRegressor, MLJDecisionTreeInterface.RandomForestRegressor, Nothing}
==========================
StatsBase.CoefTable(Any[[0.5004154443740845], [0.030645174905657768], [16.32933807373047], [6.104304311646563e-60], [0.4403520052590645], [0.5604788834891045]], ["Estimate", "Std. Error", "z value", "Pr(>|z|)", "Lower 95.0%", "Upper 95.0%"], ["d"], 4, 3)

coeftable(dml_plr_simple)

────────────────────────────────────────────────────────────────────
   Estimate  Std. Error  z value  Pr(>|z|)  Lower 95.0%  Upper 95.0%
────────────────────────────────────────────────────────────────────
d  0.500415   0.0306452    16.33    <1e-59     0.440352     0.560479
────────────────────────────────────────────────────────────────────

Advanced example: self-tuning models

begin
    # PLR with TreeParzen hyperparameter tuning

    # Set up the hyperparameter space
    space = Dict(
        :n_trees => HP.Choice(:n_trees, Float64.(10:700)),
        :max_depth => HP.Choice(:max_depth, Float64.(1:10)),
        :min_samples_leaf => HP.Choice(:min_samples_leaf, Float64.(1:15)),
        :min_purity_increase => HP.Choice(:min_purity_increase, Float64.(0:3)),
        :sampling_fraction => HP.Choice(:sampling_fraction, Float64.(0.6:0.99)),
        :feature_importance => HP.Choice(:feature_importance, [:impurity, :split]),
    )

    # Set up the self-tuning models
    tuned_ml_m = TunedModel(
        model = RandomForestRegressor(rng = StableRNG(42)),
        tuning = MLJTreeParzenTuning(random_trials = 100, max_simultaneous_draws = 5, linear_forgetting = 50),
        resampling = CV(nfolds = 3),
        range = space,
        measure = MLJ.rmse,
        acceleration = CPUProcesses(),
    )

    tuned_ml_g = TunedModel(
        model = RandomForestRegressor(rng = StableRNG(42)),
        tuning = MLJTreeParzenTuning(random_trials = 100, max_simultaneous_draws = 5, linear_forgetting = 50),
        resampling = CV(nfolds = 3),
        range = space,
        measure = MLJ.rmse,
        acceleration = CPUProcesses(),
    )

    # Pass the self-tuning models as learners to the DoubleMLPLR constructor
    dml_plr = DoubleML.DoubleMLPLR(data_plr, tuned_ml_g, tuned_ml_m, n_folds = 4, n_rep = 1)

    # Fit it
    fit!(dml_plr, verbose = 0)
end

DoubleMLPLR{Float32, MLJTuning.DeterministicTunedModel{MLJTreeParzenTuning, MLJDecisionTreeInterface.RandomForestRegressor, Nothing}, MLJTuning.DeterministicTunedModel{MLJTreeParzenTuning, MLJDecisionTreeInterface.RandomForestRegressor, Nothing}, Nothing}
==========================
StatsBase.CoefTable(Any[[0.5085684657096863], [0.029936933889985085], [16.987995147705078], [1.007752415212665e-64], [0.4498931534777588], [0.5672437779416137]], ["Estimate", "Std. Error", "z value", "Pr(>|z|)", "Lower 95.0%", "Upper 95.0%"], ["d"], 4, 3)

coeftable(dml_plr)

────────────────────────────────────────────────────────────────────
   Estimate  Std. Error  z value  Pr(>|z|)  Lower 95.0%  Upper 95.0%
────────────────────────────────────────────────────────────────────
d  0.508568   0.0299369    16.99    <1e-63     0.449893     0.567244
────────────────────────────────────────────────────────────────────

Advanced example: iterated models

# A simple example
# EvoTrees have in-built early stopping as an option; the below is just for demonstration purposes.

begin
    # Set up iteration controls
    controls = [
        Step(1),
        Patience(10),
        NumberLimit(30),
    ]

    # Set up learners with iteration control and early stopping
    ml_l_iterated = IteratedModel(
        EvoTreeRegressor(seed = 42),
        resampling = Holdout(),
        measure = rmse,
        iteration_parameter = :nrounds,
        controls = controls
    )

    ml_m_iterated = IteratedModel(
        EvoTreeRegressor(seed = 42),
        resampling = Holdout(),
        measure = rmse,
        iteration_parameter = :nrounds,
        controls = controls
    )

    # Pass the learners to the DoulbleMLPLR contructor
    dml_plr_iterated = DoubleML.DoubleMLPLR(data_plr, ml_l_iterated, ml_m_iterated, n_folds = 4, n_rep = 1)

    # Fit it
    fit!(dml_plr_iterated, verbose = 0)
end

DoubleMLPLR{Float32, MLJIteration.DeterministicIteratedModel{EvoTrees.EvoTreeRegressor, Nothing}, MLJIteration.DeterministicIteratedModel{EvoTrees.EvoTreeRegressor, Nothing}, Nothing}
==========================
StatsBase.CoefTable(Any[[0.4677792489528656], [0.03386126086115837], [13.814584732055664], [2.08151106312002e-43], [0.40141239719387933], [0.5341461007118519]], ["Estimate", "Std. Error", "z value", "Pr(>|z|)", "Lower 95.0%", "Upper 95.0%"], ["d"], 4, 3)

coeftable(dml_plr_iterated)

────────────────────────────────────────────────────────────────────
   Estimate  Std. Error  z value  Pr(>|z|)  Lower 95.0%  Upper 95.0%
────────────────────────────────────────────────────────────────────
d  0.467779   0.0338613    13.81    <1e-42     0.401412     0.534146
────────────────────────────────────────────────────────────────────

summary(dml_plr_iterated)