Loyalty Program Effects with Causal Inference

Introduction

"Correlation is not causation."

This is something everyone learns, or at least hears about, at some point in their lives. It also automatically begs the question: how can one prove - or at least approximate - causation? The answer can take two forms: you run a carefully planned experiment, or you run causal inference. Either way you suffer. While experiments are for those who can afford them (either politically or at least financially), causal inference can help us determine the effects of one or more treatment(s) over one or more outcome(s) from observational data -- from data we simply collected but without being able to assign people to groups, control their expenditure, or many of the other variables that influence not only the outcome, but the degree to which they engage with the variable whose effects you are trying to isolate (the treatment). The data is therefore biased, and it is the task of the causal inference framework to de-bias it as much as possible.

One area where these issues abound is in marketing. While at times it is perfectly feasible to simply run a tightly controlled A/B test, at others you simply cannot -- you cannot suddenly take away benefits from members just to see what happens, for example, not without facing enormous consequences which can (and will) offset any benefit from answering the question you want to answer in the first place.

This small project I put together aims to use two causal inference approaches: a more traditional one using DoWhy, and a more advanced one using EconML, to determine the member-level effect of a Loyalty Program over their total expenditure over a given time span. It uses a synthetic dataset I created for the purposes, which gives us the ability to compare our implementations against the data-generating process (DGP), while also implements two different approaches: one to determine the effect of the treatment as a whole, and one that allows a modulator to change the effect of the treatment.

But let's get on with it.

Setup and DGP

# Setup

import warnings
warnings.filterwarnings('ignore')
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

sns.set_theme(
    context='talk',
    font_scale=0.9,
    palette = ['#0F00F5', '#3061FF', '#9AB1FF', '#CDD9FF', '#E6ECFF','#E5E5E5',
               '#B6BBCB', '#878B98','#696A6F','#292A2E'],
    style = {
        'axes.facecolor': '#FFFFFF',
         'axes.edgecolor': '#000000',
         'axes.legend.edgecolor':'#FFFFFF',
         'axes.grid': False,
         'axes.axisbelow': 'line',
         'axes.labelcolor': 'black',
         'figure.facecolor': '#FFFFFF',
         'grid.color': '#b0b0b0',
         'grid.linestyle': '-',
         'text.color': 'black',
         'xtick.color': 'black',
         'ytick.color': 'black',
         'xtick.direction': 'out',
         'ytick.direction': 'out',
         'patch.edgecolor': '#FFFFFF',
         'patch.force_edgecolor': True,
         'image.cmap': 'viridis',
         'font.family': ['sans-serif'],
         'font.sans-serif': 'Helvetica Neue',
         'xtick.bottom': False,
         'xtick.top': False,
         'ytick.left': False,
         'ytick.right': False,
         'axes.spines.left': False,
         'axes.spines.bottom': False,
         'axes.spines.right': False,
         'axes.spines.top': False
    }
)

The DGP will try to capture (some of) the variables and effects that normally affect the expenditure of someone at a store: from the gender of the person, their age, income, etc. The effect that we want to isolate is a two-part problem, as we know from our conversations with our business partners (or so we are pretending) that:

• there is an effect for simply being subscribed to our loyalty programs (members get better offers), and that
• there is also an effect modulated by the number of offers the member receives.

So the DGP looks as follows.

# Create the data

np.random.seed(0) 
n = 10000 # number of customers

gender = np.random.binomial(1, 0.3, n)   # 1 = male
age = np.random.normal(40, 10, n).clip(18, 80)
avg_distance_to_store = np.random.exponential(5, n).clip(0, 50)
app_use_frequency = np.random.poisson(5, n).clip(0, 20)
shopping_frequency = np.random.poisson(8, n).clip(0, 30)

# Age influences income and SES
income = (30000 
          + 1000 * age 
          + np.random.normal(0, 5000, n) 
          + np.random.normal(0, 2000 * gender, n))

ses = 0.1 + 0.02 * age + np.random.normal(-0.15, 0, n)

# Subscription (treatment)
subscription_prob = 1 / (1 + np.exp((-4 + 5e-5 * income + 3 * ses + np.random.normal(-0.01, 0.01, n))))
subscription = np.random.binomial(1, subscription_prob, n)

# Number of offers sent
number_offers_sent = np.random.poisson(5 + 8 * subscription, n).clip(0, 50)


base = (1000 
        + 50 * shopping_frequency 
        + 20 * app_use_frequency 
        - 10 * gender 
        - 5 * avg_distance_to_store 
        + 200 * income / 10000 
        + 500 * ses)

# Direct effect of subscription
subscription_direct = (300 + 0.003 * (income - income.mean()) + 200 * (ses - ses.mean()))

# Indirect effect through number_offers_sent
offers_effect = 20 * number_offers_sent

# Sales under treatment (subscription = 1)
sales_1 = base + subscription_direct + offers_effect + np.random.normal(0, 10, n) # test

# Sales under control (subscription = 0) - no subscription effect, and the effect of offers is a fraction of those under the subscription model
sales_0 = base + 0 + offers_effect*.7 + np.random.normal(0, 10, n) # test

sales = np.where(subscription == 1, sales_1, sales_0)
sales = np.maximum(sales, 0)

# True Individual Treatment Effect (Ground Truth)
true_subscription_effect = sales_1 - sales_0

# Create DataFrame
df = pd.DataFrame({
    'gender': gender,
    'age': age,
    'ses': ses,
    'income': income,
    'avg_distance_to_store': avg_distance_to_store,
    'app_use_frequency': app_use_frequency,
    'number_offers_sent': number_offers_sent,
    'shopping_frequency': shopping_frequency,
    'subscription': subscription,
    'sales': sales,
    'true_subscription_effect': true_subscription_effect 
})

So, the data we collected looks like this:

df.head()

	gender	age	ses	income	avg_distance_to_store	app_use_frequency	number_offers_sent	shopping_frequency	subscription	sales	true_subscription_effect
0	0	32.741797	0.604836	58105.598806	1.519148	5	7	10	0	3170.097812	242.549796
1	1	45.634755	0.862695	80104.359705	3.630417	3	5	12	0	3734.250169	389.732319
2	0	35.643679	0.662874	46924.256281	4.341691	5	18	8	1	3321.260123	321.361464
3	0	38.954475	0.729089	65089.335538	0.793638	6	6	11	0	3416.325273	306.980471
4	0	18.000000	0.310000	43417.867289	21.319269	5	12	5	1	2643.909450	216.231935

And we have the luxury of calculating the ground truth and average treatment effect (ATE) directly, for reference purposes

print(f"True ATE: {df['true_subscription_effect'].mean():.2f}")

True ATE: 339.09

As I will not spend time doing the mandatory exploratory data analysis (not the point of this exercise), I will simply show you the distributions.

_, axes = plt.subplots(4,3,figsize=(20,20))

for col, ax in zip(df.columns, axes.ravel()):
    sns.histplot(
        data=df,
        x=col,
        ax=ax,
        bins=100
    )

Note a couple things from the DGP that do not translate to a real-world causal inference approach:

• On a tutorial we have perfect information - we know exactly the relationship of the variables with each other. In the real world we can only approximate this, and we will never be 100% sure we modelled the problem appropriately.

• On a tutorial we know the underlying structure of every variable (distribution, range, etc). In the real world, how do we even know if the data we captured is captured in a way that reflects the actual nature of the behaviour we wished to capture? How do we know, for example, that Socioeconomic Status (SES) truly ranges and distributes in the way our data suggests?

• On a tutorial we do not have missing variables. In the real world there will always be missing variables, even ones we did not know we missed.

• On a tutorial we know the nature of the treatment, and we select the right approach that matches it. In the real world we can only approximate it, and only make our best assumptions about it.

• On a tutorial we know the ground truth, so we can know when we estimated the wrong values. In the real world we do not know the ground truth, therefore we have to rely on other methods to assert our confidence (while still have no ground truth to assess against).

Yet causal inference is nonetheless required to understand and isolate the true effects of whatever we want to measure. If we simply calculated a naive (and biased) difference between those who had the subscription and those who did not, we end up with a very wrong estimate, because we were missing the fact that the effect over the outcome is confounded by every variable that affects both the treatment and the outcome (in this case ses and income)

naive_diff = df[df['subscription']==1]['sales'].mean() - df[df['subscription']==0]['sales'].mean()

print(f'Naive (biased) difference in means: {naive_diff:.2f}') # Biased due to confounding

Naive (biased) difference in means: 140.25

So the first step is to model the causal relationship between the variables, and capture them in a direct acyclic graph (DAG).

import networkx as nx

causal_graph = nx.DiGraph([
    ('age', 'income'),
    ('age', 'ses'),


    ('gender', 'income'),
    ('gender', 'sales'),

    ('income', 'subscription'),
    ('income', 'sales'),
    
    ('ses', 'subscription'),
    ('ses', 'sales'),

    ('avg_distance_to_store', 'sales'),

    ('app_use_frequency', 'sales'),    

    ('subscription', 'number_offers_sent'),
    ('subscription', 'sales'), 

    ('shopping_frequency', 'sales'),
        
    ('number_offers_sent', 'sales'),
    
])

And check that the DAG is indeed acyclic (because otherwise we will get an error later).

for cycle in list(nx.simple_cycles(causal_graph)):
    print(cycle, '\n') # no output means no cycles, so we're good.

from dowhy import CausalModel
causal_model = CausalModel(
    data=df,
    graph=causal_graph,
    treatment='subscription',
    mediator=['number_offers_sent'],
    outcome='sales'
)

causal_model.view_model(size=(10,7))

Once we are happy with the modelled relationships, we can proceed to have a first try at estimating the average treatment effect for the dataset. DoWhy will allow you a number of options, but an important one is whether to include effect modifiers (i.e. modulators). The choice depends on the DAG and the problem at hand. If the business knowledge indicates there should be one or more, then you include it/them. Otherwise they can be omitted. Understandably, this will change your ATE estimation, and what you want to isolate and report: if you are interested in reporting the effect of the subscription program (with the effect of the offers embedded in it), then you can omit the effect modifier(s). But if you want to isolate the effect of the subscription separate from whatever effect the offers might bring, then you should include said modifier(s).

And, as you might have noted, none of them is inherently wrong. It simply depends on what you are aiming to estimate.

estimand = causal_model.identify_effect(proceed_when_unidentifiable=True)
estimate = causal_model.estimate_effect(
    estimand,
    method_name='backdoor.linear_regression',
    effect_modifiers=[],
    confidence_intervals=True,
    control_value=0,
    treatment_value=1,
    target_units='ate'
)
print(estimate)

*** Causal Estimate ***

## Identified estimand
Estimand type: EstimandType.NONPARAMETRIC_ATE

### Estimand : 1
Estimand name: backdoor
Estimand expression:
       d                            
───────────────(E[sales|income,ses])
d[subscription]                     
Estimand assumption 1, Unconfoundedness: If U→{subscription} and U→sales then P(sales|subscription,income,ses,U) = P(sales|subscription,income,ses)

## Realized estimand
b: sales~subscription+income+ses
Target units: ate

## Estimate
Mean value: 450.14089989952527
95.0% confidence interval: [[441.56802038 458.71377942]]

The above is the ATE for the whole of the subscription effect (without decomposition). Now we can compare it with the decomposed effect.

del estimand, estimate

estimand = causal_model.identify_effect(proceed_when_unidentifiable=True)
estimate = causal_model.estimate_effect(
    estimand,
    method_name='backdoor.linear_regression',
    effect_modifiers=['number_offers_sent'],
    confidence_intervals=True,
    control_value=0,
    treatment_value=1,
    target_units='ate'
)
print(estimate)

*** Causal Estimate ***

## Identified estimand
Estimand type: EstimandType.NONPARAMETRIC_ATE

### Estimand : 1
Estimand name: backdoor
Estimand expression:
       d                            
───────────────(E[sales|income,ses])
d[subscription]                     
Estimand assumption 1, Unconfoundedness: If U→{subscription} and U→sales then P(sales|subscription,income,ses,U) = P(sales|subscription,income,ses)

## Realized estimand
b: sales~subscription+income+ses+subscription*number_offers_sent
Target units: 

## Estimate
Mean value: 322.64825163402065
95.0% confidence interval: (np.float64(307.4672562291462), np.float64(338.964849450836))
### Conditional Estimates
__categorical__number_offers_sent
(-0.001, 3.0]    239.682919
(3.0, 5.0]       282.458099
(5.0, 6.0]       312.078129
(6.0, 9.0]       346.717877
(9.0, 28.0]      461.872800
dtype: float64

Note that:

• The estimand shows you what variables you need to control for in the minimal deconfounding set to appropriately estimate the ATE, i.e. to achieve unconfoundedness: ses, and income
• The mean value calculated (ATE) changes whether you include the modulator or not, as described before.
• The effect described with minimal interaction with the effect modifier (first line of the conditional estimates), lands much closer to the ground truth. It would have landed perfectly on the ground truth if the data had not had any noise added to it, but this way it is more realistic to what happens in a real-world scenario.

If the linear approach from DoWhy is modelled via linear regression using the same four variables described in the estimand, you still land a mean effect that is much closer to the ground truth that the heavily biased estimate we saw earlier in the naive difference. Once again the effect will depend on whether you formulate the problem correctly. A simple linear approach omitting the number of offers sent will land you an estimate (the coefficient), similar to DoWhy's ATE without modulators. You would have to include the modulator and the interaction to yield the correct coefficients closer to what you can observe in the DGP.

import statsmodels.formula.api as smf

model_smf = smf.ols(formula='sales ~ subscription + income + ses', data=df)
ols_model = model_smf.fit()
print(ols_model.summary().tables[1])

================================================================================
                   coef    std err          t      P>|t|      [0.025      0.975]
--------------------------------------------------------------------------------
Intercept     1471.9111     12.730    115.629      0.000    1446.958    1496.864
subscription   450.1409      4.373    102.925      0.000     441.568     458.714
income           0.0207      0.000     65.486      0.000       0.020       0.021
ses            525.0547     17.972     29.215      0.000     489.826     560.283
================================================================================

It is always a good idea to run some refutation tests for the estimate, like the placebo test (replacing treatment with a randomly-generated placebo variable), and the random common cause (introducing a randomly generated confounder (that may have been unobserved)). In both cases, having a p < 0.05 means the test failed, suggests that the original estimator is likely biased due to unaccounted confounders or model misspecification.

placebo_results = causal_model.refute_estimate(estimand, estimate, show_progress_bar=True, method_name='placebo_treatment_refuter')
print(placebo_results) # p > 0.05 means the test passed

Refuting Estimates:   0%|          | 0/100 [00:00<?, ?it/s]Refuting Estimates: 100%|██████████| 100/100 [00:01<00:00, 54.77it/s]

Refute: Use a Placebo Treatment
Estimated effect:322.64825163402065
New effect:-0.14152878524887
p value:1.0

random_causa_results = causal_model.refute_estimate(estimand, estimate, show_progress_bar=True, method_name='random_common_cause')
print(random_causa_results) # p > 0.05 means the test passed

Refuting Estimates: 100%|██████████| 100/100 [00:01<00:00, 54.41it/s]

Refute: Add a random common cause
Estimated effect:322.64825163402065
New effect:322.6506314575354
p value:0.98

As mentioned earlier, the effect of the subscription under the above assumptions dismisses the offer x subscription interaction, leading to an inflated coefficient that does not correspond to what is observed in the DGP for the effect of the subscription alone. If we allow for the interaction effect within the linear regression modelling using statsmodels, we are able to recover effects much closer to those seen in the DGP for each variable.

model_final = smf.ols('sales ~ subscription + number_offers_sent + subscription * number_offers_sent + income + ses', data=df).fit()
print(model_final.summary().tables[1])

===================================================================================================
                                      coef    std err          t      P>|t|      [0.025      0.975]
---------------------------------------------------------------------------------------------------
Intercept                        1400.5023     12.836    109.109      0.000    1375.342    1425.663
subscription                      267.4356     13.627     19.626      0.000     240.724     294.147
number_offers_sent                 14.9606      0.780     19.179      0.000      13.431      16.490
subscription:number_offers_sent     4.8419      1.232      3.930      0.000       2.427       7.257
income                              0.0206      0.000     67.751      0.000       0.020       0.021
ses                               529.7462     17.296     30.628      0.000     495.842     563.650
===================================================================================================

from statsmodels.graphics.gofplots import qqplot

_, axes = plt.subplots(1,2, figsize=(12,5))

sns.scatterplot(x=model_final.fittedvalues, y=model_final.resid, alpha=.1, ax=axes[0], s=8)

qqplot(model_final.resid, ax=axes[1], alpha=.1, line='s');

Warning: you will only retrieve the correct effects with this approach if you set up the DAG in a manner that reflects the DGP. Even with the same data, if you make the wrong causal assumption -- i.e. draw an incorrect arrow, or fail to draw an arrow -- you will get the wrong result. Refutation tests are useful here, but that is another topic in itself.

However, DoWhy is only one (sometimes limited) way to determine effect sizes. It is usually the case that the linear regression assumption of the relationship between the treatment and outcome, for example, does not fit the data.

Fortunately, modern econometric frameworks can also leverage modern ML techniques and give us further flexibility to determine effects without having to make too many assumptions. They are also compatible with modern ML practices, such as xAI (SHAP) and bayesian hyperparameter optimisation (optuna).

EconML

Although there are many Causal Inference frameworks that can be leveraged using ML, EconML is in my opinion the one that offers the best balance between compatibility with optuna and SHAP, has the best range of models to choose from, all while allowing for heterogeneity automation. The latter is critical for industry practitioners who are often against tight deadlines and unable to delve into problems for months.

Regardless of the approach, using these modern frameworks requires the creation of two underlying so-called "nuisance models" to handle the bias in the observational data. One of the models is created to predict the outcome (model_y or model_regression), and the other one to predict the treatment (model_t, or model_propensity). The mathematical intuition of how these models work is outside the scope of this notebook (but you can read it here if you're interested).

One can use a number of models for regression or propensity (classification), and in theory use any variables available to train them and achieve good performance -- yet there are some (general) rules of thumb to obtain good results:

• Do not include mediators or colliders (to avoid Berkson's paradox).
• Optimise the model hyperparameters to reduce the error rates (width of the confidence intervals) when estimating effects.

When it comes to variable choice for model training, the following should be noted:

• Always include the confounding variables (in this case income and ses, referred more generally as W) in both models.
• If a variable is a good predictor of both Y (outcome) and T (treatment), include it in both models
• If a variable is a good predictor of Y but not of T, include it in the Y model only
• If a variable is a good predictor of T but not of Y, do not include it in any model
• If a variable is not a good predictor of either Y or T, do not include it in any model (it's just noise)

Therefore we first need to determine the above. And although it is easy to spot mediators and colliders on simple DAGs, using the naked eye on more complex diagrams is a recipe for disaster. This is why the below can be easily used to extract the list of variables we need to avoid.

def get_mediators(graph, treatment='T', outcome='Y'):
    outcome_ancestors = set(nx.ancestors(graph, outcome))
    treatment_descendants = set(nx.descendants(graph, treatment))
    return treatment_descendants.intersection(outcome_ancestors)

def get_colliders(graph, treatment='T', outcome='Y'):
    outcome_descendants = set(nx.descendants(graph, outcome))
    treatment_descendants = set(nx.descendants(graph, treatment))
    pre_colliders = outcome_descendants.intersection(treatment_descendants)

    colliders = []
    for c in pre_colliders:
        for path in nx.all_simple_paths(graph, treatment, c):
            if outcome not in path:
                colliders.append(c)
    return set(colliders)

mediators_to_exclude = get_mediators(causal_graph, 'subscription', 'sales')
colliders_to_exclude = get_colliders(causal_graph, 'subscription', 'sales')

print(f'Mediators to Exclude: {mediators_to_exclude}')
print(f'Colliders to Exclude: {colliders_to_exclude}')

Mediators to Exclude: {'number_offers_sent'}
Colliders to Exclude: set()

And now for determining the variables to use, this can be done by creating (and optimising) two models, and then comparing their feature importance (permutation and SHAP work well, but the former is much faster when there is a lot of data).

T_col = 'subscription'
y_col = 'sales'

X = df.drop(columns=[T_col, y_col, 'true_subscription_effect', 'number_offers_sent'])
W = df[['income', 'ses']]
T = df[T_col]
Y = df[y_col]

from sklearn.model_selection import train_test_split, cross_val_score, KFold
from xgboost import XGBRegressor, XGBClassifier
import optuna
import shap


X_train, X_test, W_train, W_test, T_train, T_test, y_train, y_test = train_test_split(X, W, T, Y, test_size=0.3, random_state=0)
kf = KFold(n_splits = 5, shuffle=True, random_state=0)

optuna.logging.set_verbosity(optuna.logging.WARNING)

# Objective for Outcome Model (Regression)
def objective_reg(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'random_state': 0
    }
    model = XGBRegressor(**params)
    model.fit(X_train, y_train)
    score = cross_val_score(model, X_test, y_test, cv=kf, scoring='r2', n_jobs=-1)
    return score.mean()  

study_reg = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=0), )
study_reg.optimize(objective_reg, n_trials=50, show_progress_bar=True)
best_params_reg = study_reg.best_params

Best trial: 31. Best value: 0.742021: 100%|██████████| 50/50 [00:47<00:00,  1.06it/s]

# Objective for Treatment Model (Classification)
def objective_clf(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 0.9),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'min_child_weight': trial.suggest_int('min_child_weight', 10, 200),
        'gamma': trial.suggest_int('gamma', 5, 20),
        'scale_pos_weight': trial.suggest_float('scale_pos_weight', 1.0, 4*1.5),
        'random_state': 0
    }
    model = XGBClassifier(**params)
    model.fit(X_train, T_train)
    score = cross_val_score(estimator=model, X=X_test, y=T_test, cv=kf, scoring='neg_brier_score', n_jobs=-1)
    return score.mean()  

study_clf = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=42))
study_clf.optimize(objective_clf, n_trials=50, show_progress_bar=True)
best_params_clf = study_clf.best_params

Best trial: 45. Best value: -0.132962: 100%|██████████| 50/50 [00:19<00:00,  2.59it/s]

from sklearn.inspection import permutation_importance

# Permutations: Y model
best_model_y = XGBRegressor(**best_params_reg).fit(X_train, y_train)

perm_importances_y = permutation_importance(
    best_model_y, X_test, y_test, n_repeats=10, random_state=0, n_jobs=-1
)

sorted_importances_idx_y = perm_importances_y.importances_mean.argsort()

permutation_importances_y = pd.DataFrame(
    perm_importances_y.importances[sorted_importances_idx_y].T,
    columns=X_test.columns[sorted_importances_idx_y]
)

# Permutations: T model
best_model_t = XGBClassifier(**best_params_clf).fit(X_train, T_train)

perm_importances_t = permutation_importance(
    best_model_t, X_test, T_test, n_repeats=10, random_state=0, n_jobs=-1
)

sorted_importances_idx_t = perm_importances_t.importances_mean.argsort()

permutation_importances_t = pd.DataFrame(
    perm_importances_t.importances[sorted_importances_idx_t].T,
    columns=X_test.columns[sorted_importances_idx_t]
)

# Plot the importances 
_, axes = plt.subplots(1,2,figsize=(20,7))

sns.boxplot(
    data=pd.melt(permutation_importances_y, var_name='Feature', value_name='Importance').sort_values(by='Importance', ascending=False),
    x='Importance',
    y='Feature',
    orient='h',
    linewidth=.5,
    ax=axes[0]
).set_title('Feature Importance, Y ~ X')
axes[0].axvline(x=0, color='k', linestyle='--', linewidth=.8)


sns.boxplot(
    data=pd.melt(permutation_importances_t, var_name='Feature', value_name='Importance').sort_values(by='Importance', ascending=False),
    x='Importance',
    y='Feature',
    orient='h',
    linewidth=.5,
    ax=axes[1]
).set_title('Feature Importance, T ~ X')
axes[1].axvline(x=0, color='k', linestyle='--', linewidth=.8);

From the above we can therefore conclude that:

• We have to include income and ses because they are confounders (regardless of performance).
• We can include age in both nuisance models
• shopping_frequency and app_use_frequency can be included in the Y model only.

Note that I optimised for R² (for the Y model), and for brier score (for the T model). This is because for causal inference we are more interested in model explainability and propensity calibration than in minimising prediction errors - as the models will not be used for that.

We can now re-create the final nuisance models required for the EconML models backbones, and estimate our effects.

#delete these variables to re-instantiate them with the final configuration
del X, W, T, Y, best_params_reg, best_params_clf, best_model_y, best_model_t, X_train, X_test, W_train, W_test, T_train, T_test, y_train, y_test

confounders = ['income', 'ses']
y_only_cols = ['shopping_frequency', 'app_use_frequency', 'avg_distance_to_store']

X = df[['age'] + y_only_cols + confounders]
W = df[confounders]
T = df[T_col]
Y = df[y_col]

X_train, X_test, W_train, W_test, T_train, T_test, y_train, y_test = train_test_split(X, W, T, Y, test_size=0.3, random_state=0)
kf = KFold(n_splits = 5, shuffle=True, random_state=0)

# Objective for Outcome Model (Regression)
def objective_reg(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'random_state': 0
    }
    model = XGBRegressor(**params)
    model.fit(X_train, y_train)
    score = cross_val_score(model, X_test, y_test, cv=kf, scoring='r2', n_jobs=-1)
    return score.mean()  

study_reg = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=0), )
study_reg.optimize(objective_reg, n_trials=50, show_progress_bar=True)
best_params_reg = study_reg.best_params

Best trial: 17. Best value: 0.741147: 100%|██████████| 50/50 [00:42<00:00,  1.17it/s]

# Objective for Treatment Model (Classification)
def objective_clf(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 0.9),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'min_child_weight': trial.suggest_int('min_child_weight', 10, 200),
        'gamma': trial.suggest_int('gamma', 5, 20),
        'scale_pos_weight': trial.suggest_float('scale_pos_weight', 1.0, 4*1.5),
        'random_state': 0
    }
    model = XGBClassifier(**params)
    model.fit(X_train.drop(columns=y_only_cols), T_train)
    score = cross_val_score(estimator=model, X=X_test.drop(columns=y_only_cols), y=T_test, cv=kf, scoring='neg_brier_score', n_jobs=-1)
    return score.mean()  

study_clf = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=42))
study_clf.optimize(objective_clf, n_trials=50, show_progress_bar=True)
best_params_clf = study_clf.best_params

Best trial: 46. Best value: -0.133673: 100%|██████████| 50/50 [00:14<00:00,  3.41it/s]

And now we pass the nuisance models to our causal inference model of choice, with the choice depending on a number of factors (see a flowchart here and a more detailed model comparison here).

from econml.dml import CausalForestDML

est = CausalForestDML(
    model_y=XGBRegressor(**best_params_reg),
    model_t=XGBClassifier(**best_params_clf),
    n_estimators=1000,
    max_depth=6,
    min_samples_leaf=10,
    random_state=0,
    verbose=0,
    discrete_outcome=False,
    discrete_treatment=True
)

est.tune(y_train, T_train, X=W_train, W=W_train)
est.fit(y_train, T_train, X=W_train, W=W_train)

ate_est = est.ate(X=W_test)
ate_est_ci = est.ate_interval(X=W_test)
print(f"""
ATE: {ate_est:.2f}
95% CI: ({ate_est_ci[0]:.2f}, {ate_est_ci[1]:.2f})
""")

ATE: 474.70
95% CI: (447.55, 501.84)

The CausalForestDML retrieves a similar effect size than the DoWhy approach under the initial assumptions (i.e. without decomposition), meaning this is the size of the overall effect of subscription (which includes the offers sent on average). This simple approach works well out of the box when there are no assumptions of interactions and modulators changing the effect size, unlike in our DGP. This does not mean the value is incorrect -- it is simply not decomposed with the interaction. Considering there is no knowing the true DGP on a real-world problem, once again retrieving the isolated effect of the program (decomposed from the modulator) boils down to ensuring one holds the correct assumptions when modelling the problem at hand.

Regardless, and before running away with whatever effect size the model estimated, one should briefly inspect the nuisance model performance. This is a critical step that can tell you whether you might need to collect more variables/data (e.g. you do not have enough data to give the Y or T model good performance), or you might need to reformulate the problem (e.g. your treatment definition leads to an extreme class imbalance and the propensity model is not coping).

Although "good enough performance" depends on a myriad of factors, some rules of thumbs apply:

• Avoid leakage, and ensure the performance between the train and test sets is comparable.
• Aim for R² of .4 and above.
• Do not get hung up on trying to get perfect classification metrics -- better calibration is more important than the perfect AUC.
• If you are using a binary treatment approach (like in this example), aim to have a model with reasonable overlap in propensity between classes.

# Check nuisance models performance

from sklearn.metrics import (
    r2_score,
    mean_absolute_error,
    mean_squared_error,
    roc_auc_score,
    log_loss,
    brier_score_loss,
    accuracy_score,
    precision_score,
    recall_score,
    f1_score
)

from sklearn.metrics import RocCurveDisplay
from sklearn.calibration import CalibrationDisplay


def evaluate_nuisance_models(best_params_reg = None, best_params_clf = None, X_train = None, X_test = None, 
                             T_train = None, T_test = None, y_train = None, y_test = None,
                             includes_mediator=False):

    
    y_model = XGBRegressor(**best_params_reg)
    t_model = XGBClassifier(**best_params_clf)

    y_model.fit(X_train, y_train)
    y_pred_train = y_model.predict(X_train)
    y_pred_test = y_model.predict(X_test)

    if includes_mediator == False:
        t_model.fit(X_train.drop(columns=y_only_cols), T_train)
        t_pred_train = t_model.predict(X_train.drop(columns=y_only_cols))
        t_proba_train = t_model.predict_proba(X_train.drop(columns=y_only_cols))[:, 1]

        t_pred_test = t_model.predict(X_test.drop(columns=y_only_cols))
        t_proba_test = t_model.predict_proba(X_test.drop(columns=y_only_cols))[:, 1]

    else:
        t_model.fit(X_train.drop(columns=y_only_cols + mediator), T_train)
        t_pred_train = t_model.predict(X_train.drop(columns=y_only_cols + mediator))
        t_proba_train = t_model.predict_proba(X_train.drop(columns=y_only_cols + mediator))[:, 1]

        t_pred_test = t_model.predict(X_test.drop(columns=y_only_cols + mediator))
        t_proba_test = t_model.predict_proba(X_test.drop(columns=y_only_cols + mediator))[:, 1]


    print(f"""
{'#'*10} Model Y Performance {'#'*10}
Train/Test R2: {r2_score(y_train, y_pred_train):.4f}, {r2_score(y_test, y_pred_test):.4f}
Train/Test MAE: {mean_absolute_error(y_train, y_pred_train):.2f}, {mean_absolute_error(y_test, y_pred_test):.2f}
Train/Test RMSE: {mean_squared_error(y_train, y_pred_train):.2f},  {mean_squared_error(y_test, y_pred_test):.2f}

{'#'*10} Model T Performance {'#'*10}
Train/Test AUC: {roc_auc_score(T_train, t_proba_train):.3f}, {roc_auc_score(T_test, t_proba_test):.3f}
Train/Test Log Loss: {log_loss(T_train, t_proba_train):.4f}, {log_loss(T_test, t_proba_test):.4f}
Train/Test Brier Score: {brier_score_loss(T_train, t_proba_train):.4f}, {brier_score_loss(T_test, t_proba_test):.4f}

Accuracy (Test): {accuracy_score(T_test, t_pred_test):.2f}
Precision (Test): {precision_score(T_test, t_pred_test):.2f}
Recall (Test): {recall_score(T_test, t_pred_test):.2f}
F1 (Test): {f1_score(T_test, t_pred_test):.2f}
    """)

    _, axes = plt.subplots(2,3, figsize=(20,15))

    #### Model Y ####

    # Y Preds
    sns.scatterplot(x=y_test, y=y_pred_test, alpha=.3, s=5, ax=axes[0,0]).set_title('Model Y: Preds v Ground Truth')
    axes[0,0].set_xlabel('y_true')
    axes[0,0].set_ylabel('y_pred')
    axes[0,0].plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], color='k', ls='--', lw=1)

    # Residuals
    res = y_test - y_pred_test
    sns.histplot(x=res, ax=axes[0,1], bins=100).set_title('Model Y: Residual Distribution')
    axes[0,1].axvline(0, color='k', ls='--', lw=1)

    qqplot(res, ax=axes[0,2], alpha=.1, line='s');
    axes[0,2].set_title('Model Y: QQ Plot')

    #### Model T ####


    RocCurveDisplay.from_predictions(T_test, t_proba_test, ax=axes[1,0])
    axes[1,0].set_title('Model T: ROC AUC')

    CalibrationDisplay.from_estimator(
        estimator=t_model,
        X=X_test.drop(columns=y_only_cols) if includes_mediator == False else X_test.drop(columns=y_only_cols + mediator),
        y=T_test,
        n_bins=10,
        strategy='quantile',
        ax=axes[1,1],
        **{'marker': 'o', 'markersize': 8, 'lw': 1, 'alpha': .8}
    )
    axes[1,1].set_title('Model T: Calibration')

    overlap_df = pd.DataFrame({'treatment': T_test.values, 'propensity': t_proba_test})
    sns.histplot(
        data=overlap_df,
        x='propensity',
        hue='treatment',
        bins=50,
        ax=axes[1,2],
        common_norm=False,
        stat='density',
        element='step',
        palette='bright'
    ).set_title('Model T: Propensity Overlap')

evaluate_nuisance_models(
    best_params_reg=best_params_reg, 
    best_params_clf=best_params_clf, 
    X_train=X_train, 
    X_test=X_test, 
    y_train=y_train, 
    y_test=y_test, 
    T_train=T_train,
    T_test=T_test
)

########## Model Y Performance ##########
Train/Test R2: 0.7667, 0.7533
Train/Test MAE: 120.60, 124.52
Train/Test RMSE: 28054.77,  29505.37

########## Model T Performance ##########
Train/Test AUC: 0.773, 0.776
Train/Test Log Loss: 0.4192, 0.4178
Train/Test Brier Score: 0.1323, 0.1318

Accuracy (Test): 0.81
Precision (Test): 0.57
Recall (Test): 0.20
F1 (Test): 0.29
    

The nuisance model performance is telling us two important things:

• The Y ~ X model predictions v ground truth is hinting at us that since we have two different populations in the DGP (subscribed v unsubscribed), and the model is outputting predictions v ground truth that show two distinctive clusters, we are unlikely to capture everything with the current model because we might be missing something (the mediator, num_offers_sent) -- we'll come back to this later.

• The T ~ X model propensity overlap reinforces the above, telling us that the two different populations might have very different effects (as we know from the DGP).

Despite the above, the nuisance models performance is still acceptable: neither perfect (I set up the DGP with noise), nor terrible. We can move forward.

Inspecting the nuisance model performance helps us in many ways, for example:

• Ensuring the estimate given by the EconML models are reliable
• Troubleshooting for potential data- or problem-framing-related issues (e.g. calibration deciles unevenly distributed, which are often downstream of poor propensity of receiving the treatment (e.g. high self-selection into the treatment))
• If the EconML approach yields an ATE that is substantially different to what you observed with DoWhy, who do you believe? We need to ensure it is not an issue of nuisance model performance, for example.
• If you want to try other EconML approaches, you need to ensure you can rely on the base nuisance models.

Speaking of which, it is always a good sense-check to see if other similar EconML models (or maybe slightly different ones) give you a similar ATEs.

from econml.dr import LinearDRLearner

ldml = LinearDRLearner(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cv=5,
    random_state=0
)

ldml.fit(y_train, T_train, X=W_train, W=W_train)
print(f'ATE: {ldml.ate(X=W_train):.2f}')

ATE: 486.08

from econml.dr import ForestDRLearner

fdr = ForestDRLearner(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cv=5,
    random_state=0,
    n_estimators=1000,
    max_depth=6,
    min_samples_leaf=10,
    verbose=0,
    discrete_outcome=False
)

fdr.fit(y_train, T_train, X=W_train)
print(f'ATE: {fdr.ate(X=W_test):.2f}')

ate_fdr_ci = fdr.ate_interval(X=W_test)
print(f'95% CI: {ate_fdr_ci[0]:.2f}, {ate_fdr_ci[1]:.2f}')

ATE: 485.83
95% CI: 447.38, 524.27

So far all the models have returned a similar ATE, but this does not automatically mean all models are correct, or that all perform similarly across other performance metrics.

The EconML suit implements the DRTester module, which includes some useful tests to determine the health of the models. Critically, the documentation states that these are validation tests for CATE models, basically for models where there is heterogeneity in the effect amongst individuals. If you get "bad" performance in these, it might only indicate that there is no heterogeneity in the treatment, rather than indicate a bad model specification (especially if the nuisance model performance is good and the ATEs are estimated with tight confidence intervals).

from econml.validate import DRTester

tester_est = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=est # the CATE model
)

tester_est.fit_nuisance(
    Xtrain=W_train.values,
    Dtrain=T_train.values,
    ytrain=y_train.values,
    Xval=W_test.values,
    Dval=T_test.values,
    yval=y_test.values
)

tester_est.get_cate_preds(Xtrain=W_train.values, Xval=W_test.values)
est_cal_results = tester_est.evaluate_all(Xtrain=W_train.values, Xval=W_test.values, n_groups=10)
est_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	1.164	0.176	0.0	17.895	2.751	0.0	42.296	8.289	0.0	0.392

But what do these mean? Now we are entering uplift and targeting territory.

Thinking of the ATE, we can conclude with confidence that the effect is positive and therefore there is value in getting people to subscribe -- specifically around ~$470 on average per member over whatever period we used to capture the data. We should invest in marketing to promote the subscription. But the budget is tight. We can only target about 20% of people. But whom? Should we do it randomly, or is there a way the model can tell us who will benefit more?

Remember CATE models are models for effects heterogeneity. The ATEs calculated are the effects average, but true heterogeneity means you have a meaningful distribution of effects -- some customers with higher effects, some with lower, with a distribution averaging aound the ATE. Yet while the ATE might be correct, this does not automatically mean the model was able to (correctly) pick up the heterogeneity we can see in the DGP, and make reliable predictions about it. This is where these tests come in handy. So before trusting the model to help us with targeting to get the best bang for buck, let's see if the models are worth using for specific targeting (compared to randomised targeting). Note most of these tests come with an error estimate (*_se) and a p value (*_pval) attached to them.

• The Best Linear Predictor (BLP) works in a way similar to the calibration curve for a propensity model, but asking a slightly different question: are the CATEs well calibrated against observed uplift within a segment (usually decile)? An estimate (blp_est) closer to 1 is best and preferred, though small deviations are not a huge problem.

• The QINI estimate (qini_est) captures a summaty statistic for the QINI curve (which we will see next), plotting the cumulative incremental gains (y-axis) achieved by treating the top fraction of the population. Differently put, the QINI (curve) shows how much gain we obtain from targeting the X% of the top-ranked customer based on their CATEs. The curve should ideally show an inverted U shape: rising as we capture the gains of the those who show the highest predicted uplift effects, and then decaying because we do not gain more by targeting those who have the lowest predicted uplift.

• The Targeting Operator Characteristic (TOC) curve shows the cumulative average uplift (CATE performance) as you expand targeting down the ranked list. The autoc_est is a summary statistic for the area under the TOC curve (which we will also see next). Differently put, TOC tells you about your average rate of return (ROI) at each stage of a campaign rollout. Intuitively, and if the model works appropriately, the TOC should start high on the y axis, and then slowly decay because, and similarly to the QINI curve, there is little more to gain by targeting more and more people.

• The Calibration R Squared cal_r_squared is probably the easiest to understand: similar to wider implementations, it tells us whether the model can correctly group members based on their CATEs. An estimate closer to 1 is preferred, but in practice it is like the R² metric for regression problems: it is hard to get anything past .7 or .8 with real-world, noisy data.

Critically, a model showing less-than-ideal performance on the summary statistics does not immediately equate to the model being useless or incorrectly specified. These metrics are very sensitive to sometimes one or two groups not behaving appropriately, rather than showing the whole model is bad. The visualisation of the curves helps cementing this idea.

And now, the curves.

def plot_cal_results(results=None):
    results.plot_cal(1);
    results.plot_toc(1);
    results.plot_qini(1);

plot_cal_results(est_cal_results)

tester_ldml = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=ldml
)

tester_ldml.fit_nuisance(
    Xtrain=W_train.values,
    Dtrain=T_train.values,
    ytrain=y_train.values,
    Xval=W_test.values,
    Dval=T_test.values,
    yval=y_test.values
)

tester_ldml.get_cate_preds(Xtrain=W_train.values, Xval=W_test.values)
ldml_cal_results = tester_ldml.evaluate_all(Xtrain=W_train.values, Xval=W_test.values, n_groups=10)
ldml_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	0.998	0.164	0.0	17.76	2.806	0.0	42.48	8.574	0.0	0.57

plot_cal_results(ldml_cal_results)

tester_fdr = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=fdr
)

tester_fdr.fit_nuisance(
    Xtrain=W_train.values,
    Dtrain=T_train.values,
    ytrain=y_train.values,
    Xval=W_test.values,
    Dval=T_test.values,
    yval=y_test.values
)

tester_fdr.get_cate_preds(Xtrain=W_train.values, Xval=W_test.values)
fdr_cal_results = tester_fdr.evaluate_all(Xtrain=W_train.values, Xval=W_test.values, n_groups=10)
fdr_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	1.125	0.15	0.0	19.964	2.935	0.0	57.782	9.844	0.0	0.572

plot_cal_results(fdr_cal_results)

From the above a couple of things become apparent:

• Some models seem to struggle with the bottom-most and/or the top-most deciles of members. This usually indicates poor overlap in the data space in one or more features (which was also hinted at by the propensity model calibration and propensity overlaps), and warrants further investigation that escapes the scope of this piece. Sometimes we just need to get more data (or make some exclusions) to improve model performance without sacrificing the problem at hand.

• All the models produce estimates with wide error rates - as expected due to the noise in the DGP.

• All models behave somewhat differently (especially when it comes to calibration), but show reasonable performance at least in terms of TOC and QINI.

What model would you want to move forward with? In my experience, it depends on a number of factors, for example:

• Am I comfortable enough that the model assumptions and implementation match my problem modelling?

• What is the purpose of the model? Sometimes prioritising QINI/TOC is preferable to focusing too highly on calibration. Thinking about what the metrics reflect and what we will use the model for (e.g. efficient targeting vs customer behaviour modelling/segmentation) can tell you what to prioritise on.

• Am I comfortable enough sacrificing a bit of performance to extract richer interactions and effects? The SHAP values you obtain from each model will behave differently, and I personally prefer richer information from flexible models (e.g. Forests). We will come back to SHAP values later.

The models can be further explored in terms of performance on a number of level, primarily the predictions they made, and with what level of confidence (condidence intervals for individual treatment effects), as shown below. Additionally, and since we have access to the ground truth since this is a tutorial, we can compare the predictions against the ground truth -- but in practice you do not get this luxury (and why would you be trying to estimate the ground truth if you had it anyway?).

def evaluate_cates(estimator=None, X_train=None, X_test=None):
    _, axes = plt.subplots(3,2,figsize=(15,20))

    # Calculate CATEs and plot distributions with ATE as vline
    cates_train = estimator.effect(X_train)
    cates_test = estimator.effect(X_test)
    
    cates_df = pd.concat([
        pd.DataFrame({'CATE': cates_train, 'Set': ['Train'] * len(cates_train)}),
        pd.DataFrame({'CATE': cates_test, 'Set': ['Test'] * len(cates_test)}),
    ])
    
    sns.histplot(
        data=cates_df,
        x='CATE',
        hue='Set',
        bins=50,
        ax=axes[0,0],
        common_norm=False,
        stat='density',
        element='step',
        palette='bright'
    ).set_title('CATEs distribution')

    ate_train, ate_test = np.mean(cates_train), np.mean(cates_test)
    color_train, color_test = sns.color_palette('bright')[0], sns.color_palette('bright')[1]
    
    axes[0,0].axvline(x=ate_train, color=color_train, lw=1, ls='-.')
    axes[0,0].text(x=ate_train * .98, y=axes[0,0].get_ylim()[1] * .95, s='Train ATE', ha='right', color=color_train)
    
    axes[0,0].axvline(x=ate_test, color=color_test, lw=1, ls='--')
    axes[0,0].text(x=ate_test * .98, y=axes[0,0].get_ylim()[1] * .9, s='Test ATE', ha='right', color=color_test)

    # Calculate CATE confidence intervals (Test set only) and plot distribution
    cate_ci_low, cate_ci_high = estimator.effect_interval(X_test)
    cate_interval_df = X_test.copy()
    cate_interval_df['cate_pred'] = cates_test
    cate_interval_df['interval_width'] = cate_ci_high - cate_ci_low
    cate_interval_df['interval_includes_zero'] = [1 if i <=0 else 0 for i in cate_ci_low]
    cate_interval_df['cate_decile'] = pd.qcut(cate_interval_df['cate_pred'], q=10, precision=2)

    sns.histplot(data=cate_interval_df, x='interval_width', ax=axes[0,1], bins=50, color=color_test).set_title('Confidence Intervals Distribution (Test Set)')

    mean_interval_width = np.mean(cate_interval_df['interval_width'])
    median_interval_width = np.median(cate_interval_df['interval_width'])
    
    axes[0,1].text(x=axes[0,1].get_xlim()[1] * 0.98, y=axes[0,1].get_ylim()[1] * .95, s=f'Mean:{mean_interval_width:.1f}, Median: {median_interval_width:.1f}', ha='right', color='k')

    # Check if there are any deciles where there are members with CIs including zero
    ci_grouped = cate_interval_df.groupby(['cate_decile', 'interval_includes_zero']).count().loc[:, ['cate_pred']].reset_index()
    
    ci_grouped['mbrs_cate_group'] = ci_grouped.groupby('cate_decile')['cate_pred'].transform(np.sum)
    ci_grouped['prop_of_cate_group'] = ci_grouped['cate_pred'] * 100 / ci_grouped['mbrs_cate_group']
    

    plot_zeros = ci_grouped[ci_grouped['interval_includes_zero']==1]

    if len(plot_zeros) == 0:
        axes[1,0].text(x=0.5, y=0.5, s='No Group Contains\nPredictions With Zero\n\n=)!', ha='center', va='center')

    else:
        sns.barplot(data=plot_zeros, x='cate_decile', y='cate_pred', ax=axes[1,0], color=color_test).set_title('Groups Containing\nPredictions With Zero')
        for container in axes[1,0].containers:
            axes[1,0].bar_label(container, labels = [f'{i:.1f}%' for i in plot_zeros['prop_of_cate_group']], padding=5, fontsize=12)

        axes[1,0].set_xticks(ticks=[i for i in axes[1,0].get_xticks()], labels=[i for i in axes[1,0].get_xticklabels()], rotation=45, ha='right')

    # KDE Plot combining predictions and intervals
    plot_non_zeros = cate_interval_df[cate_interval_df['interval_includes_zero']==0]
    if len(plot_non_zeros) == 0:
        axes[1,1].text(x=0.5, y=0.5, s='All Confidence Intervals\nInclude Zero!', ha='center', va='center')

    else:
        sns.kdeplot(
            data=plot_non_zeros, 
            x='cate_pred', 
            y='interval_width',  
            linewidths=.5, 
            ax=axes[1,1]).set_title('Predictions vs CI Widths')

    # Plot CATE predictions vs ground truth
    # Note that here I am using 6 x number of offers because that is the difference in slopes between subscribed and ubsubscribed for this effect
    # between those with subscription and those without
    true_cate_structural = 300 + 0.003 * (df['income'] - df['income'].mean()) + 200 * (df['ses'] - df['ses'].mean()) + 6.0 * df['number_offers_sent']

    y_true = true_cate_structural.loc[X_test.index].to_numpy()
    y_pred = np.asarray(estimator.effect(X_test)).ravel()

    sns.scatterplot(
        x=y_true,
        y=y_pred,
        alpha=.3,
        s=10,
        ax=axes[2, 0]
    ).set_title('Predicted CATE vs structural ground truth (Test Set)')

    lims = [min(y_true.min(), y_pred.min()), max(y_true.max(), y_pred.max())]
    axes[2, 0].plot(lims, lims, color='k', ls='--', lw=1)
    axes[2, 0].set_xlabel('Structural ground-truth CATE')
    axes[2, 0].set_ylabel('Predicted CATE')

    plt.subplots_adjust(hspace=.3)



# Note that you have to pass the train/test sets that match whatever you trained the estimator on
evaluate_cates(est, W_train, W_test) 

The above are the results for the CausalForestDML approach. Note that:

• The model produces reasonable and feasible effects (top left plot): the spread of effects is neither too wide nor too narrow, with good overlap between the train/test split.
• The model is mostly confident in its predictions (the distribution is right-skewed), with only some few outliers (top right plot). The confidence interval width represents the +/- error the prediction comes with, with 95% confidence.
• No predictions had intervals spanning zero (middle left plot), which usually happens when the confidence of the model is poor.
• When you combine the distribution of effects with their confidence intervals you get a resonable distribution of members over whom you could confidently act for targeting (middle right plot). In the above case, I would include those with higher-than-average CATE and less than 50-60 in error range, so high-value, high-confidence predictions.
• The model is however not too good at lining up its predictions with the ground truth we can see/extract in the DGP. But it's not terrible either (bottom left plot). Again, in practice you do not get this plot, so don't get too fixated on it.

We can do the same evaluations for the other models we set up.

evaluate_cates(ldml, W_train, W_test) 

evaluate_cates(fdr, W_train, W_test) 

But what about the mediator? So far we have completely skipped a critical part of the puzzle. This is on purpose to (i) explore the setup when the DAG does not include a mediator -- so we have basically either assumed there is no mediator and the effect of subscription is similar for the whole population regardless of how many offers we send, or we want to estimate the effect regardless of mediator -- , and (ii) to have a baseline to compare against when we add the mediator next for the same data and models.

Just to reiterate: not including the modulator is not inherently wrong, it is simply a different approach. If we are interested in calculating the effect of the subscription without any factors allowing it to vary (e.g. we just want a top-line estimate for reporting purposes), we can omit the mediator. If however we want to make a decomposition analysis like we did in the earlier section using DoWhy/OLS, then we have to include it. Both estimates are correct, they just focus on different things: one gives you the estimate of the effect of being subscribed (which includes being sent offers, and will be larger because the effect of the offers is included in it), and the other one controls for these offers, and gives you the estimate of simply being subscribed (and therefore lands much closer to what you can observe in the DGP).

But where do you add the modulator? It depends once again on the approach:

Approach: Direct effect (no mediator):

• Outcome Nuisance Y ~ X: No
• Propensity Nuisance T ~ X: No
• EconML X: No
• EconML W: No

Approach: Average controlled effect:

• Outcome Nuisance Y ~ X: Yes
• Propensity Nuisance T ~ X: No
• EconML X: No
• EconML W: Yes

Approach: Effect varies with offers sent:

• Outcome Nuisance Y ~ X: Yes
• Propensity Nuisance T ~ X: No
• EconML X: Yes
• EconML W: No

We already saw and therefore now understand the first and last approaches. The middle approach (average controlled effect) can be interpreted as: "we know there is an effect driven by offers being sent, but we want to control for that (i.e. we do not want the effect to vary with offers, but we just want to account for them)".

Since we want to get closer to the ground truth (where we know the effect varies with offers sent), we will implement that approach. Note that the propensity nuisance model never gets the mediator included, no matter the approach, because the mediator is downstream of the treatment.

del est, ldml, fdr, tester_est, est_cal_results, ldml_cal_results, fdr_cal_results, X, W, Y, T, X_train, X_test, W_train, W_test, T_train, T_test, y_train, y_test, best_params_reg, best_params_clf

confounders = ['income', 'ses']
y_only_cols = ['shopping_frequency', 'app_use_frequency', 'avg_distance_to_store']
mediator = ['number_offers_sent']

X = df[['age'] + y_only_cols + confounders + mediator]
W = df[confounders]
T = df[T_col]
Y = df[y_col]

X_train, X_test, W_train, W_test, T_train, T_test, y_train, y_test = train_test_split(X, W, T, Y, test_size=0.3, random_state=0)
kf = KFold(n_splits = 5, shuffle=True, random_state=0)

def objective_reg(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'random_state': 0
    }
    model = XGBRegressor(**params)
    model.fit(X_train, y_train)
    score = cross_val_score(model, X_test, y_test, cv=kf, scoring='r2', n_jobs=-1)
    return score.mean()  

study_reg = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=0), )
study_reg.optimize(objective_reg, n_trials=50, show_progress_bar=True)
best_params_reg = study_reg.best_params

Best trial: 38. Best value: 0.957518: 100%|██████████| 50/50 [00:34<00:00,  1.45it/s]

def objective_clf(trial):
    params = {
        'learning_rate': trial.suggest_float('learning_rate', 0.0001, 0.1, log=True),
        'n_estimators': trial.suggest_int('n_estimators', 100, 1000),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 0.9),
        'max_depth': trial.suggest_int('max_depth', 3, 6),
        'min_child_weight': trial.suggest_int('min_child_weight', 10, 200),
        'gamma': trial.suggest_int('gamma', 5, 20),
        'scale_pos_weight': trial.suggest_float('scale_pos_weight', 1.0, 4*1.5),
        'random_state': 0
    }
    model = XGBClassifier(**params)
    model.fit(X_train.drop(columns=y_only_cols + mediator), T_train)
    score = cross_val_score(estimator=model, X=X_test.drop(columns=y_only_cols + mediator), y=T_test, cv=kf, scoring='neg_brier_score', n_jobs=-1)
    return score.mean()  

study_clf = optuna.create_study(direction='maximize', sampler=optuna.samplers.TPESampler(seed=42))
study_clf.optimize(objective_clf, n_trials=50, show_progress_bar=True)
best_params_clf = study_clf.best_params

Best trial: 46. Best value: -0.133673: 100%|██████████| 50/50 [00:14<00:00,  3.42it/s]

evaluate_nuisance_models(
    best_params_reg=best_params_reg, 
    best_params_clf=best_params_clf, 
    X_train=X_train, 
    X_test=X_test, 
    y_train=y_train, 
    y_test=y_test, 
    T_train=T_train,
    T_test=T_test,
    includes_mediator=True
)

########## Model Y Performance ##########
Train/Test R2: 0.9707, 0.9629
Train/Test MAE: 30.83, 34.88
Train/Test RMSE: 3527.47,  4441.65

########## Model T Performance ##########
Train/Test AUC: 0.773, 0.776
Train/Test Log Loss: 0.4192, 0.4178
Train/Test Brier Score: 0.1323, 0.1318

Accuracy (Test): 0.81
Precision (Test): 0.57
Recall (Test): 0.20
F1 (Test): 0.29
    

est = CausalForestDML(
    model_y=XGBRegressor(**best_params_reg),
    model_t=XGBClassifier(**best_params_clf),
    n_estimators=1000,
    max_depth=6,
    min_samples_leaf=10,
    random_state=0,
    verbose=0,
    discrete_outcome=False,
    discrete_treatment=True
)

est.tune(y_train, T_train, X=X_train, W=W_train)
est.fit(y_train, T_train, X=X_train, W=W_train)

ate_est = est.ate(X=X_test)
ate_est_ci = est.ate_interval(X=X_test)
print(f"""
ATE: {ate_est:.2f}
95% CI: ({ate_est_ci[0]:.2f}, {ate_est_ci[1]:.2f})
""")

ATE: 314.71
95% CI: (288.46, 340.95)

tester_est = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=est # the CATE model
)

tester_est.fit_nuisance(
    Xtrain=X_train.values,
    Dtrain=T_train.values,
    ytrain=y_train.values,
    Xval=X_test.values,
    Dval=T_test.values,
    yval=y_test.values
)

tester_est.get_cate_preds(Xtrain=X_train.values, Xval=X_test.values)
est_cal_results = tester_est.evaluate_all(Xtrain=X_train.values, Xval=X_test.values, n_groups=10)
est_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	1.115	0.076	0.0	13.681	0.842	0.0	33.842	2.282	0.0	0.008

plot_cal_results(est_cal_results)

ldml = LinearDRLearner(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cv=5,
    random_state=0
)

ldml.fit(y_train, T_train, X=X_train, W=W_train)
print(f'ATE: {ldml.ate(X=X_train):.2f}')

ATE: 337.38

tester_ldml = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=ldml # the CATE model
)

tester_ldml.fit_nuisance(
    Xtrain=X_train.values,
    Dtrain=T_train.values,
    ytrain=y_train.values,
    Xval=X_test.values,
    Dval=T_test.values,
    yval=y_test.values
)

tester_ldml.get_cate_preds(Xtrain=X_train.values, Xval=X_test.values)
ldml_cal_results = tester_ldml.evaluate_all(Xtrain=X_train.values, Xval=X_test.values, n_groups=10)
ldml_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	0.639	0.046	0.0	13.831	0.739	0.0	31.529	1.929	0.0	0.228

plot_cal_results(ldml_cal_results)

fdr = ForestDRLearner(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cv=5,
    random_state=0,
    n_estimators=1000,
    max_depth=6,
    min_samples_leaf=10,
    verbose=0,
    discrete_outcome=False
)

fdr.fit(y_train, T_train, X=X_train, W=W_train)
print(f'ATE: {fdr.ate(X=X_test):.2f}')

ate_fdr_ci = fdr.ate_interval(X=X_test)
print(f'95% CI: {ate_fdr_ci[0]:.2f}, {ate_fdr_ci[1]:.2f}')

ATE: 336.48
95% CI: 317.11, 355.85

tester_fdr = DRTester(
    model_regression=XGBRegressor(**best_params_reg),
    model_propensity=XGBClassifier(**best_params_clf),
    cate=fdr
)

tester_fdr.fit_nuisance(
    Xtrain=X_train.values,
        Dtrain=T_train.values,
        ytrain=y_train.values,
        Xval=X_test.values,
        Dval=T_test.values,
        yval=y_test.values
)

tester_fdr.get_cate_preds(Xtrain=X_train.values, Xval=X_test.values)
fdr_cal_results = tester_fdr.evaluate_all(Xtrain=X_train.values, Xval=X_test.values, n_groups=10)
fdr_cal_results.summary()

	treatment	blp_est	blp_se	blp_pval	qini_est	qini_se	qini_pval	autoc_est	autoc_se	autoc_pval	cal_r_squared
0	1	0.814	0.042	0.0	16.162	0.773	0.0	45.496	1.943	0.0	0.519

plot_cal_results(fdr_cal_results)

evaluate_cates(est, X_train=X_train, X_test=X_test)

evaluate_cates(ldml, X_train=X_train, X_test=X_test)

evaluate_cates(fdr, X_train=X_train, X_test=X_test)

Now that we've added the mediator, we can see that all models recover an ATE closer to the DGP for the effect subscription only, and that the performance of all models (nuisance and EconML) has improved. There is still an important level of noise/error, but that is by design, and the CATE models do not capture this noise. Once again the infamous question: what model to use? I mentioned above that the choice also depends on what SHAP values you might want to deal with -- SHAP values from flexible models tend to be richer and capture more complex relationships, compared to a linear model, for example.

So below you can see the SHAP implementation for the three models explored so far, and how they differ in some basic outputs (scatter/dependence, and summary plots).

def cate_shap(estimator=None):
    def cate_fn(X):
        return estimator.effect(X)
    return cate_fn

background = shap.sample(X_train, 700, random_state=0)

explainer_est = shap.Explainer(cate_shap(est), background)
est_shap_vals = explainer_est(X_test)

ExactExplainer explainer: 3001it [10:15,  4.80it/s]                          

explainer_ldml = shap.Explainer(cate_shap(ldml), background)
ldml_shap_vals = explainer_ldml(X_test)

explainer_fdr = shap.Explainer(cate_shap(fdr), background)
fdr_shap_vals = explainer_fdr(X_test)

ExactExplainer explainer: 3001it [09:16,  5.29it/s]                          

models = ['CausalForestDML', 'LinearDRLearner', 'CausalForestDML']
shap_vals = [est_shap_vals, ldml_shap_vals, fdr_shap_vals]

for model, vals in zip(models, shap_vals):
    print(f'{model}:')
    shap.summary_plot(vals)
    print('#'*40, '\n')

CausalForestDML:

######################################## 

LinearDRLearner:

######################################## 

CausalForestDML:

######################################## 

for model, vals in zip(models, shap_vals):
    print(f'{model}: ')
    for col in X_test.columns.tolist():
        shap.plots.scatter(vals[:, col], color=vals)
    print('#'*40, '\n')

CausalForestDML: 

######################################## 

LinearDRLearner: 

######################################## 

CausalForestDML: 

######################################## 

What do the SHAP values show? In the case of EconML, and provided the background data we provided as sample is representative of the larger dataframe, these SHAP values show what is the uplift driven by the relevant feature x distribution interaction compared to the ATE (the zero line in each feature). So this is different to what we usually see with SHAP values applied to other ML problems. In a classic ML problem predicting sales, a positive linear relationship between income and SHAP values for income could be interpreted as "higher income increases sales or predicts higher sales". Here, the interpretation would be "those with higher income are predicted to have a higher conditional effect as a consequence of the treatment". Differently put, that higher sales should come from converting higher income people into the subscription.

The interesting thing to notice is that, despite showing disagreements on the relevant contribution of some features as shown in the summary plot (and earlier difference in CATE performance which we didn't discuss), they all converge on a number of things:

• Income and age are by far the largest drivers of variability in the uplift effect (which is understandable considering the DGP).
• The people that would show the highest uplift from subscription are (i) those 30-40 years old and over, (ii) those with a medium shopping frequency (7-11 transactions), those who have a more moderate or higher app usage frequency (6-7+), (iii) those who live closer to the stores, (iv) those with higher income (70k+), those with moderate/higher SES, (.5-.7+), and those who are sent a moderate/higher number of offers (9+).

So that gives us the basis to create some further segmentation based not only on the features above, but also leveraging the CATEs produced by the model we choose in the end.

Thanks for reading!