Skip to content

Commit 9575229

Browse files
SvenKlaassenSvenKlaassen
authored
Merge pull request #368 from DoubleML/j-optuna
Hyperparameter Tuning using Optuna
2 parents 0c04ef6 + 80c731f commit 9575229

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+4614
-37
lines changed

‎.github/workflows/pytest.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ jobs:
5858
matrix.config.os != 'ubuntu-latest' ||
5959
matrix.config.python-version != '3.12'
6060
run: |
61+
pytest --doctest-modules --ignore-glob="doubleml/**/tests/*" --ignore-glob="doubleml/tests/*"
6162
pytest -m ci
6263
pytest -m ci_rdd
6364

‎doubleml/data/tests/test_dml_data.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -569,6 +569,7 @@ def test_dml_data_w_missings(generate_data_irm_w_missings):
569569
assert dml_data.force_all_x_finite == "allow-nan"
570570

571571

572+
@pytest.mark.ci
572573
def test_dml_data_w_missing_d(generate_data1):
573574
data = generate_data1
574575
np.random.seed(3141)

‎doubleml/did/datasets/dgp_did_CS2021.py‎

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,11 +105,35 @@ def make_did_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, time_typ
105105
106106
6. Treatment assignment:
107107
108-
For non-experimental settings (DGP 1-4), the probability of being in treatment group :math:`g` is:
108+
For non-experimental settings (DGP 1-4), the probability of being in treatment group :math:`g` is computed as follows:
109109
110-
.. math::
110+
- Compute group-specific logits for each observation:
111+
112+
.. math::
113+
114+
\\text{logit}_{i,g} = f_{ps,g}(W_{ps})
115+
116+
The logits are clipped to the range [-2.5, 2.5] for numerical stability.
117+
118+
- Convert logits to uncapped probabilities via softmax:
119+
120+
.. math::
121+
122+
p^{\\text{uncapped}}_{i,g} = \\frac{\\exp(\\text{logit}_{i,g})}{\\sum_{g'} \\exp(\\text{logit}_{i,g'})}
123+
124+
- Clip uncapped probabilities to the range [0.05, 0.95]:
125+
126+
.. math::
127+
128+
p^{\\text{clipped}}_{i,g} = \\min(\\max(p^{\\text{uncapped}}_{i,g}, 0.05), 0.95)
129+
130+
- Renormalize clipped probabilities so they sum to 1 for each observation:
131+
132+
.. math::
133+
134+
p_{i,g} = \\frac{p^{\text{clipped}}_{i,g}}{\\sum_{g'} p^{\\text{clipped}}_{i,g'}}
111135
112-
P(G_i = g) = \\frac{\\exp(f_{ps,g}(W_{ps}))}{\\sum_{g'} \\exp(f_{ps,g'}(W_{ps}))}
136+
- Assign each observation to a treatment group by sampling from the categorical distribution defined by :math:`p_{i,g}`.
113137
114138
For experimental settings (DGP 5-6), each treatment group (including never-treated) has equal probability:
115139
@@ -159,7 +183,7 @@ def make_did_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, time_typ
159183
`dim_x` (int, default=4):
160184
Dimension of feature vectors.
161185
162-
`xi` (float, default=0.9):
186+
`xi` (float, default=0.5):
163187
Scale parameter for the propensity score function.
164188
165189
`n_periods` (int, default=5):
@@ -188,7 +212,7 @@ def make_did_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, time_typ
188212

189213
c = kwargs.get("c", 0.0)
190214
dim_x = kwargs.get("dim_x", 4)
191-
xi = kwargs.get("xi", 0.9)
215+
xi = kwargs.get("xi", 0.75)
192216
n_periods = kwargs.get("n_periods", 5)
193217
anticipation_periods = kwargs.get("anticipation_periods", 0)
194218
n_pre_treat_periods = kwargs.get("n_pre_treat_periods", 2)
@@ -228,8 +252,11 @@ def make_did_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, time_typ
228252
p = np.ones(n_treatment_groups) / n_treatment_groups
229253
d_index = np.random.choice(n_treatment_groups, size=n_obs, p=p)
230254
else:
231-
unnormalized_p = np.exp(_f_ps_groups(features_ps, xi, n_groups=n_treatment_groups))
232-
p = unnormalized_p / unnormalized_p.sum(1, keepdims=True)
255+
logits = np.clip(_f_ps_groups(features_ps, xi, n_groups=n_treatment_groups), a_min=-2.5, a_max=2.5)
256+
unnormalized_p = np.exp(logits)
257+
p_uncapped = unnormalized_p / unnormalized_p.sum(1, keepdims=True)
258+
p_clipped = np.clip(p_uncapped, a_min=0.05, a_max=0.95)
259+
p = p_clipped / p_clipped.sum(1, keepdims=True)
233260
d_index = np.array([np.random.choice(n_treatment_groups, p=p_row) for p_row in p])
234261

235262
# fixed effects (shape (n_obs, n_time_periods))

‎doubleml/did/datasets/dgp_did_cs_CS2021.py‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,11 +85,35 @@ def make_did_cs_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, lambd
8585
8686
6. Treatment assignment:
8787
88-
For non-experimental settings (DGP 1-4), the probability of being in treatment group :math:`g` is:
88+
For non-experimental settings (DGP 1-4), the probability of being in treatment group :math:`g` is computed as follows:
8989
90-
.. math::
90+
- Compute group-specific logits for each observation:
91+
92+
.. math::
93+
94+
\\text{logit}_{i,g} = f_{ps,g}(W_{ps})
95+
96+
The logits are clipped to the range [-2.5, 2.5] for numerical stability.
97+
98+
- Convert logits to uncapped probabilities via softmax:
99+
100+
.. math::
101+
102+
p^{\\text{uncapped}}_{i,g} = \\frac{\\exp(\\text{logit}_{i,g})}{\\sum_{g'} \\exp(\\text{logit}_{i,g'})}
103+
104+
- Clip uncapped probabilities to the range [0.05, 0.95]:
105+
106+
.. math::
107+
108+
p^{\\text{clipped}}_{i,g} = \\min(\\max(p^{\\text{uncapped}}_{i,g}, 0.05), 0.95)
109+
110+
- Renormalize clipped probabilities so they sum to 1 for each observation:
111+
112+
.. math::
113+
114+
p_{i,g} = \\frac{p^{\text{clipped}}_{i,g}}{\\sum_{g'} p^{\\text{clipped}}_{i,g'}}
91115
92-
P(G_i = g) = \\frac{\\exp(f_{ps,g}(W_{ps}))}{\\sum_{g'} \\exp(f_{ps,g'}(W_{ps}))}
116+
- Assign each observation to a treatment group by sampling from the categorical distribution defined by :math:`p_{i,g}`.
93117
94118
For experimental settings (DGP 5-6), each treatment group (including never-treated) has equal probability:
95119
@@ -148,7 +172,7 @@ def make_did_cs_CS2021(n_obs=1000, dgp_type=1, include_never_treated=True, lambd
148172
`dim_x` (int, default=4):
149173
Dimension of feature vectors.
150174
151-
`xi` (float, default=0.9):
175+
`xi` (float, default=0.5):
152176
Scale parameter for the propensity score function.
153177
154178
`n_periods` (int, default=5):

‎doubleml/did/did.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from doubleml.double_ml_score_mixins import LinearScoreMixin
1010
from doubleml.utils._checks import _check_finite_predictions, _check_is_propensity, _check_score
1111
from doubleml.utils._estimation import _dml_cv_predict, _dml_tune, _get_cond_smpls
12+
from doubleml.utils._tune_optuna import _dml_tune_optuna
1213

1314

1415
# TODO: Remove DoubleMLDIDData with version 0.12.0
@@ -427,6 +428,83 @@ def _nuisance_tuning(
427428

428429
return res
429430

431+
def _nuisance_tuning_optuna(
432+
self,
433+
optuna_params,
434+
scoring_methods,
435+
cv,
436+
optuna_settings,
437+
):
438+
"""
439+
Optuna-based hyperparameter tuning for DID nuisance models.
440+
441+
Performs tuning once on the whole dataset using cross-validation,
442+
returning the same optimal parameters for all folds.
443+
"""
444+
445+
x, y = check_X_y(self._dml_data.x, self._dml_data.y, ensure_all_finite=False)
446+
x, d = check_X_y(x, self._dml_data.d, ensure_all_finite=False)
447+
448+
if scoring_methods is None:
449+
if self.score == "observational":
450+
scoring_methods = {"ml_g0": None, "ml_g1": None, "ml_m": None}
451+
else:
452+
scoring_methods = {"ml_g0": None, "ml_g1": None}
453+
454+
# Separate data by treatment status for conditional mean tuning
455+
mask_d0 = d == 0
456+
mask_d1 = d == 1
457+
458+
x_d0 = x[mask_d0, :]
459+
y_d0 = y[mask_d0]
460+
g0_tune_res = _dml_tune_optuna(
461+
y_d0,
462+
x_d0,
463+
self._learner["ml_g"],
464+
optuna_params["ml_g0"],
465+
scoring_methods["ml_g0"],
466+
cv,
467+
optuna_settings,
468+
learner_name="ml_g",
469+
params_name="ml_g0",
470+
)
471+
472+
x_d1 = x[mask_d1, :]
473+
y_d1 = y[mask_d1]
474+
g1_tune_res = _dml_tune_optuna(
475+
y_d1,
476+
x_d1,
477+
self._learner["ml_g"],
478+
optuna_params["ml_g1"],
479+
scoring_methods["ml_g1"],
480+
cv,
481+
optuna_settings,
482+
learner_name="ml_g",
483+
params_name="ml_g1",
484+
)
485+
486+
# Tune propensity score on full dataset for observational score
487+
m_tune_res = None
488+
if self.score == "observational":
489+
m_tune_res = _dml_tune_optuna(
490+
d,
491+
x,
492+
self._learner["ml_m"],
493+
optuna_params["ml_m"],
494+
scoring_methods["ml_m"],
495+
cv,
496+
optuna_settings,
497+
learner_name="ml_m",
498+
params_name="ml_m",
499+
)
500+
501+
if self.score == "observational":
502+
results = {"ml_g0": g0_tune_res, "ml_g1": g1_tune_res, "ml_m": m_tune_res}
503+
else:
504+
results = {"ml_g0": g0_tune_res, "ml_g1": g1_tune_res}
505+
506+
return results
507+
430508
def sensitivity_benchmark(self, benchmarking_set, fit_args=None):
431509
"""
432510
Computes a benchmark for a given set of features.

‎doubleml/did/did_binary.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
_check_score,
2424
)
2525
from doubleml.utils._estimation import _dml_cv_predict, _dml_tune, _get_cond_smpls
26+
from doubleml.utils._tune_optuna import _dml_tune_optuna
2627
from doubleml.utils.propensity_score_processing import PSProcessorConfig, init_ps_processor
2728

2829

@@ -666,6 +667,79 @@ def _nuisance_tuning(
666667

667668
return res
668669

670+
def _nuisance_tuning_optuna(
671+
self,
672+
optuna_params,
673+
scoring_methods,
674+
cv,
675+
optuna_settings,
676+
):
677+
678+
x, y = check_X_y(self._x_data_subset, self._y_data_subset, ensure_all_finite=False)
679+
x, d = check_X_y(x, self._g_data_subset, ensure_all_finite=False)
680+
681+
if scoring_methods is None:
682+
if self.score == "observational":
683+
scoring_methods = {"ml_g0": None, "ml_g1": None, "ml_m": None}
684+
else:
685+
scoring_methods = {"ml_g0": None, "ml_g1": None}
686+
687+
mask_d0 = d == 0
688+
mask_d1 = d == 1
689+
690+
x_d0 = x[mask_d0, :]
691+
y_d0 = y[mask_d0]
692+
g0_param_grid = optuna_params["ml_g0"]
693+
g0_scoring = scoring_methods["ml_g0"]
694+
g0_tune_res = _dml_tune_optuna(
695+
y_d0,
696+
x_d0,
697+
self._learner["ml_g"],
698+
g0_param_grid,
699+
g0_scoring,
700+
cv,
701+
optuna_settings,
702+
learner_name="ml_g",
703+
params_name="ml_g0",
704+
)
705+
706+
x_d1 = x[mask_d1, :]
707+
y_d1 = y[mask_d1]
708+
g1_param_grid = optuna_params["ml_g1"]
709+
g1_scoring = scoring_methods["ml_g1"]
710+
g1_tune_res = _dml_tune_optuna(
711+
y_d1,
712+
x_d1,
713+
self._learner["ml_g"],
714+
g1_param_grid,
715+
g1_scoring,
716+
cv,
717+
optuna_settings,
718+
learner_name="ml_g",
719+
params_name="ml_g1",
720+
)
721+
722+
m_tune_res = None
723+
if self.score == "observational":
724+
m_tune_res = _dml_tune_optuna(
725+
d,
726+
x,
727+
self._learner["ml_m"],
728+
optuna_params["ml_m"],
729+
scoring_methods["ml_m"],
730+
cv,
731+
optuna_settings,
732+
learner_name="ml_m",
733+
params_name="ml_m",
734+
)
735+
736+
if self.score == "observational":
737+
results = {"ml_g0": g0_tune_res, "ml_g1": g1_tune_res, "ml_m": m_tune_res}
738+
else:
739+
results = {"ml_g0": g0_tune_res, "ml_g1": g1_tune_res}
740+
741+
return results
742+
669743
def _sensitivity_element_est(self, preds):
670744
y = self._y_data_subset
671745
d = self._g_data_subset

0 commit comments

Comments
 (0)