DoubleML

The DoubleML scalar estimation logic is organized into a hierarchical layer system:

• DoubleMLBase (ABC): The foundation layer. It handles data storage and delegates result reporting (like summary, confint, bootstrap, p_adjust, and sensitivity_analysis) to a DoubleMLFramework object.
• DoubleMLScalar (ABC): The orchestration layer. It manages the lifecycle of a single-parameter estimation, including fit(), draw_sample_splitting(), fit_nuisance_models(), and learner management via set_learners().
• LinearScoreMixin: A specialized layer for models using linear scores. It implements closed-form parameter estimation ($\hat{\theta} = -E[\psi_b] / E[\psi_a]$) and standard error computation.
• Subclasses (e.g., PLR, IRM): The implementation layer. These define specific nuisance estimation logic (_nuisance_est()), required learner names (e.g., ml_l, ml_m), and score element computation (_get_score_elements()).

DoubleML uses an object-oriented design where specific model classes inherit from an abstract base class DoubleML.

Model Classes

The following classes implement the estimation of nuisance functions via machine learning and the computation of the Neyman orthogonal score function:

• DoubleMLPLR: Partially linear regression models (PLR)
• DoubleMLPLIV: Partially linear IV regression models (PLIV)
• DoubleMLIRM: Interactive regression models (IRM)
• DoubleIIVM: Interactive IV regression models (IIVM)

Base Class Functionality

The abstract base class DoubleML provides core functionalities for all models, including:

• fit: Estimate the double machine learning models.
• bootstrap: Perform bootstrapping for statistical inference.
• confint: Compute confidence intervals.
• p_adjust: Perform p-value adjustment.
• tune: Perform hyperparameter tuning.

This structure allows users to flexibly specify machine learners for nuisance functions, resampling schemes, the double machine learning algorithm, and Neyman orthogonal score functions.

Ensure the model raises appropriate errors for invalid inputs. Use pytest.raises(Error, match=r"regex pattern") to verify both the exception type and the error message.

Common Exceptions to Test:

• TypeError: Passing non-DoubleMLData objects.
• ValueError: Invalid score strings, n_folds < 2, or n_rep < 1.
• ValueError: Calling fit_nuisance_models() before draw_sample_splitting().
• ValueError: Calling estimate_causal_parameters() before fit_nuisance_models().
• ValueError: Calling fit() without setting required learners.
• TypeError: Passing a class instead of an instance to learners.

Model-Specific Examples:

• PLR: Check instrumental variables (z_cols) and ml_g warnings.
• IRM: Check binary treatment, instruments, and that ml_m is a classifier.

@pytest.mark.ci
def test_exception_data():
    msg = r"The data must be of DoubleMLData type\."
    with pytest.raises(TypeError, match=msg):
        <Model>(pd.DataFrame())

To use a scalar DoubleML model (like PLR), follow these five steps:

1. Initialize the model: Pass your DoubleMLData object and the desired score to the constructor.
2. Set learners: Use set_learners() to provide the machine learning models for nuisance parameter estimation. Note that learners are separated from the constructor.
3. Draw sample splitting: Call draw_sample_splitting(n_folds=..., n_rep=...) to define the cross-fitting strategy. This can be called independently of fitting.
4. Fit the model: Call fit() to execute the nuisance estimation and causal parameter estimation.
5. Retrieve results: Access results through methods inherited from DoubleMLBase, such as .summary, .confint(), or .bootstrap().

This workflow follows a template method pattern where the orchestration is handled by the base classes, while specific estimation logic is provided by subclasses.

# 1. Define model (data + score)
plr = PLR(obj_dml_data, score="partialling out")

# 2. Set learners
plr.set_learners(ml_l=RandomForestRegressor(), ml_m=RandomForestRegressor())

# 3. Draw sample splitting
plr.draw_sample_splitting(n_folds=5, n_rep=1)

# 4. Fit
plr.fit()

# 5. Results (delegated to DoubleMLFramework via DoubleMLBase)
print(plr.summary)
plr.confint()
plr.bootstrap()

Core estimation tests ensure the model produces statistically reasonable estimates. Use the following assertion patterns:

• When the true parameter matches the DGP theta: Use the 3-sigma rule: abs(coef - true_theta) <= 3.0 * se.
• When the true parameter is unknown (e.g., ATTE with heterogeneous effects): Check for finiteness and reasonable magnitude: np.isfinite(coef) and abs(coef) < 10.0.
• Standard Error check: Always verify se > 0.

Example fixture pattern for fitting a model for testing:

@pytest.fixture(scope="module")
def fitted_fixture(score, option):
    np.random.seed(3141)
    data = make_<model>_data(theta=true_theta, n_obs=500, ...)
    dml_obj = <Model>(data, score=score, option=option)
    dml_obj.set_learners(...)
    dml_obj.draw_sample_splitting(n_folds=5, n_rep=1)
    dml_obj.fit()
    return {"coef": dml_obj.coef[0], "se": dml_obj.se[0], "true_theta": true_theta, "score": score}

To ensure backward compatibility, verify that the new model produces numerically equivalent results to the legacy DoubleML implementation.

Crucial Step: You must share the sample splits (smpls) from the old model to the new model to ensure they consume the same random state, as they may handle initialization differently.

# Inside a comparison fixture
dml_new._smpls = dml_old.smpls
dml_new.fit()

Assertions: Use np.testing.assert_allclose(new, old, rtol=1e-9) for exact matches. Note that property names differ: the new implementation uses all_thetas/all_ses while the old one uses all_coef/all_se.

def test_coef_equal(comparison_fixture):
    np.testing.assert_allclose(new.coef, old.coef, rtol=1e-9)

def test_se_equal(comparison_fixture):
    np.testing.assert_allclose(new.se, old.se, rtol=1e-9)

def test_all_coef_equal(comparison_fixture):
    np.testing.assert_allclose(new.all_thetas, old.all_coef, rtol=1e-9)

def test_all_se_equal(comparison_fixture):
    np.testing.assert_allclose(new.all_ses, old.all_se, rtol=1e-9)

Verify that providing pre-computed predictions via the external_predictions argument in .fit() produces results equivalent to the reference model.

Workflow:

1. Fit a reference model.
2. Extract predictions from the reference model (e.g., dml_ref.predictions["ml_x"]).
3. Fit the new model using the same sample splits (_smpls) and passing the extracted predictions into fit(external_predictions=...).

Assertion Pattern: Because small numerical differences can accumulate when mixing external and fitted predictions, use math.isclose with a small absolute tolerance instead of np.testing.assert_allclose:

import math
assert math.isclose(ref.coef[0], ext.coef[0], rel_tol=1e-9, abs_tol=1e-4)

import math

def test_coef(ext_pred_fixture):
    assert math.isclose(ref.coef[0], ext.coef[0], rel_tol=1e-9, abs_tol=1e-4)

def test_se(ext_pred_fixture):
    assert math.isclose(ref.se[0], ext.se[0], rel_tol=1e-9, abs_tol=1e-4)