# Define reusable helpers for the Helper Functions section.
def save_table(df, name):
"""
Idea: Save a result table to the notebook table directory and return the same DataFrame for display.
Parameters
----------
df : pd.DataFrame
Rows used by this helper.
name : str
Short name used for the scenario label.
Returns
-------
pd.DataFrame
Same table after writing it to the tutorial table directory.
"""
path = TABLE_DIR / f"{NOTEBOOK_PREFIX}_{name}.csv"
df.to_csv(path, index=False)
return df
def sigmoid(x):
"""
Idea: Map a linear score to a probability on the 0 to 1 scale for treatment assignment or simulation logic.
Parameters
----------
x : str or array-like
First variable, column name, or numeric input in the comparison.
Returns
-------
np.ndarray
Probability-scale values obtained by applying the logistic transform elementwise.
"""
return 1.0 / (1.0 + np.exp(-x))
def rmse(y_true, y_pred):
"""
Idea: Compute root mean squared error between an oracle target and an estimated or predicted value.
Parameters
----------
y_true : array-like
Reference outcome, oracle value, or known target used for evaluation.
y_pred : array-like
Predicted outcome or estimated value being evaluated.
Returns
-------
float
Root mean squared error between the reference values and predictions.
"""
return float(np.sqrt(mean_squared_error(y_true, y_pred)))
def mean_difference_summary(y, group, label, target="difference"):
"""
Idea: Assemble a summary of the mean difference summary with the quantities needed for interpretation.
Parameters
----------
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
group : object
Group or segment label being summarized.
label : str
Short label attached to a scenario, method, or plotted result.
target : str or array-like
Outcome, metric, or variable whose transformed value is needed.
Returns
-------
dict[str, float]
Treated-control mean-difference summary with estimate, standard error, interval, and target label.
"""
group = np.asarray(group)
y = np.asarray(y)
high = y[group == 1]
low = y[group == 0]
theta_hat = float(high.mean() - low.mean())
std_error = float(np.sqrt(high.var(ddof=1) / len(high) + low.var(ddof=1) / len(low)))
return {
"estimator": label,
"target": target,
"theta_hat": theta_hat,
"std_error": std_error,
"ci_95_lower": theta_hat - 1.96 * std_error,
"ci_95_upper": theta_hat + 1.96 * std_error,
"p_value": np.nan,
}
def treatment_ols_summary(y, X, treatment_col, label):
"""
Idea: Assemble a summary of the treatment OLS summary with the quantities needed for interpretation.
Parameters
----------
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
X : pd.DataFrame or np.ndarray
Feature matrix containing pre-treatment covariates or effect modifiers.
treatment_col : str
Name of the treatment or intervention column.
label : str
Short label attached to a scenario, method, or plotted result.
Returns
-------
dict[str, float]
Treatment-slope summary row from the fitted OLS model.
"""
X_design = sm.add_constant(X, has_constant="add")
fit = sm.OLS(y, X_design).fit(cov_type="HC1")
row = fit.summary2().tables[1].loc[treatment_col]
return {
"estimator": label,
"target": "observed-control treatment slope",
"theta_hat": float(row["Coef."]),
"std_error": float(row["Std.Err."]),
"ci_95_lower": float(row["[0.025"]),
"ci_95_upper": float(row["0.975]"]),
"p_value": float(row["P>|z|"]),
}
def wald_summary(y, d, z, label):
"""
Idea: Assemble the Wald ratio estimate, first-stage strength, and reduced-form evidence for an IV design.
Parameters
----------
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
d : object
Treatment or decision indicator used by the causal score.
z : str or array-like
Instrument, conditioning variable, or third input used by the model.
label : str
Short label attached to a scenario, method, or plotted result.
Returns
-------
dict[str, float]
Wald-style IV summary row with numerator, denominator, estimate, and uncertainty fields.
"""
y = np.asarray(y)
d = np.asarray(d)
z = np.asarray(z)
numerator = y[z == 1].mean() - y[z == 0].mean()
denominator = d[z == 1].mean() - d[z == 0].mean()
theta_hat = float(numerator / denominator)
influence = (z / z.mean()) * (y - theta_hat * d) - ((1 - z) / (1 - z).mean()) * (y - theta_hat * d)
std_error = float(np.std(influence, ddof=1) / (abs(denominator) * np.sqrt(len(y))))
return {
"estimator": label,
"target": "raw Wald LATE-style ratio",
"theta_hat": theta_hat,
"std_error": std_error,
"ci_95_lower": theta_hat - 1.96 * std_error,
"ci_95_upper": theta_hat + 1.96 * std_error,
"p_value": np.nan,
"wald_numerator": float(numerator),
"wald_denominator": float(denominator),
}
def residualized_wald_summary(y, d, z, X, label):
"""
Idea: Compute a Wald-style IV summary after residualizing treatment, outcome, and instrument on controls.
Parameters
----------
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
d : object
Treatment or decision indicator used by the causal score.
z : str or array-like
Instrument, conditioning variable, or third input used by the model.
X : pd.DataFrame or np.ndarray
Feature matrix containing pre-treatment covariates or effect modifiers.
label : str
Short label attached to a scenario, method, or plotted result.
Returns
-------
dict[str, float]
Residualized IV summary row after adjusting outcome, treatment, and instrument for controls.
"""
y_resid = y - LinearRegression().fit(X, y).predict(X)
d_resid = d - LinearRegression().fit(X, d).predict(X)
z_resid = z - LinearRegression().fit(X, z).predict(X)
numerator = float(np.mean(z_resid * y_resid))
denominator = float(np.mean(z_resid * d_resid))
theta_hat = numerator / denominator
psi = z_resid * (y_resid - theta_hat * d_resid)
std_error = float(np.sqrt(np.mean(psi ** 2) / (denominator ** 2 * len(y))))
return {
"estimator": label,
"target": "linear residualized Wald ratio",
"theta_hat": float(theta_hat),
"std_error": std_error,
"ci_95_lower": theta_hat - 1.96 * std_error,
"ci_95_upper": theta_hat + 1.96 * std_error,
"p_value": np.nan,
"wald_numerator": numerator,
"wald_denominator": denominator,
}
def iivm_ratio_summary(y, d, z, g0_hat, g1_hat, m_hat, r0_hat, r1_hat, label):
"""
Idea: Assemble a summary of the interactive IV model ratio summary with the quantities needed for interpretation.
Parameters
----------
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
d : object
Treatment or decision indicator used by the causal score.
z : str or array-like
Instrument, conditioning variable, or third input used by the model.
g0_hat : object
Estimated control-outcome regression prediction.
g1_hat : object
Estimated treated-outcome regression prediction.
m_hat : object
Estimated treatment or propensity nuisance prediction used by the orthogonal score.
r0_hat : object
Estimated nuisance prediction for the control or untreated state.
r1_hat : object
Estimated nuisance prediction for the treated state.
label : str
Short label attached to a scenario, method, or plotted result.
Returns
-------
dict[str, float]
IIVM Wald-ratio summary with numerator, denominator, estimate, and interval.
"""
y = np.asarray(y)
d = np.asarray(d)
z = np.asarray(z)
m_hat = np.clip(np.asarray(m_hat), 0.02, 0.98)
g0_hat = np.asarray(g0_hat)
g1_hat = np.asarray(g1_hat)
r0_hat = np.clip(np.asarray(r0_hat), 0.01, 0.99)
r1_hat = np.clip(np.asarray(r1_hat), 0.01, 0.99)
numerator_score = g1_hat - g0_hat + z * (y - g1_hat) / m_hat - (1 - z) * (y - g0_hat) / (1 - m_hat)
denominator_score = r1_hat - r0_hat + z * (d - r1_hat) / m_hat - (1 - z) * (d - r0_hat) / (1 - m_hat)
numerator = float(np.mean(numerator_score))
denominator = float(np.mean(denominator_score))
theta_hat = numerator / denominator
ratio_score = numerator_score - theta_hat * denominator_score
std_error = float(np.std(ratio_score, ddof=1) / (abs(denominator) * np.sqrt(len(y))))
return {
"estimator": label,
"target": "LATE",
"theta_hat": float(theta_hat),
"std_error": std_error,
"ci_95_lower": theta_hat - 1.96 * std_error,
"ci_95_upper": theta_hat + 1.96 * std_error,
"p_value": np.nan,
"iivm_numerator": numerator,
"iivm_denominator": denominator,
}
def crossfit_iivm_nuisances(X, y, d, z, outcome_learner, instrument_learner, treatment_learner, n_splits=5, random_state=RANDOM_STATE):
"""
Idea: Cross-fit the nuisance models for the crossfit interactive IV model nuisances design so evaluation uses out-of-fold predictions.
Parameters
----------
X : pd.DataFrame or np.ndarray
Feature matrix containing pre-treatment covariates or effect modifiers.
y : str or array-like
Outcome, second variable, or numeric input in the comparison.
d : object
Treatment or decision indicator used by the causal score.
z : str or array-like
Instrument, conditioning variable, or third input used by the model.
outcome_learner : object
Model used to learn the outcome regression nuisance function.
instrument_learner : object
Model used to learn the instrument-related nuisance function.
treatment_learner : object
Model used to learn the treatment or propensity nuisance function.
n_splits : object
Number of folds or sample splits used for cross-fitting.
random_state : int
Random-state value used to make sampling or model fitting reproducible.
Returns
-------
tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]
Out-of-fold IIVM nuisance predictions for outcome, treatment, and instrument components.
"""
cv = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
g0_hat = np.zeros(len(y))
g1_hat = np.zeros(len(y))
m_hat = np.zeros(len(y))
r0_hat = np.zeros(len(y))
r1_hat = np.zeros(len(y))
for train_idx, test_idx in cv.split(X):
X_train = X.iloc[train_idx]
X_test = X.iloc[test_idx]
y_train = y[train_idx]
d_train = d[train_idx]
z_train = z[train_idx]
g0_model = clone(outcome_learner).fit(X_train[z_train == 0], y_train[z_train == 0])
g1_model = clone(outcome_learner).fit(X_train[z_train == 1], y_train[z_train == 1])
m_model = clone(instrument_learner).fit(X_train, z_train)
r0_model = clone(treatment_learner).fit(X_train[z_train == 0], d_train[z_train == 0])
r1_model = clone(treatment_learner).fit(X_train[z_train == 1], d_train[z_train == 1])
g0_hat[test_idx] = g0_model.predict(X_test)
g1_hat[test_idx] = g1_model.predict(X_test)
m_hat[test_idx] = m_model.predict_proba(X_test)[:, 1]
r0_hat[test_idx] = r0_model.predict_proba(X_test)[:, 1]
r1_hat[test_idx] = r1_model.predict_proba(X_test)[:, 1]
return g0_hat, g1_hat, np.clip(m_hat, 0.02, 0.98), np.clip(r0_hat, 0.01, 0.99), np.clip(r1_hat, 0.01, 0.99)
def prediction_vector(doubleml_model, learner_key):
"""
Idea: Return predictions from a learner as a flat numeric vector with shape compatible with later diagnostics.
Parameters
----------
doubleml_model : object
Model object used for `doubleml` fitting, prediction, or comparison.
learner_key : object
Dictionary key identifying which learner configuration is being evaluated.
Returns
-------
np.ndarray
One-dimensional prediction array aligned with the evaluation rows.
"""
arr = np.asarray(doubleml_model.predictions[learner_key])
if arr.ndim != 3:
raise ValueError(f"Expected a 3D prediction array, got shape {arr.shape}")
return arr[:, 0, 0]
def learner_loss_table(doubleml_model, model_label):
"""
Idea: Assemble a DataFrame for the learner loss table so the result can be displayed and saved consistently.
Parameters
----------
doubleml_model : object
Model object used for `doubleml` fitting, prediction, or comparison.
model_label : str
Display label for `model`.
Returns
-------
pd.DataFrame
Out-of-fold nuisance prediction loss table by learner and nuisance role.
"""
losses = doubleml_model.evaluate_learners()
rows = []
for learner_name, values in losses.items():
arr = np.asarray(values)
rows.append(
{
"model": model_label,
"learner": learner_name,
"mean_loss": float(np.mean(arr)),
"min_loss": float(np.min(arr)),
"max_loss": float(np.max(arr)),
}
)
return pd.DataFrame(rows)
