Skip to content

Commit

Permalink
fix #717 tests KeyError['tau'] not in index
Browse files Browse the repository at this point in the history
  • Loading branch information
@rolandrmgservices
rolandrmgservices committed Nov 29, 2023
1 parent 3615bc8 commit 614cb2b
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 28 deletions.
4 changes: 1 addition & 3 deletions tests/test_cevae.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,7 @@ def test_CEVAE():
# check the accuracy of the ite accuracy
ite = cevae.predict(X).flatten()

auuc_metrics = pd.DataFrame(
{"ite": ite, "W": treatment, "y": y, "treatment_effect_col": tau}
)
auuc_metrics = pd.DataFrame({"ite": ite, "W": treatment, "y": y, "tau": tau})

cumgain = get_cumgain(
auuc_metrics, outcome_col="y", treatment_col="W", treatment_effect_col="tau"
Expand Down
2 changes: 1 addition & 1 deletion tests/test_ivlearner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def test_drivlearner():
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down
48 changes: 24 additions & 24 deletions tests/test_meta_learners.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def test_BaseSRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -179,7 +179,7 @@ def test_BaseTLearner(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -229,7 +229,7 @@ def test_BaseTRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -268,7 +268,7 @@ def test_MLPTRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -307,7 +307,7 @@ def test_XGBTRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -346,7 +346,7 @@ def test_BaseXLearner(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -398,7 +398,7 @@ def test_BaseXRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -439,7 +439,7 @@ def test_BaseXLearner_without_p(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -478,7 +478,7 @@ def test_BaseXRegressor_without_p(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -517,7 +517,7 @@ def test_BaseRLearner(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -568,7 +568,7 @@ def test_BaseRRegressor(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -607,7 +607,7 @@ def test_BaseRLearner_without_p(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -646,7 +646,7 @@ def test_BaseRRegressor_without_p(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down Expand Up @@ -698,15 +698,15 @@ def test_BaseSClassifier(generate_classification_data):
"tau_pred": tau_pred.flatten(),
"W": df_test["treatment_group_key"].values,
CONVERSION: df_test[CONVERSION].values,
"treatment_effect_col": df_test["treatment_effect"].values,
"tau": df_test["treatment_effect"].values,
}
)

cumgain = get_cumgain(
auuc_metrics,
outcome_col=CONVERSION,
treatment_col="W",
treatment_effect_col="treatment_effect_col",
treatment_effect_col="tau",
)

# Check if the cumulative gain when using the model's prediction is
Expand Down Expand Up @@ -742,15 +742,15 @@ def test_BaseTClassifier(generate_classification_data):
"tau_pred": tau_pred.flatten(),
"W": df_test["treatment_group_key"].values,
CONVERSION: df_test[CONVERSION].values,
"treatment_effect_col": df_test["treatment_effect"].values,
"tau": df_test["treatment_effect"].values,
}
)

cumgain = get_cumgain(
auuc_metrics,
outcome_col=CONVERSION,
treatment_col="W",
treatment_effect_col="treatment_effect_col",
treatment_effect_col="tau",
)

# Check if the cumulative gain when using the model's prediction is
Expand Down Expand Up @@ -812,15 +812,15 @@ def test_BaseXClassifier(generate_classification_data):
"tau_pred": tau_pred.flatten(),
"W": df_test["treatment_group_key"].values,
CONVERSION: df_test[CONVERSION].values,
"treatment_effect_col": df_test["treatment_effect"].values,
"tau": df_test["treatment_effect"].values,
}
)

cumgain = get_cumgain(
auuc_metrics,
outcome_col=CONVERSION,
treatment_col="W",
treatment_effect_col="treatment_effect_col",
treatment_effect_col="tau",
)

# Check if the cumulative gain when using the model's prediction is
Expand Down Expand Up @@ -861,15 +861,15 @@ def test_BaseRClassifier(generate_classification_data):
"tau_pred": tau_pred.flatten(),
"W": df_test["treatment_group_key"].values,
CONVERSION: df_test[CONVERSION].values,
"treatment_effect_col": df_test["treatment_effect"].values,
"tau": df_test["treatment_effect"].values,
}
)

cumgain = get_cumgain(
auuc_metrics,
outcome_col=CONVERSION,
treatment_col="W",
treatment_effect_col="treatment_effect_col",
treatment_effect_col="tau",
)

# Check if the cumulative gain when using the model's prediction is
Expand Down Expand Up @@ -912,15 +912,15 @@ def test_BaseRClassifier_with_sample_weights(generate_classification_data):
"tau_pred": tau_pred.flatten(),
"W": df_test["treatment_group_key"].values,
CONVERSION: df_test[CONVERSION].values,
"treatment_effect_col": df_test["treatment_effect"].values,
"tau": df_test["treatment_effect"].values,
}
)

cumgain = get_cumgain(
auuc_metrics,
outcome_col=CONVERSION,
treatment_col="W",
treatment_effect_col="treatment_effect_col",
treatment_effect_col="tau",
)

# Check if the cumulative gain when using the model's prediction is
Expand Down Expand Up @@ -1005,7 +1005,7 @@ def test_BaseDRLearner(generate_regression_data):
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
"tau": tau,
}
)

Expand Down

0 comments on commit 614cb2b

Please sign in to comment.