[1]:
import pandas as pd
import matplotlib.pyplot as plt

from odtlearn.flow_opt import FlowOPT_IPW, FlowOPT_DM, FlowOPT_DR
from odtlearn.datasets import prescriptive_ex_data

FlowOPT Examples

We will look at the different methods of learning optimal prescriptive trees: inverse probability weighting (IPW), direct method (DM), and doubly robust (DR)

Example 0: Preparing the Data

We will first load a synthetic dataset to use as our running example. In this example, we have precomputed the learned inverse propensity weights (IPW) and predicted outcomes for the direct method (DM). At least one of these values must be computed in order to use any of the prescriptive tree classifiers in ODTLearn. The advantages and disadvantages of using each type of prescriptive tree is given in the subsequent examples, and we refer to Dudik et al. (2011) and Jo et al. (2021) for more details.

The column names in this dataset that correspond to the IPW weights and predicted outcomes for DM are:

  • prob_t_pred_log is the inverse propensity weight learned through logistic regression

  • prob_t_pred_tree is the inverse propensity weight learned through a decision tree

  • linear0 is the predicted outcome under treatment 0 learned through linear regression

  • linear1 is the predicted outcome under treatment 1 learned through linear regression

  • lasso0 is the predicted outcome under treatment 0 learned through lasso regression

  • lasso0 is the predicted outcome under treatment 2 learned through lasso regression

[2]:
# Read data
train_data, test_data  = prescriptive_ex_data()
print(f'shape{train_data.shape}')
train_data = train_data.sample(frac=0.1,random_state=42).reset_index(drop=True) # Taking a small subset of the data for example purposes
test_data = test_data.sample(frac=0.2, random_state=42).reset_index(drop=True)

train_data.head()
shape(500, 31)
[2]:
V1.1 V1.2 V1.3 V1.4 V1.5 V1.6 V1.7 V1.8 V1.9 V1.10 ... y1 t prob_t y prob_t_pred_log prob_t_pred_tree linear0 lasso0 linear1 lasso1
0 0 0 0 0 0 0 0 1 1 1 ... 0.601624 1 0.5 0.601624 0.466833 0.466 0.491137 0.485787 0.574565 0.539485
1 0 0 1 1 1 1 1 1 1 1 ... 0.219783 1 0.5 0.219783 0.483647 0.466 0.152827 0.327543 0.219158 0.433684
2 0 0 0 0 0 0 1 1 1 1 ... 0.454380 0 0.5 0.319519 0.525567 0.534 0.336105 0.401297 0.443978 0.490118
3 0 0 0 1 1 1 1 1 1 1 ... 0.425231 0 0.5 0.390287 0.528118 0.534 0.391595 0.452416 0.434434 0.509328
4 0 1 1 1 1 1 1 1 1 1 ... 0.575465 0 0.5 0.663447 0.540226 0.534 0.642607 0.611545 0.585660 0.590906

5 rows × 31 columns

Note: In practice, users should learn IPW weights and predicted outcomes in the following manner:

  • IPW weights can be learned using a standard machine learning method that predicts the probability of treatment assignments given covariates (e.g., logistic regression), then taking its inverse.

  • Predicted outcomes for DM are learned by training a machine learning method for each treatment t:

    • Take the subset of training data with treatment t

    • Learn a machine learning model that can predict the outcome y under treatment t using that subset of training data through, e.g., lasso regression

    • Predict the outcome y under treatment t using the learned model

Example 1: Inverse Propensity Weighting (IPW)

Inverse Propensity Weighting (IPW) works by creating a pseudo-population where treatment assignments are randomized, allowing us to estimate the average outcome under a learned prescriptive tree.

IPW requires learning propensities, which are the probability that each sample received its assigned treatment. This can be estimated via standard machine learning methods (e.g., logistic regression), where the covariates and assigned treatments are used to create a propensity model that predicts the probability of receiving a treatment.

IPW is only a consistent estimator when the learned propensity model is correctly specified, and may suffer from high variance and sensitivity to small propensity scores (leading to very high inverse propensity weights).

[3]:
X = train_data.iloc[:, 15:20] # Taking a subset of covariates for use with Gurobi restricted license limits
t = train_data["t"] # treatment
y = train_data["y"] # outcome
ipw = train_data["prob_t_pred_tree"] # IPW weights

opt_ipw = FlowOPT_IPW(solver="gurobi", depth=2, time_limit=300)

opt_ipw.fit(X, t, y, ipw) # pass in IPW weights in the fit function
Restricted license - for non-production use only - expires 2027-11-29
Set parameter TimeLimit to value 300
Gurobi Optimizer version 13.0.2 build v13.0.2rc1 (linux64 - "Ubuntu 24.04.4 LTS")

CPU model: AMD EPYC 9V74 80-Core Processor, instruction set [SSE2|AVX|AVX2]
Thread count: 2 physical cores, 4 logical processors, using up to 4 threads

Non-default parameters:
TimeLimit  300

Optimize a model with 1014 rows, 736 columns and 2803 nonzeros (Max)
Model fingerprint: 0xbf6dc368
Model has 50 linear objective coefficients
Variable types: 714 continuous, 22 integer (22 binary)
Coefficient statistics:
  Matrix range     [1e+00, 1e+00]
  Objective range  [4e-01, 2e+00]
  Bounds range     [1e+00, 1e+00]
  RHS range        [1e+00, 1e+00]

Found heuristic solution: objective -0.0000000
Presolve removed 605 rows and 506 columns
Presolve time: 0.00s
Presolved: 409 rows, 230 columns, 1292 nonzeros
Found heuristic solution: objective 27.2886911
Variable types: 212 continuous, 18 integer (18 binary)

Root relaxation: objective 4.557964e+01, 303 iterations, 0.00 seconds (0.00 work units)

    Nodes    |    Current Node    |     Objective Bounds      |     Work
 Expl Unexpl |  Obj  Depth IntInf | Incumbent    BestBd   Gap | It/Node Time

     0     0   45.57964    0    2   27.28869   45.57964  67.0%     -    0s
H    0     0                      31.4587260   45.57964  44.9%     -    0s
     0     0   45.57964    0    4   31.45873   45.57964  44.9%     -    0s
     0     0   45.57964    0    7   31.45873   45.57964  44.9%     -    0s
     0     0   45.57964    0    7   31.45873   45.57964  44.9%     -    0s
     0     0   45.57964    0    7   31.45873   45.57964  44.9%     -    0s
     0     0   45.57964    0    7   31.45873   45.57964  44.9%     -    0s
     0     2   45.57964    0    7   31.45873   45.57964  44.9%     -    0s
*   12    10               4      31.9674681   43.40149  35.8%  20.3    0s
*   13    10               5      32.9958399   43.40149  31.5%  19.0    0s
*   27    15              10      33.5045819   41.89721  25.0%  15.8    0s
*   64     6               9      33.7293800   37.37083  10.8%  12.8    0s

Cutting planes:
  Gomory: 1
  MIR: 19
  Flow cover: 4
  RLT: 4
  Relax-and-lift: 1
  BQP: 1

Explored 73 nodes (1334 simplex iterations) in 0.07 seconds (0.04 work units)
Thread count was 4 (of 4 available processors)

Solution count 7: 33.7294 33.5046 32.9958 ... -0

Optimal solution found (tolerance 1.00e-04)
Best objective 3.372937999418e+01, best bound 3.372937999418e+01, gap 0.0000%
[3]:
FlowOPT_IPW(solver=gurobi,depth=2,time_limit=300,num_threads=None,verbose=False)
[4]:
opt_ipw.print_tree()
fig, ax = plt.subplots(figsize=(4, 4))
opt_ipw.plot_tree(ax=ax, fontsize=10, color_dict={"node": None, "leaves": []})
plt.show()
#########node  1
branch on V2.7
#########node  2
branch on V2.8
#########node  3
branch on V2.6
#########node  4
leaf 0
#########node  5
leaf 1
#########node  6
leaf 0
#########node  7
leaf 1
../_images/notebooks_FlowOPT_9_1.png

Example 2: Direct Method

The Direct Method (DM) works by modeling the outcomes under each treatment directly, which allows for an estimate of the outcome of each data sample under a different treatment assignment.

DM requires learning an outcome model for each treatment in the data. This can be estimated via standard machine learning methods (e.g., lasso regression), where the subset of individuals in the data that received one treatment assignment is used to create the outcome model under that treatment.

DM is only a consistent estimator when the outcome models are correctly specified, and suffers from high bias as it requires extrapolating the outcome beyond the individuals that historically received that treatment.

[5]:
X = train_data.iloc[:, 15:20] # covariates
t = train_data["t"] # treatment
y = train_data["y"] # outcome
y_hat = train_data[["linear0", "linear1"]] # estimated outcomes

opt_dm = FlowOPT_DM(solver="gurobi", depth=2, time_limit=300)

opt_dm.fit(X, t, y, y_hat) # pass in estimated outcomes in the fit function
Set parameter TimeLimit to value 300
Gurobi Optimizer version 13.0.2 build v13.0.2rc1 (linux64 - "Ubuntu 24.04.4 LTS")

CPU model: AMD EPYC 9V74 80-Core Processor, instruction set [SSE2|AVX|AVX2]
Thread count: 2 physical cores, 4 logical processors, using up to 4 threads

Non-default parameters:
TimeLimit  300

Optimize a model with 1414 rows, 1086 columns and 3903 nonzeros (Max)
Model fingerprint: 0xde838f85
Model has 700 linear objective coefficients
Variable types: 1064 continuous, 22 integer (22 binary)
Coefficient statistics:
  Matrix range     [1e+00, 1e+00]
  Objective range  [2e-01, 9e-01]
  Bounds range     [1e+00, 1e+00]
  RHS range        [1e+00, 1e+00]

Found heuristic solution: objective 26.4357477
Presolve removed 634 rows and 534 columns
Presolve time: 0.01s
Presolved: 780 rows, 552 columns, 2510 nonzeros
Variable types: 534 continuous, 18 integer (18 binary)

Root relaxation: objective 2.834794e+01, 272 iterations, 0.00 seconds (0.00 work units)

    Nodes    |    Current Node    |     Objective Bounds      |     Work
 Expl Unexpl |  Obj  Depth IntInf | Incumbent    BestBd   Gap | It/Node Time

     0     0   28.34794    0    2   26.43575   28.34794  7.23%     -    0s
H    0     0                      27.7780800   28.34794  2.05%     -    0s
     0     0   28.34794    0    9   27.77808   28.34794  2.05%     -    0s
H    0     0                      27.9719234   28.34794  1.34%     -    0s
     0     0   28.34794    0    6   27.97192   28.34794  1.34%     -    0s
     0     0   28.34794    0    6   27.97192   28.34794  1.34%     -    0s
     0     0   28.34794    0    8   27.97192   28.34794  1.34%     -    0s
     0     0   28.34794    0    6   27.97192   28.34794  1.34%     -    0s
     0     0   28.34794    0    6   27.97192   28.34794  1.34%     -    0s
     0     2   28.34794    0    6   27.97192   28.34794  1.34%     -    0s
*   23    12               7      28.0371003   28.34794  1.11%  29.8    0s

Cutting planes:
  MIR: 20
  Flow cover: 4
  RLT: 2
  Relax-and-lift: 1

Explored 56 nodes (1719 simplex iterations) in 0.12 seconds (0.08 work units)
Thread count was 4 (of 4 available processors)

Solution count 4: 28.0371 27.9719 27.7781 26.4357

Optimal solution found (tolerance 1.00e-04)
Best objective 2.803710031053e+01, best bound 2.803710031053e+01, gap 0.0000%
[5]:
FlowOPT_DM(solver=gurobi,depth=2,time_limit=300,num_threads=None,verbose=False)
[6]:
opt_dm.print_tree()
fig, ax = plt.subplots(figsize=(4, 4))
opt_dm.plot_tree(ax=ax, fontsize=10, color_dict={"node": None, "leaves": []})
plt.show()
#########node  1
branch on V2.7
#########node  2
branch on V2.8
#########node  3
branch on V2.6
#########node  4
leaf 0
#########node  5
leaf 1
#########node  6
leaf 0
#########node  7
leaf 1
../_images/notebooks_FlowOPT_12_1.png

Example 3: Doubly Robust Method

The Doubly Robust Method (DR) combines the IPW and DM methods. The DR method uses an estimated outcome model (similar to DM), and corrects any bias from the outcome model using IPW.

DR is consistent if either the propensity model OR the outcome model is correctly specified. However, it does require learning both the IPW weights and the outcome models under each treatment.

[7]:
X = train_data.iloc[:, 15:20] # covariates
t = train_data["t"] # treatment
y = train_data["y"] # outcome
ipw = train_data["prob_t_pred_tree"] # IPW weights
y_hat = train_data[["linear0", "linear1"]] # estimated outcomes

opt_dr = FlowOPT_DR(solver="gurobi", depth=2, time_limit=300)

opt_dr.fit(X, t, y, ipw, y_hat) # Pass in both IPW weights and estimated outcomes in the fit function
Set parameter TimeLimit to value 300
Gurobi Optimizer version 13.0.2 build v13.0.2rc1 (linux64 - "Ubuntu 24.04.4 LTS")

CPU model: AMD EPYC 9V74 80-Core Processor, instruction set [SSE2|AVX|AVX2]
Thread count: 2 physical cores, 4 logical processors, using up to 4 threads

Non-default parameters:
TimeLimit  300

Optimize a model with 1414 rows, 1086 columns and 3903 nonzeros (Max)
Model fingerprint: 0x2a6fcff6
Model has 700 linear objective coefficients
Variable types: 1064 continuous, 22 integer (22 binary)
Coefficient statistics:
  Matrix range     [1e+00, 1e+00]
  Objective range  [2e-01, 9e-01]
  Bounds range     [1e+00, 1e+00]
  RHS range        [1e+00, 1e+00]

Found heuristic solution: objective 26.2414664
Presolve removed 634 rows and 534 columns
Presolve time: 0.01s
Presolved: 780 rows, 552 columns, 2510 nonzeros
Variable types: 534 continuous, 18 integer (18 binary)

Root relaxation: objective 2.869269e+01, 282 iterations, 0.00 seconds (0.00 work units)

    Nodes    |    Current Node    |     Objective Bounds      |     Work
 Expl Unexpl |  Obj  Depth IntInf | Incumbent    BestBd   Gap | It/Node Time

     0     0   28.69269    0    4   26.24147   28.69269  9.34%     -    0s
H    0     0                      27.9231182   28.69269  2.76%     -    0s
     0     0   28.69269    0    7   27.92312   28.69269  2.76%     -    0s
H    0     0                      28.0904572   28.69269  2.14%     -    0s
     0     0   28.69269    0    7   28.09046   28.69269  2.14%     -    0s
     0     0   28.69269    0    7   28.09046   28.69269  2.14%     -    0s
     0     0   28.69269    0    7   28.09046   28.69269  2.14%     -    0s
H    0     0                      28.1093786   28.69269  2.08%     -    0s
     0     0   28.69269    0    6   28.10938   28.69269  2.08%     -    0s
     0     0   28.69269    0    6   28.10938   28.69269  2.08%     -    0s
     0     2   28.69269    0    6   28.10938   28.69269  2.08%     -    0s
*    9     7               6      28.2406652   28.67666  1.54%  43.0    0s

Cutting planes:
  Gomory: 2
  MIR: 3
  Flow cover: 2
  RLT: 4

Explored 37 nodes (1415 simplex iterations) in 0.11 seconds (0.08 work units)
Thread count was 4 (of 4 available processors)

Solution count 5: 28.2407 28.1094 28.0905 ... 26.2415

Optimal solution found (tolerance 1.00e-04)
Best objective 2.824066517198e+01, best bound 2.824066517198e+01, gap 0.0000%
[7]:
FlowOPT_DR(solver=gurobi,depth=2,time_limit=300,num_threads=None,verbose=False)
[8]:
opt_dr.print_tree()
fig, ax = plt.subplots(figsize=(4, 4))
opt_dr.plot_tree(ax=ax, fontsize=10, color_dict={"node": None, "leaves": []})
plt.show()
#########node  1
branch on V2.7
#########node  2
branch on V2.9
#########node  3
branch on V2.6
#########node  4
leaf 0
#########node  5
leaf 1
#########node  6
leaf 0
#########node  7
leaf 1
../_images/notebooks_FlowOPT_15_1.png

Example 4: Comparing Performance

It may be of interest to compare the performance of each type of estimator. We can compute summary statistics on the average outcome from each learned tree.

In practice, DR is the typical preferred choice as it is consistent if either IPW or DR methods are consistent. IPW also remains a popular choice since it relies on just one propensity model to learn. We refer to Dudik et al. (2011) and Jo et al. (2021) for more details.

[9]:
# Read test data
print(f'shape{test_data.shape}')
test_data.columns

test_data.head()
shape(20, 27)
[9]:
V1.1 V1.2 V1.3 V1.4 V1.5 V1.6 V1.7 V1.8 V1.9 V1.10 ... V2.8 V2.9 V2.10 y0 y1 t prob_t y prob_t_pred_log prob_t_pred_tree
0 0 0 1 1 1 1 1 1 1 1 ... 1 1 1 0.530883 0.492669 0 0.5 0.530883 0.537961 0.534
1 0 1 1 1 1 1 1 1 1 1 ... 1 1 1 0.411497 0.356187 1 0.5 0.356187 0.470692 0.466
2 0 0 0 0 0 0 1 1 1 1 ... 1 1 1 0.517887 0.501766 0 0.5 0.517887 0.536472 0.534
3 1 1 1 1 1 1 1 1 1 1 ... 1 1 1 0.242120 0.146275 1 0.5 0.146275 0.481979 0.466
4 1 1 1 1 1 1 1 1 1 1 ... 0 0 1 0.649355 0.527296 0 0.5 0.649355 0.544013 0.534

5 rows × 27 columns

[10]:
# Predict
X_test = test_data.iloc[:, 15:20] # covariates
y_0_test = test_data["y0"] # estimated outcome under treatment 0
y_1_test = test_data["y1"] # estimated outcome under treatment 1

predict_ipw = opt_ipw.predict(X_test)
predict_dm = opt_dm.predict(X_test)
predict_dr = opt_dr.predict(X_test)
[11]:
# Check the outcome
def calculate_outcome(y0, y1, predict):
    total_outcome = 0
    for i in range(len(predict)):
        if predict[i] == 0:
            total_outcome += y0[i]
        else:
            total_outcome += y1[i]
    return total_outcome / len(predict)

print("IPW test estimated outcome: ", calculate_outcome(y_0_test, y_1_test, predict_ipw))
print("DM test estimated outcome: ", calculate_outcome(y_0_test, y_1_test, predict_dm))
print("DR test estimated outcome: ", calculate_outcome(y_0_test, y_1_test, predict_dr))
IPW test estimated outcome:  0.5109253363383204
DM test estimated outcome:  0.5109253363383204
DR test estimated outcome:  0.5123693781552128

References

  • Dudík, M., Langford, J., & Li, L. (2011). Doubly robust policy evaluation and learning. In Proceedings of the 28th International Conference on International Conference on Machine Learning (pp. 1097-1104).

  • Jo, N., Aghaei, S., Gómez, A., & Vayanos, P. (2021). Learning optimal prescriptive trees from observational data. arXiv preprint arXiv:2108.13628.

[ ]: