%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from pycausalml import (
CausalTree,
InteractionTree,
RLearner,
SLearner,
TLearner,
UpliftRandomForest,
UpliftXGBoost,
)
from pycausalml.tree_based_models import plot_causal_tree
from pycausalml.uplift_learners import (
format_uplift_tree,
plot_uplift_curve,
plot_uplift_tree,
qini_coefficient,
qini_curve,
)
plt.rcParams["figure.dpi"] = 110
def f_filter(X, y, w, names, order=1):
"""Rank features by the partial F statistic of treatment interactions.
Order 1 tests ``Y ~ X + W + X:W``. Higher orders add powers of ``X`` and
the matching treatment interactions, following the F-filter idea in
Zhao et al. (2020).
"""
y = np.asarray(y, dtype=float)
w = np.asarray(w, dtype=float)
rows = []
for j, name in enumerate(names):
xj = np.asarray(X[:, j], dtype=float)
scale = float(np.std(xj))
if scale < 1e-8:
rows.append({"feature": name, "score": np.nan})
continue
xj = (xj - float(np.mean(xj))) / scale
main = [np.ones(len(y)), w, xj]
for power in range(2, order + 1):
main.append(xj ** power)
n_main = len(main)
interactions = [w * (xj ** power) for power in range(1, order + 1)]
design = np.column_stack(main + interactions)
try:
full = sm.OLS(y, design).fit()
restricted = sm.OLS(y, design[:, :n_main]).fit()
q = design.shape[1] - n_main
score = ((restricted.ssr - full.ssr) / q) / (full.ssr / full.df_resid)
except Exception:
score = np.nan
rows.append({"feature": name, "score": float(score)})
return pd.DataFrame(rows).sort_values("score", ascending=False).reset_index(drop=True)
def plot_qini_models(y, treatment, scores, title):
fig, ax = plt.subplots(figsize=(6.8, 4.4))
table = []
for name, score in scores.items():
curve = qini_curve(y, treatment, score)
ax.plot(curve["fraction"], curve["qini"], label=name)
table.append({"model": name, "qini": qini_coefficient(y, treatment, score)})
ax.set_xlabel("Fraction targeted")
ax.set_ylabel("Cumulative incremental gain")
ax.set_title(title)
ax.legend(frameon=False)
fig.tight_layout()
return fig, pd.DataFrame(table)
print("Imported PyCausalML")
