Add tested MakeTables integration (continuation of #383)

doubleml/double_ml.py

@@ -16,11 +16,12 @@
 from doubleml.utils._sensitivity import _compute_sensitivity_bias
 from doubleml.utils._tune_optuna import OPTUNA_GLOBAL_SETTING_KEYS, TUNE_ML_MODELS_DOC, resolve_optuna_cv
 from doubleml.utils.gain_statistics import gain_statistics
+from doubleml.utils.maketables_mixin import MakeTablesMixin
 
 _implemented_data_backends = ["DoubleMLData", "DoubleMLClusterData", "DoubleMLDIDData", "DoubleMLSSMData", "DoubleMLRDDData"]
 
 
-class DoubleML(SampleSplittingMixin, ABC):
+class DoubleML(MakeTablesMixin, SampleSplittingMixin, ABC):
     """Double Machine Learning."""
 
     def __init__(self, obj_dml_data, n_folds, n_rep, score, draw_sample_splitting, double_sample_splitting=False):

doubleml/plm/tests/test_lplr_tune_ml_models.py

@@ -1,93 +1,87 @@
+"""Validate LPLR tuning against its cross-validation objective and applied parameters."""
+
 import numpy as np
+import optuna
 import pytest
+from sklearn.base import clone
+from sklearn.model_selection import KFold, cross_val_score
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 
 import doubleml as dml
 from doubleml.plm.datasets import make_lplr_LZZ2020
-from doubleml.tests._utils_tune_optuna import (
-    _SAMPLER_CASES,
-    _assert_tree_params,
-    _basic_optuna_settings,
-    _small_tree_params,
-)
+from doubleml.tests._utils_tune_optuna import _assert_tree_params, _basic_optuna_settings, _small_tree_params
 
 
 @pytest.fixture(scope="module", params=["nuisance_space", "instrument"])
 def score(request):
     return request.param
 
 
-@pytest.fixture(
-    scope="module",
-    params=[DecisionTreeRegressor(random_state=567), None],
-)
+@pytest.fixture(scope="module", params=[DecisionTreeRegressor(random_state=567), None])
 def ml_a(request):
     return request.param
 
 
 @pytest.mark.ci
-@pytest.mark.parametrize("sampler_name,optuna_sampler", _SAMPLER_CASES, ids=[case[0] for case in _SAMPLER_CASES])
-def test_doubleml_lplr_optuna_tune(sampler_name, optuna_sampler, score, ml_a):
+@pytest.mark.parametrize("sampler_class", [optuna.samplers.RandomSampler, optuna.samplers.TPESampler], ids=["random", "tpe"])
+def test_doubleml_lplr_optuna_tune(sampler_class, score, ml_a):
+    """Tuning improves its own CV objective and applies the selected parameters."""
     np.random.seed(3141)
-    alpha = 0.5
-    dml_data = make_lplr_LZZ2020(n_obs=500, dim_x=15, alpha=alpha)
-
-    ml_M = DecisionTreeClassifier(random_state=123)
-    ml_t = DecisionTreeRegressor(random_state=234)
-    ml_m = DecisionTreeRegressor(random_state=456)
-
+    dml_data = make_lplr_LZZ2020(n_obs=500, dim_x=15, alpha=0.5)
     dml_lplr = dml.DoubleMLLPLR(
         dml_data,
-        ml_M=ml_M,
-        ml_t=ml_t,
-        ml_m=ml_m,
+        ml_M=DecisionTreeClassifier(random_state=123),
+        ml_t=DecisionTreeRegressor(random_state=234),
+        ml_m=DecisionTreeRegressor(random_state=456),
         ml_a=ml_a,
         n_folds=2,
         n_folds_inner=2,
         score=score,
     )
-    dml_lplr.fit()
-    untuned_score = dml_lplr.evaluate_learners()
-
-    optuna_params = {
-        "ml_M": _small_tree_params,
-        "ml_m": _small_tree_params,
-        "ml_t": _small_tree_params,
-        "ml_a": _small_tree_params,
+    cv = KFold(n_splits=5, shuffle=True, random_state=42)
+    x, y, d = dml_data.x, dml_data.y, dml_data.d
+    mask_m = y == 0 if score == "nuisance_space" else np.ones(len(y), dtype=bool)
+    tuning_data = {
+        "ml_M": (np.column_stack((d, x)), y, "neg_log_loss"),
+        "ml_m": (x[mask_m], d[mask_m], "neg_root_mean_squared_error"),
+        "ml_a": (x, d, "neg_root_mean_squared_error"),
     }
-
+    baseline_scores = {
+        name: cross_val_score(clone(dml_lplr.learner[name]), features, target, cv=cv, scoring=scorer).mean()
+        for name, (features, target, scorer) in tuning_data.items()
+    }
+    optuna_params = {name: _small_tree_params for name in ["ml_M", "ml_m", "ml_a", "ml_t"]}
     tune_res = dml_lplr.tune_ml_models(
         ml_param_space=optuna_params,
-        optuna_settings=_basic_optuna_settings({"sampler": optuna_sampler, "n_trials": 5}),
+        cv=cv,
+        optuna_settings=_basic_optuna_settings({"sampler": sampler_class(seed=3141), "n_trials": 5}),
         return_tune_res=True,
     )
 
-    dml_lplr.fit()
-    tuned_score = dml_lplr.evaluate_learners()
-
-    tuned_params_M = tune_res[0]["ml_M"].best_params
-    tuned_params_t = tune_res[0]["ml_t"].best_params
-    tuned_params_m = tune_res[0]["ml_m"].best_params
-    tuned_params_a = tune_res[0]["ml_a"].best_params
+    assert len(tune_res) == 1
+    assert set(tune_res[0]) == set(optuna_params)
+    for name, result in tune_res[0].items():
+        assert result.tuned
+        _assert_tree_params(result.best_params)
+        assert dml_lplr.get_params(name)["d"] == [[result.best_params] * dml_lplr.n_folds]
+        for key, value in result.best_params.items():
+            assert result.best_estimator.get_params()[key] == value
+
+    for name, (features, target, scorer) in tuning_data.items():
+        result = tune_res[0][name]
+        verified_score = cross_val_score(clone(result.best_estimator), features, target, cv=cv, scoring=scorer).mean()
+        assert result.scoring_method == scorer
+        assert result.best_score == pytest.approx(verified_score)
+        assert verified_score > baseline_scores[name]
 
-    _assert_tree_params(tuned_params_M)
-    _assert_tree_params(tuned_params_t)
-    _assert_tree_params(tuned_params_m)
-    _assert_tree_params(tuned_params_a)
-
-    # ensure results contain optuna objects and best params
-    assert isinstance(tune_res[0], dict)
-    assert set(tune_res[0].keys()) == {"ml_M", "ml_m", "ml_t", "ml_a"}
-    assert hasattr(tune_res[0]["ml_M"], "best_params")
-    assert tune_res[0]["ml_M"].best_params["max_depth"] == tuned_params_M["max_depth"]
-    assert hasattr(tune_res[0]["ml_t"], "best_params")
-    assert tune_res[0]["ml_t"].best_params["max_depth"] == tuned_params_t["max_depth"]
-    assert hasattr(tune_res[0]["ml_m"], "best_params")
-    assert tune_res[0]["ml_m"].best_params["max_depth"] == tuned_params_m["max_depth"]
-    assert hasattr(tune_res[0]["ml_a"], "best_params")
-    assert tune_res[0]["ml_a"].best_params["max_depth"] == tuned_params_a["max_depth"]
-
-    # ensure tuning improved RMSE #  not actually possible for ml_t as the targets are not available
-    assert tuned_score["ml_M"] < untuned_score["ml_M"]
-    assert tuned_score["ml_m"] < untuned_score["ml_m"]
-    assert tuned_score["ml_a"] < untuned_score["ml_a"]
+    dml_lplr.fit()
+    assert np.isfinite(dml_lplr.coef).all()
+    assert np.isfinite(dml_lplr.se).all()
+
+    expected_a = np.empty_like(d)
+    for (_, test), inner_splits in zip(dml_lplr.smpls[0], dml_lplr.smpls_inner[0]):
+        predictions = [
+            clone(tune_res[0]["ml_a"].best_estimator).fit(x[train], d[train]).predict(x[test]) for train, _ in inner_splits
+        ]
+        expected_a[test] = np.mean(predictions, axis=0)
+    np.testing.assert_allclose(dml_lplr.predictions["ml_a"][:, 0, 0], expected_a)

doubleml/tests/test_model_defaults.py

@@ -1,3 +1,5 @@
+"""Test model defaults independently of the execution order."""
+
 import numpy as np
 import pytest
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
@@ -224,6 +226,7 @@ def test_apos_defaults():
 
 @pytest.mark.ci
 def test_sensitivity_defaults():
+    dml_plr = dml.DoubleMLPLR(dml_data_plr, Lasso(), Lasso()).fit()
     input_dict = {"cf_y": 0.03, "cf_d": 0.03, "rho": 1.0, "level": 0.95, "null_hypothesis": np.array([0.0])}
 
     dml_plr.sensitivity_analysis()

doubleml/utils/maketables_mixin.py

@@ -0,0 +1,82 @@
+"""MakeTables plug-in support for models inheriting from DoubleML."""
+
+import numpy as np
+import pandas as pd
+
+
+class MakeTablesMixin:
+    """Expose fitted DoubleML results through the MakeTables plug-in protocol.
+
+    MakeTables detects these attributes without DoubleML importing it. Users
+    install ``maketables`` separately to create tables from fitted models.
+    """
+
+    @property
+    def __maketables_coef_table__(self) -> pd.DataFrame:
+        """Return estimates and pointwise inference for MakeTables.
+
+        Returns
+        -------
+        pd.DataFrame
+            Treatment-indexed estimates (``b``), standard errors (``se``),
+            t-statistics (``t``), p-values (``p``), and pointwise 95% confidence
+            interval bounds (``ci95l``, ``ci95u``). Confidence intervals use
+            DoubleML's aggregation across repeated sample splits.
+
+        Notes
+        -----
+        Before fitting or after resetting sample splits, returns an empty table
+        with the same columns, consistent with ``DoubleML.summary``.
+        """
+        if self.framework is None or np.isnan(self.coef).all():
+            return pd.DataFrame(columns=["b", "se", "t", "p", "ci95l", "ci95u"])
+
+        ci = self.confint(joint=False, level=0.95)
+        return pd.DataFrame(
+            {
+                "b": self.coef,
+                "se": self.se,
+                "t": self.t_stat,
+                "p": self.pval,
+                "ci95l": ci.iloc[:, 0],
+                "ci95u": ci.iloc[:, 1],
+            },
+            index=ci.index,
+        )
+
+    def __maketables_stat__(self, key: str) -> int | None:
+        """Return a model statistic for MakeTables.
+
+        Parameters
+        ----------
+        key : str
+            Statistic identifier. Only ``'N'`` (sample size) is supported.
+
+        Returns
+        -------
+        int or None
+            Number of observations, or ``None`` for an unsupported statistic.
+        """
+        return self.n_obs if key == "N" else None
+
+    @property
+    def __maketables_depvar__(self) -> str:
+        """Return the outcome variable name.
+
+        Returns
+        -------
+        str
+            Outcome column of the model's data object.
+        """
+        return self._dml_data.y_col
+
+    @property
+    def __maketables_default_stat_keys__(self) -> list[str]:
+        """Return the default model statistics for MakeTables.
+
+        Returns
+        -------
+        list[str]
+            ``['N']``. Users can override this through ``ETable(model_stats=...)``.
+        """
+        return ["N"]

doubleml/utils/tests/test_maketables_mixin.py

@@ -0,0 +1,130 @@
+"""Test MakeTables integration against fitted DoubleML estimates and rendered tables."""
+
+from copy import deepcopy
+
+import numpy as np
+import pandas as pd
+import pytest
+from maketables import ETable
+from maketables.extractors import get_extractor
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+import doubleml as dml
+
+
+@pytest.fixture(scope="module", params=["plr", "multi", "irm", "pliv", "cluster"])
+def model_kind(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def fitted_model(model_kind, n_rep):
+    np.random.seed(3141)
+    n_obs = 200
+    x = np.random.normal(size=(n_obs, 3))
+    z = np.random.normal(size=n_obs)
+    d = x[:, 0] + z + np.random.normal(size=n_obs)
+    if model_kind == "irm":
+        d = (d > 0).astype(float)
+    second = x[:, 1] + np.random.normal(size=n_obs)
+    y = 0.5 * d + 0.3 * second + x[:, 2] + np.random.normal(size=n_obs)
+    frame = pd.DataFrame(x, columns=["x1", "x2", "x3"])
+    frame["outcome"] = y
+    frame["treatment_z"] = d
+    frame["treatment_a"] = second
+    frame["instrument"] = z
+    d_cols = ["treatment_z", "treatment_a"] if model_kind == "multi" else ["treatment_z"]
+    kwargs = {"y_col": "outcome", "d_cols": d_cols, "x_cols": ["x1", "x2", "x3"]}
+    if model_kind == "cluster":
+        frame["cluster"] = np.repeat(np.arange(40), 5)
+        data = dml.DoubleMLData(frame, cluster_cols="cluster", **kwargs)
+    else:
+        data = dml.DoubleMLData(frame, z_cols="instrument" if model_kind == "pliv" else None, **kwargs)
+    if model_kind == "irm":
+        model = dml.DoubleMLIRM(data, LinearRegression(), LogisticRegression(), n_folds=2, n_rep=n_rep)
+    elif model_kind == "pliv":
+        model = dml.DoubleMLPLIV(data, LinearRegression(), LinearRegression(), LinearRegression(), n_folds=2, n_rep=n_rep)
+    else:
+        model = dml.DoubleMLPLR(data, LinearRegression(), LinearRegression(), n_folds=2, n_rep=n_rep)
+    return model.fit()
+
+
+@pytest.mark.ci
+def test_maketables_estimates(fitted_model):
+    """The detected extractor preserves estimates, inference, and treatment order."""
+    model = fitted_model
+    before = model.summary.copy(deep=True)
+    table = get_extractor(model).coef_table(model)
+    expected = before.rename(columns={"coef": "b", "std err": "se", "P>|t|": "p", "2.5 %": "ci95l", "97.5 %": "ci95u"})
+    expected.index.name = "Coefficient"
+    pd.testing.assert_frame_equal(table[expected.columns], expected)
+    pd.testing.assert_frame_equal(model.summary, before)
+    table.iloc[0, 0] = 999
+    pd.testing.assert_frame_equal(model.summary, before)
+
+
+@pytest.mark.ci
+def test_maketables_metadata(fitted_model):
+    """MakeTables receives the outcome name, sample size, and supported defaults."""
+    extractor = get_extractor(fitted_model)
+    assert extractor.depvar(fitted_model) == "outcome"
+    assert extractor.stat(fitted_model, "N") == 200
+    assert extractor.default_stat_keys(fitted_model) == ["N"]
+    assert extractor.stat(fitted_model, "r2") is None
+    assert extractor.stat(fitted_model, "unknown") is None
+
+
+@pytest.mark.ci
+@pytest.mark.parametrize("output_format", ["tex", "html"])
+def test_maketables_render(fitted_model, output_format):
+    """Actual exported tables contain estimates and pointwise confidence intervals."""
+    model = fitted_model
+    table = ETable(model, coef_fmt="b:.6f (se:.6f) [ci95l:.6f, ci95u:.6f]", labels={"treatment_z": "Treatment"})
+    output = table.make(output_format)
+    assert "Treatment" in output
+    assert "outcome" in output
+    assert "200" in output
+    for value in [model.coef[0], model.se[0], *model.confint().iloc[0]]:
+        assert f"{value:.6f}" in output
+
+
+@pytest.mark.ci
+def test_maketables_multiple_models(fitted_model):
+    """The same estimate is rendered in each requested model column."""
+    output = ETable([fitted_model, fitted_model], coef_fmt="b:.6f").make("tex")
+    assert output.count(f"{fitted_model.coef[0]:.6f}") == 2
+
+
+@pytest.mark.ci
+def test_maketables_before_fit():
+    """An unfitted model produces no estimates, matching DoubleML.summary."""
+    np.random.seed(3141)
+    data = dml.DoubleMLData.from_arrays(np.random.normal(size=(20, 2)), np.arange(20.0), np.arange(20.0))
+    model = dml.DoubleMLPLR(data, LinearRegression(), LinearRegression(), n_folds=2)
+    table = model.__maketables_coef_table__
+    assert table.empty
+    assert list(table.columns) == ["b", "se", "t", "p", "ci95l", "ci95u"]
+    output = ETable(model).make("tex")
+    assert "20" in output
+    assert "nan" not in output.lower()
+    assert model.framework is None
+
+
+@pytest.mark.ci
+def test_maketables_reset_and_refit(fitted_model):
+    """Resetting sample splits must hide old intervals until the model is refitted."""
+    model = deepcopy(fitted_model)
+    before = model.__maketables_coef_table__
+    model.set_sample_splitting(model.smpls, model.smpls_cluster)
+    assert model.summary.empty
+    assert model.__maketables_coef_table__.empty
+    output = ETable(model, coef_fmt="b:.6f [ci95l:.6f, ci95u:.6f]").make("tex")
+    for value in before["ci95l"]:
+        assert f"{value:.6f}" not in output
+    model.fit()
+    pd.testing.assert_frame_equal(model.__maketables_coef_table__, before)

pyproject.toml

@@ -52,6 +52,7 @@ dev = [
     "ruff>=0.11.1",
     "pre-commit>=4.2.0",
     "mypy>=1.18.0",
+    "maketables>=0.2.0",
 ]
 
 [project.urls]