Xiang Ao

Gender Wage Gap Analysis Using US Current Population Survey Data

Xiang Ao — Mon, 17 Nov 2025 05:00:00 GMT

Traditionally Oaxaca-Blinder decomposition is used to decompose the wage gap into explained and unexplained parts. The explained part is due to differences in observable characteristics, such as education, experience, occupation, etc. The unexplained part is due to discrimination or other unobserved factors.

It has a problem that when decomposing, using female as reference group or male as reference group will give different results.

Słoczyński (2015) and Słoczyński (2013) have shown that a regression adjustment type of analysis is equivalent to Oaxaca-Blinder decomposition in a nonstandard way.

Oaxaca-Blinder decomposition

Assume the regression model is different for male and female:

where is gender (female is 1), ’s are person characteristics.

These two equations are different forms of the Oaxaca-Blinder decomposition. The first element is the “unexplained component”, which is the difference in the coefficients, and the second element is the “explained component”, which is the difference in the means of the covariates.

These two equations can give different results for the two components, one has a baseline of male, and other one female.

We can further shows that

That is, the observed difference is the sum of and selection bias. It is also know in the OB decomposition literature as unexplained component and explained component. The former as “discrimination”.

We can simply use Regression Adjustment for . The easiest would be an OLS of on demeaned and and interaction, then the coefficient on is .

He aslo advocates the “population average gender effect” (PAGE):

which is basically ATE, or conditional average effect (CATE), then average over the ’s.

In this version, the two forms of OB decomposition can be and , for men and women.

The nice thing about this estimator is that first, it is consistent with the treatment effect (causal inference) literature, and it is invariant to the choice of reference group.

The second advantage is that we also do not need to assume homogeneity of treatment effect, which is a strong assumption in the Oaxaca-Blinder decomposition.

The third advantage is that we can use nonparametric methods to estimate the treatment effect. Traditional OB decomposition is based on linear regression, which assumes a linear relationship between the covariates and the outcome.

Example

I am using the data from https://github.com/d2cml-ai/CausalAI-Course/tree/main/data.

“The data set we consider is from the 2015 March Supplement of the U.S. Current Population Survey. We select white non-hispanic individuals, aged 25 to 64 years, and working more than 35 hours per week for at least 50 weeks of the year. We exclude self-employed workers; individuals living in group quarters; individuals in the military, agricultural or private household sectors; individuals with inconsistent reports on earnings and employment status; individuals with allocated or missing information in any of the variables used in the analysis; and individuals with hourly wage below 3.”

I am using this data to study gender wage gap (GWP).

We can model GWP using potential outcome framework. Here treatment is gender discrimination, and outcome is logged wage. Treatment is 1 if female, 0 if male.

The ATE would be

where is the potential outcome if treated, and is the potential outcome if not treated.

Regression adjustment

Since this is to estimate ATE, we can use regression adjustment to estimate the ATE. First let’s use de-meaned ’s. Then the coefficient on gender is the ATE.

We use These variables as ’s:

“shs”,“hsg”,“scl”,“clg”,“ad”,“ne”,“mw”,“so”,“we”,“exp1”, “occ”, “occ2”, “ind”, and “ind2”.

They are:

“Less then High School”,“High School Graduate”,“Some College”,“College Graduate”,“Advanced Degree”, “Northeast”,“Midwest”,“South”,“West”,“Experience”, “occupation”, “occupation 2”, “industry”, and “industry 2”.

I decide to use “occ2” and “ind2” as they are broader categories of occupation and industry. We don’t have enough data for finer categories “occ” and “ind”.

“exp1” seems to be the only variable that is continuous, let’s demean it first. Let’s also input “occ2” and “ind2” as factors.

Let’s include demeaned “exp1”, sohs”, “hsg”, “socl”, “clg”, “mw”, “sout”, and “we” first. I leave out “occ2” and “ind2” since there are too many levels to demean.

library(tidyverse)
library(broom)


data <- read_csv("~/Downloads/wage2015_subsample_inference.csv") |> 
  rename(socl = scl, sohs = shs, sout = so) |> 
  mutate(exp_dm = exp1 - mean(exp1, na.rm = TRUE), female=sex, occ2 = factor(occ2), ind2 = factor(ind2)) |> 
  mutate(sohs_dm = sohs - mean(sohs, na.rm = TRUE),
         hsg_dm = hsg - mean(hsg, na.rm = TRUE),
         socl_dm = socl - mean(socl, na.rm = TRUE),
         clg_dm = clg - mean(clg, na.rm = TRUE),
         mw_dm = mw - mean(mw, na.rm = TRUE),
         sout_dm = sout - mean(sout, na.rm = TRUE),
         we_dm = we - mean(we, na.rm = TRUE)) 

glimpse(data)

Rows: 5,150
Columns: 30
$ rownames  10, 12, 15, 18, 19, 30, 43, 44, 47, 71, 73, 77, 84, 89, 96, 1…
$ wage      9.615385, 48.076923, 11.057692, 13.942308, 28.846154, 11.7307…
$ lwage     2.263364, 3.872802, 2.403126, 2.634928, 3.361977, 2.462215, 2…
$ sex       1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1…
$ sohs      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0…
$ hsg       0, 0, 1, 0, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0…
$ socl      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1…
$ clg       1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0…
$ ad        0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0…
$ mw        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0…
$ sout      0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0…
$ we        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0…
$ ne        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1…
$ exp1      7.0, 31.0, 18.0, 25.0, 22.0, 1.0, 42.0, 37.0, 31.0, 4.0, 7.0,…
$ exp2      0.4900, 9.6100, 3.2400, 6.2500, 4.8400, 0.0100, 17.6400, 13.6…
$ exp3      0.343000, 29.791000, 5.832000, 15.625000, 10.648000, 0.001000…
$ exp4      0.24010000, 92.35210000, 10.49760000, 39.06250000, 23.4256000…
$ occ       3600, 3050, 6260, 420, 2015, 1650, 5120, 5240, 4040, 3255, 40…
$ occ2      11, 10, 19, 1, 6, 5, 17, 17, 13, 10, 13, 14, 11, 11, 1, 19, 1…
$ ind       8370, 5070, 770, 6990, 9470, 7460, 7280, 5680, 8590, 8190, 82…
$ ind2      18, 9, 4, 12, 22, 14, 14, 9, 19, 18, 18, 18, 18, 18, 17, 4, 4…
$ exp_dm    -6.760583, 17.239417, 4.239417, 11.239417, 8.239417, -12.7605…
$ female    1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1…
$ sohs_dm   -0.02330097, -0.02330097, -0.02330097, -0.02330097, -0.023300…
$ hsg_dm    -0.2438835, -0.2438835, 0.7561165, -0.2438835, -0.2438835, -0…
$ socl_dm   -0.2780583, -0.2780583, -0.2780583, -0.2780583, -0.2780583, -…
$ clg_dm    0.6823301, 0.6823301, -0.3176699, -0.3176699, 0.6823301, 0.68…
$ mw_dm     -0.2596117, -0.2596117, -0.2596117, -0.2596117, -0.2596117, -…
$ sout_dm   -0.2965049, -0.2965049, -0.2965049, -0.2965049, -0.2965049, -…
$ we_dm     -0.2161165, -0.2161165, -0.2161165, -0.2161165, -0.2161165, -…

# construct matrices for estimation from the data 

# 1. basic model
reg1 <- lm(lwage ~ female, data=data)
tidy(reg1)

# A tibble: 2 × 5
  term        estimate std.error statistic p.value
                         
1 (Intercept)   2.99      0.0107    280.    0     
2 female       -0.0383    0.0160     -2.40  0.0165

reg2 <- lm(lwage ~ female*(exp_dm + sohs_dm + hsg_dm + socl_dm + clg_dm + mw_dm + sout_dm + we_dm), data=data)
tidy(reg2)

# A tibble: 18 × 5
   term           estimate std.error statistic  p.value
                              
 1 (Intercept)     3.02     0.00974    310.    0       
 2 female         -0.115    0.0147      -7.80  7.39e-15
 3 exp_dm          0.00801  0.000957     8.37  7.31e-17
 4 sohs_dm        -0.811    0.0620     -13.1   1.63e-38
 5 hsg_dm         -0.706    0.0350     -20.1   6.09e-87
 6 socl_dm        -0.553    0.0351     -15.7   1.45e-54
 7 clg_dm         -0.256    0.0345      -7.41  1.42e-13
 8 mw_dm           0.0337   0.0280       1.20  2.29e- 1
 9 sout_dm         0.0148   0.0271       0.548 5.84e- 1
10 we_dm           0.0555   0.0290       1.91  5.59e- 2
11 female:exp_dm   0.00212  0.00140      1.51  1.30e- 1
12 female:sohs_dm -0.0180   0.117       -0.154 8.78e- 1
13 female:hsg_dm   0.0398   0.0507       0.786 4.32e- 1
14 female:socl_dm  0.0466   0.0482       0.968 3.33e- 1
15 female:clg_dm   0.0992   0.0468       2.12  3.41e- 2
16 female:mw_dm   -0.124    0.0417      -2.98  2.93e- 3
17 female:sout_dm -0.0530   0.0403      -1.31  1.89e- 1
18 female:we_dm   -0.0655   0.0435      -1.51  1.32e- 1

We see the raw gap is 3.8%. Adjusting for some covariates we have 11% gap.

Or, we can use “marginaleffects” instead of demeaning.

reg3 <- lm(lwage ~ female*(exp1 + sohs + hsg+ socl + clg + mw + sout + we ), data=data)

library(marginaleffects)

avg_comparisons(reg3,
                variables = "female",
                vcov = ~subclass)


 Estimate Std. Error    z Pr(>|z|)    S  2.5 %  97.5 %
   -0.115     0.0147 -7.8   <0.001 47.2 -0.143 -0.0858

Term: female
Type: response
Comparison: 1 - 0

We get the same result as the demeaned regression.

If we want , then we can use the “newdata” argument to specify the treatment group,

avg_comparisons(reg3,
                variables = "female",
                vcov = ~subclass,
                newdata = subset(female == 1))


 Estimate Std. Error     z Pr(>|z|)    S  2.5 %  97.5 %
   -0.113     0.0148 -7.65   <0.001 45.5 -0.142 -0.0844

Term: female
Type: response
Comparison: 1 - 0

Next, we adjust for occupation and industry. We can use “occ2” and “ind2” as factors.

reg4 <- lm(lwage ~ female*(exp1 + sohs + hsg+ socl + clg + mw + sout + we + occ2 + ind2), data=data)
avg_comparisons(reg4,
                variables = "female",
                vcov = ~subclass,
                newdata = subset(female == 1))


 Estimate Std. Error     z Pr(>|z|)    S  2.5 %  97.5 %
   -0.075     0.0162 -4.62   <0.001 18.0 -0.107 -0.0432

Term: female
Type: response
Comparison: 1 - 0

So in the model with full interaction, we have a gap of 7.5%.

Now let’s try nonparametric estimation.

Nonparametric estimation

If we’d like to go for nonparametric, relaxing linearity assumption, then we can use “npcausal”. Note this may take a while, depending on which machine learning packages you decide to include. We can also use other nonparametric estimators, such as TMLE or doubleML.

library(npcausal)

#SL.library <- c("SL.earth","SL.glmnet","SL.glm.interaction", "SL.mean","SL.ranger", "SL.xgboost")


SL.library <- c("SL.earth","SL.glmnet","SL.mean","SL.ranger", "SL.xgboost")

# construct W from model matrix of reg3 
reg3 <- lm(lwage ~ exp_dm + sohs + hsg+ socl + clg + mw + sout + we + occ2 + ind2, data=data)
W <- as.data.frame(model.matrix(reg3)[, -1]) # remove intercept

Y <- data$lwage
A <- data$female
set.seed(123)
aipw_ate <- ate(y=Y, a=A, x=W, nsplits=5, sl.lib=SL.library)


  |                                                                            
  |                                                                      |   0%


  |                                                                            
  |====                                                                  |   5%
  |                                                                            
  |=======                                                               |  10%
  |                                                                            
  |==========                                                            |  15%
  |                                                                            
  |==============                                                        |  20%


  |                                                                            
  |==================                                                    |  25%
  |                                                                            
  |=====================                                                 |  30%


  |                                                                            
  |========================                                              |  35%
  |                                                                            
  |============================                                          |  40%
  |                                                                            
  |================================                                      |  45%
  |                                                                            
  |===================================                                   |  50%
  |                                                                            
  |======================================                                |  55%
  |                                                                            
  |==========================================                            |  60%
  |                                                                            
  |==============================================                        |  65%
  |                                                                            
  |=================================================                     |  70%
  |                                                                            
  |====================================================                  |  75%
  |                                                                            
  |========================================================              |  80%
  |                                                                            
  |============================================================          |  85%
  |                                                                            
  |===============================================================       |  90%
  |                                                                            
  |==================================================================    |  95%
  |                                                                            
  |======================================================================| 100%
     parameter         est         se       ci.ll       ci.ul pval
1      E{Y(0)}  2.99614039 0.01093002  2.97471755  3.01756324    0
2      E{Y(1)}  2.93954753 0.01328396  2.91351096  2.96558410    0
3 E{Y(1)-Y(0)} -0.05659286 0.01598082 -0.08791527 -0.02527045    0

aipw_att <- att(y=Y, a=A, x=W, nsplits=5, sl.lib=SL.library)


  |                                                                            
  |                                                                      |   0%
  |                                                                            
  |=======                                                               |  10%
  |                                                                            
  |==============                                                        |  20%


  |                                                                            
  |=====================                                                 |  30%
  |                                                                            
  |============================                                          |  40%
  |                                                                            
  |===================================                                   |  50%
  |                                                                            
  |==========================================                            |  60%
  |                                                                            
  |=================================================                     |  70%
  |                                                                            
  |========================================================              |  80%


  |                                                                            
  |===============================================================       |  90%
  |                                                                            
  |======================================================================| 100%
      parameter         est         se       ci.ll       ci.ul pval
1      E(Y|A=1)  2.94948490 0.01160804  2.92673314  2.97223666    0
2   E{Y(0)|A=1}  3.01543445 0.01199017  2.99193371  3.03893519    0
3 E{Y-Y(0)|A=1} -0.06594955 0.01392369 -0.09323998 -0.03865913    0

We can see this effect is about -6.6%, similar to the regression adjustment result -7.5%.

More on Correlated Random Effects (CRE)

Xiang Ao — Tue, 11 Nov 2025 05:00:00 GMT

Jeff Wooldridge suggested using Correlated Random Effect (CRE) for binary outcomes in this twitter post: https://x.com/jmwooldridge/status/1986100627454206220

Here I summarize what I learned so far:

linear model

For panel data, the usual set up is:

Wooldridge (2021) paper on DiD shows that these models are the same: 1. Two way fixed effect OLS. Dummy variable approach is the same too. 2. Mundlak’s CRE approach. This includes a pooled OLS with Mundlak device (means of ’s by group), and CRE with random effects.

The idea of Chamberlain’s device is that since we don’t have , we can project onto . Basically, we can replace with projection of onto . Mundlak’s device is a special case, that we put all ’s the same weight, thus the mean of for each is the projection of onto . Replace with that projection, and what is left then by definition is uncorrelated with .

Here is an example:

example 1

library(fixest)
library(bacondecomp)
data("castle")

# Fixed-effects model with individual and time fixed effects
fe_model_twoway <- feols(l_homicide ~ poverty + l_police | state + year, data = castle)
summary(fe_model_twoway)

OLS estimation, Dep. Var.: l_homicide
Observations: 550
Fixed-effects: state: 50,  year: 11
Standard-errors: IID 
          Estimate Std. Error   t value Pr(>|t|)    
poverty  -0.027068   0.013306 -2.034206 0.042471 *  
l_police  0.066033   0.104588  0.631364 0.528098    
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
RMSE: 0.176686     Adj. R2: 0.898978
                 Within R2: 0.009286

# random effect
library(lme4)
re_model <- lmer(l_homicide ~ poverty + l_police +  factor(year) +  (1 | state), data = castle)
summary(re_model)

Linear mixed model fit by REML ['lmerMod']
Formula: l_homicide ~ poverty + l_police + factor(year) + (1 | state)
   Data: castle

REML criterion at convergence: 3.9

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-4.5357 -0.4657  0.0151  0.4601  3.6973 

Random effects:
 Groups   Name        Variance Std.Dev.
 state    (Intercept) 0.30041  0.5481  
 Residual             0.03551  0.1884  
Number of obs: 550, groups:  state, 50

Fixed effects:
                   Estimate Std. Error t value
(Intercept)       0.4395941  0.6019389   0.730
poverty          -0.0007991  0.0120087  -0.067
l_police          0.1665913  0.1015581   1.640
factor(year)2001  0.0225820  0.0379597   0.595
factor(year)2002  0.0005681  0.0389059   0.015
factor(year)2003  0.0475588  0.0395876   1.201
factor(year)2004  0.0419220  0.0411546   1.019
factor(year)2005  0.0612232  0.0435803   1.405
factor(year)2006  0.0806576  0.0418469   1.927
factor(year)2007  0.0790850  0.0409066   1.933
factor(year)2008  0.0401268  0.0414964   0.967
factor(year)2009 -0.0424481  0.0488959  -0.868
factor(year)2010 -0.0959167  0.0470940  -2.037

# Mundlak 
castle2 <- castle |> 
  group_by(state) |>
  mutate(poverty_mean = mean(poverty, na.rm = TRUE), l_police_mean= mean(l_police, na.rm = TRUE))
cre_model <- feols(l_homicide ~ poverty + poverty_mean + l_police + l_police_mean + factor(year), data = castle2)
summary(cre_model)

OLS estimation, Dep. Var.: l_homicide
Observations: 550
Standard-errors: IID 
                  Estimate Std. Error    t value   Pr(>|t|)    
(Intercept)      -7.800873   0.531186 -14.685774  < 2.2e-16 ***
poverty          -0.027068   0.030004  -0.902125 3.6740e-01    
poverty_mean      0.125283   0.030674   4.084272 5.0983e-05 ***
l_police          0.066033   0.235836   0.279996 7.7959e-01    
l_police_mean     1.318037   0.253323   5.202983 2.7987e-07 ***
factor(year)2001  0.033071   0.085349   0.387481 6.9855e-01    
factor(year)2002  0.022870   0.087972   0.259968 7.9499e-01    
factor(year)2003  0.074569   0.089852   0.829907 4.0696e-01    
factor(year)2004  0.079080   0.094152   0.839920 4.0133e-01    
factor(year)2005  0.109800   0.100737   1.089973 2.7622e-01    
factor(year)2006  0.118822   0.095991   1.237857 2.1631e-01    
factor(year)2007  0.115412   0.093475   1.234679 2.1749e-01    
factor(year)2008  0.077190   0.095054   0.812069 4.1711e-01    
factor(year)2009  0.027528   0.114987   0.239400 8.1089e-01    
factor(year)2010 -0.034286   0.110153  -0.311255 7.5573e-01    
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
RMSE: 0.417154   Adj. R2: 0.486344

# Mundlak 
cre_model2 <- lmer(l_homicide ~ poverty + poverty_mean + l_police + l_police_mean + factor(year) +  (1 | state), data = castle2)
summary(cre_model2)

Linear mixed model fit by REML ['lmerMod']
Formula: l_homicide ~ poverty + poverty_mean + l_police + l_police_mean +  
    factor(year) + (1 | state)
   Data: castle2

REML criterion at convergence: -30.1

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-4.5991 -0.4336  0.0219  0.4307  3.6228 

Random effects:
 Groups   Name        Variance Std.Dev.
 state    (Intercept) 0.14872  0.3856  
 Residual             0.03518  0.1876  
Number of obs: 550, groups:  state, 50

Fixed effects:
                 Estimate Std. Error t value
(Intercept)      -7.80087    1.60978  -4.846
poverty          -0.02707    0.01331  -2.034
poverty_mean      0.12528    0.02360   5.309
l_police          0.06603    0.10459   0.631
l_police_mean     1.31804    0.30140   4.373
factor(year)2001  0.03307    0.03785   0.874
factor(year)2002  0.02287    0.03901   0.586
factor(year)2003  0.07457    0.03985   1.871
factor(year)2004  0.07908    0.04175   1.894
factor(year)2005  0.10980    0.04467   2.458
factor(year)2006  0.11882    0.04257   2.791
factor(year)2007  0.11541    0.04145   2.784
factor(year)2008  0.07719    0.04215   1.831
factor(year)2009  0.02753    0.05099   0.540
factor(year)2010 -0.03429    0.04885  -0.702

In this example, fixed effect model, CRE1 and CRE2 all give the same coefficient on poverty.

As we know, RE model assumes is uncorrelated with , while CRE and FE allow for correlation between and .

nonlinear model

What I did not realize before is that this is only valid for linear case. For OLS to be consistent, we only need contemporaneous exogeneity, which means . This is not the case for nonlinear cases, such as logit or probit models.

Binary outcome

For binary outcomes, the model is:

I used to think it’s the best to use conditional logit. However, it needs “conditional independence” (serial independence could be a better name) to be consistent. That is, we have to assume the series of is conditionally independent of each other given and . This is a strong assumption. If we think we have serial correlation in the error term, then this would not hold. The other disadvantage of conditional logit (also called fixed effect logit) is that ’s are not estimated, it is just a nuisance parameter. Any partial effects cannot be calculated, because any partial effects are functions of ’s.

Instead, Wooldridge suggests using CRE probit. Basically adding Mundlak device (group means of ’s) to the model, and then use pooled probit. Note that we do not use RE probit or logit, since that would need conditional independence too.

For nonlinear panel data with unobserved heterogeneity (meaning there is that we don’t observe), we have a few assumptions.

First, strict exogeneity, This means that the distribution of only depends on , given .

Second, conditional independence,

This means that the distribution of joint distriubtion of ’s can be modeled by each independently, given and . We need this in most nonlinear cases, but not in linear case.

Then we need to specify . The random effect assumption is saying , that is, is independent of . This is a strong assumption. We should try to avoid this assumption.

Now, most nonlinear models make the second assumption, in the panel setting, when we have . That includes RE model, conditional logit. However, that does not include pooled probit/logit, or CRE with pooled logit/probit. That is why Wooldridge recommend using CRE with pooled probit.

Example 2

# generate a binary outcome, high_homocide

castle2 <- castle2 |> 
  mutate(high_homicide = ifelse(l_homicide > mean(l_homicide, na.rm = TRUE), 1, 0))

# CRE pooled probit
cre_probit <- glm(high_homicide ~ poverty + poverty_mean + l_police + l_police_mean + factor(year), data = castle2, family = binomial(link = "probit"))
# Compute Average Marginal Effects (AMEs)
library(marginaleffects)
ame_cre_probit <- avg_slopes(cre_probit)
ame_cre_probit


          Term    Contrast Estimate Std. Error       z Pr(>|z|)   S     2.5 %
 l_police      dY/dX        0.01924     0.2622  0.0734   0.9415 0.1 -0.494686
 l_police_mean dY/dX       -0.00304     0.2824 -0.0108   0.9914 0.0 -0.556572
 poverty       dY/dX       -0.06477     0.0334 -1.9417   0.0522 4.3 -0.130144
 poverty_mean  dY/dX        0.06620     0.0341  1.9424   0.0521 4.3 -0.000597
 year          2001 - 2000  0.06596     0.0990  0.6659   0.5054 1.0 -0.128159
 year          2002 - 2000  0.09470     0.1019  0.9293   0.3527 1.5 -0.105016
 year          2003 - 2000  0.22958     0.1020  2.2518   0.0243 5.4  0.029753
 year          2004 - 2000  0.15492     0.1082  1.4315   0.1523 2.7 -0.057193
 year          2005 - 2000  0.22388     0.1133  1.9766   0.0481 4.4  0.001880
 year          2006 - 2000  0.26184     0.1068  2.4526   0.0142 6.1  0.052595
 year          2007 - 2000  0.24899     0.1048  2.3756   0.0175 5.8  0.043564
 year          2008 - 2000 -0.03718     0.1099 -0.3383   0.7352 0.4 -0.252646
 year          2009 - 2000 -0.02471     0.1345 -0.1837   0.8542 0.2 -0.288280
 year          2010 - 2000 -0.21753     0.1186 -1.8339   0.0667 3.9 -0.450000
   97.5 %
 0.533162
 0.550499
 0.000608
 0.132987
 0.260070
 0.294406
 0.429403
 0.367039
 0.445887
 0.471080
 0.454421
 0.178276
 0.238865
 0.014950

Type: response

# linear model
cre_lm <- lm(high_homicide ~ poverty + poverty_mean + l_police + l_police_mean + factor(year), data = castle2)
ame_lm <- avg_slopes(cre_lm)
ame_lm


          Term    Contrast Estimate Std. Error       z Pr(>|z|)   S    2.5 %
 l_police      dY/dX        0.01234     0.2666  0.0463   0.9631 0.1 -0.51010
 l_police_mean dY/dX        0.00427     0.2865  0.0149   0.9881 0.0 -0.55724
 poverty       dY/dX       -0.06406     0.0339 -1.8891   0.0589 4.1 -0.13053
 poverty_mean  dY/dX        0.06578     0.0347  1.8975   0.0578 4.1 -0.00217
 year          2001 - 2000  0.06384     0.0965  0.6618   0.5081 1.0 -0.12522
 year          2002 - 2000  0.09082     0.0994  0.9134   0.3610 1.5 -0.10406
 year          2003 - 2000  0.22452     0.1016  2.2108   0.0270 5.2  0.02548
 year          2004 - 2000  0.14789     0.1064  1.3898   0.1646 2.6 -0.06068
 year          2005 - 2000  0.21659     0.1139  1.9023   0.0571 4.1 -0.00656
 year          2006 - 2000  0.25687     0.1085  2.3677   0.0179 5.8  0.04423
 year          2007 - 2000  0.24413     0.1056  2.3108   0.0208 5.6  0.03706
 year          2008 - 2000 -0.04737     0.1074 -0.4410   0.6592 0.6 -0.25794
 year          2009 - 2000 -0.03438     0.1300 -0.2645   0.7914 0.3 -0.28910
 year          2010 - 2000 -0.20933     0.1245 -1.6814   0.0927 3.4 -0.45335
 97.5 %
 0.5348
 0.5658
 0.0024
 0.1337
 0.2529
 0.2857
 0.4236
 0.3565
 0.4397
 0.4695
 0.4512
 0.1632
 0.2203
 0.0347

Type: response

# CRE RE probit
cre_re_probit <- glmer(high_homicide ~ poverty + poverty_mean + l_police + l_police_mean + factor(year) +  (1 | state), data = castle2, family = binomial(link = "probit"))
ame_cre_re_probit <- avg_slopes(cre_re_probit)
ame_cre_re_probit


          Term    Contrast Estimate Std. Error       z Pr(>|z|)   S    2.5 %
 l_police      dY/dX        0.01927     0.2499  0.0771   0.9385 0.1 -0.47055
 l_police_mean dY/dX       -0.00307     0.2720 -0.0113   0.9910 0.0 -0.53625
 poverty       dY/dX       -0.06477     0.0336 -1.9275   0.0539 4.2 -0.13063
 poverty_mean  dY/dX        0.06619     0.0343  1.9285   0.0538 4.2 -0.00108
 year          2001 - 2000  0.06595     0.0995  0.6631   0.5072 1.0 -0.12898
 year          2002 - 2000  0.09469     0.1021  0.9279   0.3535 1.5 -0.10532
 year          2003 - 2000  0.22958     0.1018  2.2554   0.0241 5.4  0.03007
 year          2004 - 2000  0.15492     0.1085  1.4278   0.1534 2.7 -0.05775
 year          2005 - 2000  0.22388     0.1136  1.9714   0.0487 4.4  0.00130
 year          2006 - 2000  0.26184     0.1068  2.4524   0.0142 6.1  0.05258
 year          2007 - 2000  0.24899     0.1047  2.3784   0.0174 5.8  0.04380
 year          2008 - 2000 -0.03719     0.1101 -0.3378   0.7355 0.4 -0.25298
 year          2009 - 2000 -0.02471     0.1346 -0.1836   0.8543 0.2 -0.28848
 year          2010 - 2000 -0.21752     0.1188 -1.8308   0.0671 3.9 -0.45039
  97.5 %
 0.50910
 0.53011
 0.00109
 0.13347
 0.26089
 0.29471
 0.42908
 0.36759
 0.44646
 0.47110
 0.45418
 0.17860
 0.23906
 0.01535

Type: response

Ultimately we are interested in the partial effect, therefore we can compare across models. In this case, we should prefer average marginal effect by CRE pooled probit.

Using NumPyro

Xiang Ao — Fri, 17 Oct 2025 04:00:00 GMT

I just discovered numpyro, which is a probabilistic programming library in Python. It can be coupled with Jax and is very powerful.

Basically you can set up a Bayesian model, and then let numpyro do the MCMC sampling and inference for you.

Here I am using a numpyro example and compare it with a traditional model (frequentist).

Example 1

I am using the numpyro example here: https://num.pyro.ai/en/stable/tutorials/bayesian_hierarchical_linear_regression.html

The data set is from https://www.kaggle.com/c/osic-pulmonary-fibrosis-progression

“Pulmonary fibrosis is a disorder with no known cause and no known cure, created by scarring of the lungs. In this competition, we were asked to predict a patient’s severity of decline in lung function. Lung function is assessed based on output from a spirometer, which measures the forced vital capacity (FVC), i.e. the volume of air exhaled.

In medical applications, it is useful to evaluate a model’s confidence in its decisions. Accordingly, the metric used to rank the teams was designed to reflect both the accuracy and certainty of each prediction.”

I read it in R first.

library(tidyverse)
# read in csv file
library(readr)

data <- read_csv("./osic_pulmonary_fibrosis.csv") |> 
  arrange(Patient, Weeks)

data

# A tibble: 1,549 × 7
   Patient                   Weeks   FVC Percent   Age Sex   SmokingStatus
                                       
 1 ID00007637202177411956430    -4  2315    58.3    79 Male  Ex-smoker    
 2 ID00007637202177411956430     5  2214    55.7    79 Male  Ex-smoker    
 3 ID00007637202177411956430     7  2061    51.9    79 Male  Ex-smoker    
 4 ID00007637202177411956430     9  2144    54.0    79 Male  Ex-smoker    
 5 ID00007637202177411956430    11  2069    52.1    79 Male  Ex-smoker    
 6 ID00007637202177411956430    17  2101    52.9    79 Male  Ex-smoker    
 7 ID00007637202177411956430    29  2000    50.3    79 Male  Ex-smoker    
 8 ID00007637202177411956430    41  2064    51.9    79 Male  Ex-smoker    
 9 ID00007637202177411956430    57  2057    51.8    79 Male  Ex-smoker    
10 ID00009637202177434476278     8  3660    85.3    69 Male  Ex-smoker    
# ℹ 1,539 more rows

a random effect model

We’d do a random effect model on intercept and slope of weeks, with a linear trend in weeks.

# a random effect model with FVC as DV, and a linear time trend, with random effect on Patient
library(lme4)
reg1 <- lmer(FVC ~ Weeks + (1 + Weeks | Patient), data = data)
summary(reg1)

Linear mixed model fit by REML ['lmerMod']
Formula: FVC ~ Weeks + (1 + Weeks | Patient)
   Data: data

REML criterion at convergence: 20891.4

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-9.3209 -0.4257  0.0063  0.4458  5.6990 

Random effects:
 Groups   Name        Variance  Std.Dev. Corr 
 Patient  (Intercept) 686691.11 828.668       
          Weeks           25.88   5.087  -0.14
 Residual              18592.11 136.353       
Number of obs: 1549, groups:  Patient, 176

Fixed effects:
             Estimate Std. Error t value
(Intercept) 2810.2843    62.9497  44.643
Weeks         -4.2629     0.4377  -9.739

Correlation of Fixed Effects:
      (Intr)
Weeks -0.174

Then we do the same with a Bayesian hierarchical model. Should have similar results.

hierarchical model in numpyro

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

train = pd.read_csv(
    "https://gist.githubusercontent.com/ucals/"
    "2cf9d101992cb1b78c2cdd6e3bac6a4b/raw/"
    "43034c39052dcf97d4b894d2ec1bc3f90f3623d9/"
    "osic_pulmonary_fibrosis.csv"
)


from jax import random

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive

assert numpyro.__version__.startswith("0.19.0")

def model(patient_code, Weeks, FVC_obs=None):
    μ_α = numpyro.sample("μ_α", dist.Normal(0.0, 500.0))
    σ_α = numpyro.sample("σ_α", dist.HalfNormal(100.0))
    μ_β = numpyro.sample("μ_β", dist.Normal(0.0, 3.0))
    σ_β = numpyro.sample("σ_β", dist.HalfNormal(3.0))

    n_patients = len(np.unique(patient_code))

    with numpyro.plate("plate_i", n_patients):
        α = numpyro.sample("α", dist.Normal(μ_α, σ_α))
        β = numpyro.sample("β", dist.Normal(μ_β, σ_β))

    σ = numpyro.sample("σ", dist.HalfNormal(100.0))
    FVC_est = α[patient_code] + β[patient_code] * Weeks

    with numpyro.plate("data", len(patient_code)):
        numpyro.sample("obs", dist.Normal(FVC_est, σ), obs=FVC_obs)
        
from sklearn.preprocessing import LabelEncoder

patient_encoder = LabelEncoder()
train["patient_code"] = patient_encoder.fit_transform(train["Patient"].values)

FVC_obs = train["FVC"].values
Weeks = train["Weeks"].values
patient_code = train["patient_code"].values


nuts_kernel = NUTS(model)

mcmc = MCMC(nuts_kernel, num_samples=2000, num_warmup=2000, progress_bar=False)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, patient_code, Weeks, FVC_obs=FVC_obs)

posterior_samples = mcmc.get_samples()
mcmc.print_summary()


                mean       std    median      5.0%     95.0%     n_eff     r_hat
      α[0]   2181.96     60.74   2182.26   2080.02   2276.12   3478.15      1.00
      α[1]   3783.49     73.96   3784.85   3667.70   3913.58   4928.66      1.00
      α[2]   3265.39     58.09   3264.75   3175.63   3366.59   4088.11      1.00
      α[3]   3478.90     66.67   3477.34   3371.61   3589.76   3993.87      1.00
      α[4]   3700.88    118.15   3699.46   3507.03   3892.74   4173.49      1.00
      α[5]   3724.77     59.33   3725.87   3623.35   3817.91   5760.40      1.00
      α[6]   3050.36     87.13   3052.33   2901.09   3187.95   4301.39      1.00
      α[7]   2151.55     78.65   2149.32   2031.82   2287.03   4117.25      1.00
      α[8]   2344.04     89.13   2344.16   2197.93   2493.25   4350.38      1.00
      α[9]   1496.10     59.90   1495.72   1400.21   1595.97   3821.66      1.00
     α[10]   3147.59     59.34   3146.73   3054.35   3247.87   3823.34      1.00
     α[11]   3061.16    134.88   3059.77   2859.88   3302.93   4108.39      1.00
     α[12]   2776.40     81.49   2777.66   2638.22   2904.36   5336.77      1.00
     α[13]   2681.60    125.34   2682.50   2470.12   2876.24   5991.18      1.00
     α[14]   5195.74    116.71   5196.11   5008.86   5385.35   4844.13      1.00
     α[15]   2780.42     75.71   2780.33   2663.12   2909.89   4404.98      1.00
     α[16]   3924.87     78.16   3923.05   3795.71   4053.21   4959.80      1.00
     α[17]   3613.71     61.09   3613.50   3513.57   3715.78   3498.32      1.00
     α[18]   3276.12     66.24   3276.51   3162.19   3382.03   2883.28      1.00
     α[19]   1477.57     69.84   1478.19   1352.44   1579.33   3962.45      1.00
     α[20]   1703.49     56.27   1704.27   1612.68   1796.76   2952.87      1.00
     α[21]   2984.97    218.20   2988.79   2625.00   3337.65   5150.51      1.00
     α[22]   3392.43     93.03   3391.44   3252.03   3554.48   3935.83      1.00
     α[23]   3637.04    103.11   3639.29   3468.56   3799.16   4211.06      1.00
     α[24]   2483.06    144.52   2483.37   2247.02   2721.52   5069.97      1.00
     α[25]   3453.93     66.54   3455.23   3345.72   3565.21   4795.24      1.00
     α[26]   2834.53     75.77   2834.91   2716.25   2961.84   4413.27      1.00
     α[27]   2419.90     83.80   2420.70   2289.84   2568.34   4094.23      1.00
     α[28]   2495.48     68.58   2494.98   2379.58   2601.87   5322.98      1.00
     α[29]   1828.49    109.39   1829.96   1650.90   2013.15   3854.51      1.00
     α[30]   2387.46     59.38   2387.19   2291.45   2484.52   4065.98      1.00
     α[31]   3258.15     64.09   3259.99   3156.46   3366.40   4339.93      1.00
     α[32]   1885.49     63.12   1886.79   1779.17   1986.83   3903.91      1.00
     α[33]   2946.67    109.14   2945.00   2761.04   3113.03   4230.53      1.00
     α[34]   3373.80     62.44   3373.85   3277.11   3478.82   4635.49      1.00
     α[35]   2537.18     72.11   2536.82   2427.51   2662.56   4037.66      1.00
     α[36]   3219.40     92.97   3219.59   3067.16   3361.41   3682.34      1.00
     α[37]   3578.86     64.44   3579.33   3473.77   3685.27   4499.79      1.00
     α[38]   4806.80     67.82   4806.16   4704.59   4926.04   4021.18      1.00
     α[39]   3142.70    128.24   3140.92   2944.80   3357.44   4483.11      1.00
     α[40]   2303.84     64.02   2304.47   2191.04   2403.73   4225.49      1.00
     α[41]   2567.85    104.45   2568.37   2408.97   2752.63   4547.57      1.00
     α[42]   2746.64    132.81   2749.36   2525.71   2960.18   4787.09      1.00
     α[43]   2856.82     79.62   2858.73   2734.38   2990.30   5177.49      1.00
     α[44]   2553.31    161.52   2556.90   2279.42   2807.01   4763.51      1.00
     α[45]   2028.61     99.29   2029.85   1870.11   2198.04   4007.81      1.00
     α[46]   2415.84     85.37   2412.85   2274.98   2551.99   4389.82      1.00
     α[47]   2559.61     91.70   2560.16   2415.14   2713.06   4703.63      1.00
     α[48]   2903.48     67.59   2903.82   2783.52   3003.76   3074.21      1.00
     α[49]   2916.36     61.66   2917.28   2815.96   3018.28   3231.64      1.00
     α[50]   2433.25     62.36   2433.10   2333.02   2535.01   4112.43      1.00
     α[51]   1874.57    149.88   1874.53   1636.37   2125.48   4466.82      1.00
     α[52]   3885.27    152.80   3884.64   3634.18   4129.37   4764.01      1.00
     α[53]   2453.30     74.16   2454.34   2332.25   2572.95   4032.64      1.00
     α[54]   2289.13     98.67   2291.88   2113.48   2432.82   4228.26      1.00
     α[55]   1968.45    148.75   1967.96   1727.68   2223.98   4738.61      1.00
     α[56]   2448.41    168.45   2449.80   2177.20   2731.69   4271.57      1.00
     α[57]   2284.98     61.21   2285.08   2188.96   2386.97   5302.78      1.00
     α[58]   1690.54     75.88   1692.69   1572.17   1820.44   4384.45      1.00
     α[59]   3652.09    175.83   3651.85   3371.15   3952.49   5137.58      1.00
     α[60]   2659.02     72.32   2658.87   2540.28   2776.74   4343.19      1.00
     α[61]   3162.24     59.04   3164.49   3066.44   3260.57   3628.39      1.00
     α[62]   2697.13     68.00   2697.47   2589.83   2809.31   3821.55      1.00
     α[63]   3943.73    119.42   3946.85   3729.60   4121.03   4423.43      1.00
     α[64]   3339.19     64.10   3340.30   3229.28   3435.80   3789.94      1.00
     α[65]   4422.48     85.06   4424.31   4280.59   4564.35   6713.01      1.00
     α[66]   4017.49     92.23   4016.28   3868.26   4169.53   4683.90      1.00
     α[67]   4089.30     82.06   4091.63   3959.81   4223.71   4983.39      1.00
     α[68]   3549.15    138.06   3549.95   3327.27   3776.16   3890.86      1.00
     α[69]   2901.59     73.02   2900.42   2790.57   3030.32   4683.16      1.00
     α[70]   2181.65    219.45   2175.94   1841.71   2541.47   4550.12      1.00
     α[71]   4770.96     74.12   4771.11   4654.07   4898.68   4209.58      1.00
     α[72]   3148.67     68.49   3149.51   3038.27   3261.98   4719.43      1.00
     α[73]   2568.25     84.10   2568.06   2422.14   2691.78   4620.57      1.00
     α[74]   2071.70    135.79   2069.03   1851.62   2289.57   4258.39      1.00
     α[75]   2805.05    126.29   2808.18   2587.92   3004.50   4962.74      1.00
     α[76]   2851.84     84.64   2854.08   2708.30   2982.13   4178.09      1.00
     α[77]   3069.15     75.47   3067.51   2947.36   3191.41   4765.39      1.00
     α[78]   2473.79    134.71   2472.45   2246.52   2689.86   4316.09      1.00
     α[79]   3176.65    112.70   3181.20   2992.57   3366.85   4780.32      1.00
     α[80]   2918.38     75.56   2919.69   2795.13   3045.32   4288.53      1.00
     α[81]   4380.73     73.78   4380.60   4256.85   4500.40   4449.85      1.00
     α[82]   1901.83     73.17   1900.47   1776.18   2016.42   3629.52      1.00
     α[83]   3939.29     74.29   3939.85   3822.65   4061.64   4840.37      1.00
     α[84]   2302.36     64.27   2301.19   2197.80   2406.84   3878.34      1.00
     α[85]   2803.07     66.43   2803.56   2706.09   2922.83   4427.60      1.00
     α[86]   4430.48     94.38   4430.86   4275.12   4581.68   4876.03      1.00
     α[87]   2862.12    122.95   2862.95   2663.85   3066.27   5159.40      1.00
     α[88]   2831.12     59.21   2830.93   2738.42   2929.85   3505.62      1.00
     α[89]   3174.82     69.33   3175.22   3065.44   3291.17   2897.17      1.00
     α[90]   2907.47     63.66   2908.82   2807.63   3013.57   3894.72      1.00
     α[91]   2321.38     73.67   2321.39   2197.67   2435.79   4540.13      1.00
     α[92]   2534.82     92.63   2533.36   2380.92   2685.90   4466.08      1.00
     α[93]   5897.89     61.90   5898.57   5798.89   6001.60   3792.77      1.00
     α[94]   1493.36     77.54   1494.47   1357.99   1613.48   4334.43      1.00
     α[95]   2553.11     64.25   2552.00   2446.32   2654.63   5198.04      1.00
     α[96]   2703.26     61.69   2704.33   2597.79   2800.33   3588.30      1.00
     α[97]   1254.25     84.35   1254.03   1111.30   1386.08   4096.89      1.00
     α[98]   1977.65     81.96   1975.93   1850.62   2117.53   4584.07      1.00
     α[99]   3626.00     94.04   3624.16   3453.35   3765.02   6152.51      1.00
    α[100]   1960.98     74.47   1960.78   1839.09   2081.26   5508.20      1.00
    α[101]   3618.52     56.55   3618.62   3520.58   3705.30   3322.28      1.00
    α[102]   1974.77     65.12   1974.23   1860.57   2073.28   3503.83      1.00
    α[103]   2105.64    128.71   2105.25   1886.33   2306.98   6616.46      1.00
    α[104]   2910.32     68.85   2910.23   2789.15   3014.76   4584.25      1.00
    α[105]   1401.37     63.49   1401.44   1303.52   1512.56   4920.60      1.00
    α[106]   1840.26     65.45   1840.61   1734.82   1946.79   4199.05      1.00
    α[107]   1381.50     80.50   1379.99   1243.57   1515.55   5144.95      1.00
    α[108]   3918.88    134.89   3916.91   3674.47   4118.55   4558.49      1.00
    α[109]   2231.49    103.29   2231.06   2077.49   2414.70   4677.99      1.00
    α[110]   1565.15     66.49   1567.76   1456.32   1674.67   3762.40      1.00
    α[111]   2075.41     87.26   2076.29   1941.11   2223.85   4193.44      1.00
    α[112]   1065.46     73.31   1067.54    938.10   1176.44   4130.45      1.00
    α[113]   2863.31     68.76   2865.82   2749.60   2974.66   4549.16      1.00
    α[114]   2517.66    146.55   2517.04   2289.99   2768.97   4708.89      1.00
    α[115]   2973.36     73.24   2973.90   2847.22   3086.66   4239.57      1.00
    α[116]   4160.73     70.88   4160.89   4043.80   4277.77   4756.33      1.00
    α[117]   1981.23     59.08   1981.86   1887.80   2081.02   3847.90      1.00
    α[118]   1402.20     63.46   1401.69   1303.28   1512.53   3938.85      1.00
    α[119]   2096.95    108.78   2098.23   1910.34   2266.26   4640.21      1.00
    α[120]   2363.06     80.42   2361.92   2237.11   2505.73   3919.36      1.00
    α[121]   4055.78     59.60   4055.97   3955.94   4151.83   3641.06      1.00
    α[122]   2289.47     66.27   2286.09   2188.57   2405.43   4131.95      1.00
    α[123]   2191.83     74.29   2192.15   2071.21   2315.32   3979.08      1.00
    α[124]   1923.23     66.87   1922.47   1815.24   2032.27   3807.17      1.00
    α[125]   2250.75    100.50   2249.12   2086.16   2409.58   5053.45      1.00
    α[126]   2915.61     75.51   2915.58   2794.77   3037.53   5279.26      1.00
    α[127]   3239.20     64.65   3237.52   3144.70   3354.14   4017.07      1.00
    α[128]   3222.20     66.91   3222.45   3109.25   3328.30   3365.91      1.00
    α[129]   1859.73     61.92   1858.81   1758.87   1964.10   4568.50      1.00
    α[130]   2614.03    115.54   2615.68   2420.21   2804.06   4947.98      1.00
    α[131]   3910.90     81.49   3911.47   3779.04   4048.18   4029.07      1.00
    α[132]   2618.84     63.93   2618.52   2509.26   2720.43   4529.40      1.00
    α[133]   1617.22    126.03   1614.94   1406.29   1819.48   4226.61      1.00
    α[134]   2891.23    130.63   2892.29   2684.80   3108.29   4700.87      1.00
    α[135]   3080.86    116.70   3078.24   2867.59   3246.76   5356.82      1.00
    α[136]   2736.10     65.37   2737.26   2632.36   2845.87   4798.27      1.00
    α[137]   2616.46     70.43   2617.43   2499.96   2731.58   4693.34      1.00
    α[138]   3931.41    113.71   3931.37   3744.14   4120.55   6092.02      1.00
    α[139]   3540.65     66.25   3541.71   3435.03   3652.82   4975.09      1.00
    α[140]   2319.48     64.50   2319.73   2214.62   2429.24   3795.68      1.00
    α[141]   2933.86    114.23   2931.55   2746.68   3115.25   4674.73      1.00
    α[142]   2405.33     99.42   2404.73   2248.27   2568.26   4378.19      1.00
    α[143]   1608.70     94.94   1607.20   1452.57   1764.13   4365.56      1.00
    α[144]   2065.78     66.48   2066.25   1962.33   2177.87   4630.61      1.00
    α[145]   2608.69     80.14   2607.56   2476.55   2740.32   5272.79      1.00
    α[146]   4202.96    105.58   4201.76   4034.35   4378.06   5225.38      1.00
    α[147]   2052.41     75.69   2052.85   1930.84   2175.75   4737.47      1.00
    α[148]   2888.84     66.31   2888.21   2785.91   3000.87   4069.44      1.00
    α[149]   3538.52    130.67   3538.85   3311.72   3734.16   5014.65      1.00
    α[150]   2671.02    152.86   2672.30   2421.11   2920.64   5058.42      1.00
    α[151]   1734.95    130.10   1735.16   1529.26   1966.11   4134.36      1.00
    α[152]   3036.62     77.43   3037.41   2913.63   3163.50   5443.08      1.00
    α[153]   2280.62     80.84   2280.98   2157.90   2419.70   6519.05      1.00
    α[154]   2535.44     67.69   2535.98   2424.18   2645.61   4251.41      1.00
    α[155]   4401.42    142.25   4403.37   4172.91   4646.13   4468.32      1.00
    α[156]   3468.31    103.94   3469.87   3318.57   3657.42   4614.79      1.00
    α[157]   2300.55     74.07   2301.65   2186.87   2430.85   4202.74      1.00
    α[158]   3354.23    157.23   3350.28   3104.44   3608.42   4578.01      1.00
    α[159]   3142.29     62.26   3142.13   3032.66   3240.83   4841.13      1.00
    α[160]   2204.44     70.52   2202.80   2089.87   2320.99   4692.09      1.00
    α[161]   3737.23     61.33   3736.79   3634.25   3833.27   3300.82      1.00
    α[162]   3000.90     86.93   3000.32   2858.64   3146.88   4496.44      1.00
    α[163]   3965.14    100.36   3964.14   3793.29   4120.08   3782.56      1.00
    α[164]   2020.35    142.58   2021.99   1791.25   2264.88   5086.22      1.00
    α[165]   1580.13     64.85   1579.95   1486.46   1700.85   2607.56      1.00
    α[166]   3250.36     94.84   3251.18   3103.08   3407.85   4652.29      1.00
    α[167]   2455.41     89.25   2455.88   2307.87   2606.59   4002.42      1.00
    α[168]   1770.46     70.57   1770.27   1659.72   1888.74   4205.60      1.00
    α[169]   2892.62     76.82   2890.41   2774.02   3024.54   3921.34      1.00
    α[170]   3357.47     69.56   3356.52   3252.30   3478.19   4614.54      1.00
    α[171]   2836.92     67.01   2839.11   2729.38   2944.73   4768.27      1.00
    α[172]   2828.49     81.06   2828.37   2687.76   2949.13   4523.36      1.00
    α[173]   1989.75     77.96   1989.84   1854.52   2106.57   3484.11      1.00
    α[174]   3007.70     85.53   3010.05   2869.75   3152.74   4488.10      1.00
    α[175]   2931.63     67.48   2931.62   2813.62   3032.52   4880.61      1.00
      β[0]     -3.49      2.13     -3.52     -6.86      0.11   3777.86      1.00
      β[1]     -8.02      2.36     -8.00    -11.68     -4.02   4356.88      1.00
      β[2]    -14.27      2.28    -14.26    -17.84    -10.21   3486.20      1.00
      β[3]     -4.27      2.31     -4.27     -7.77     -0.22   4751.57      1.00
      β[4]     -7.13      2.19     -7.13    -10.57     -3.43   4758.71      1.00
      β[5]    -11.69      2.32    -11.69    -15.47     -7.84   4411.74      1.00
      β[6]     -6.71      2.27     -6.70    -10.40     -3.02   4800.55      1.00
      β[7]     -5.62      2.20     -5.63     -9.54     -2.30   4401.94      1.00
      β[8]     -6.81      2.27     -6.88    -10.33     -2.83   4538.73      1.00
      β[9]     -3.06      2.28     -3.04     -6.67      0.83   5195.12      1.00
     β[10]     -9.32      2.33     -9.33    -13.17     -5.60   3699.40      1.00
     β[11]     -3.07      2.15     -3.07     -6.56      0.49   4320.79      1.00
     β[12]     -5.17      2.26     -5.14     -8.90     -1.54   4824.94      1.00
     β[13]     -4.64      2.92     -4.56     -9.56     -0.14   4782.43      1.00
     β[14]    -11.04      2.19    -11.04    -14.47     -7.42   4283.43      1.00
     β[15]     -8.45      2.30     -8.48    -12.12     -4.57   4260.51      1.00
     β[16]     -1.18      2.36     -1.12     -5.03      2.76   4967.28      1.00
     β[17]     -0.16      2.14     -0.19     -3.62      3.29   3726.29      1.00
     β[18]     -6.65      3.82     -6.71    -13.21     -0.86   3985.31      1.00
     β[19]     -2.16      2.20     -2.18     -5.76      1.43   3160.09      1.00
     β[20]     -3.71      2.20     -3.68     -7.30     -0.02   3877.82      1.00
     β[21]    -11.26      3.68    -11.34    -17.13     -5.13   4327.97      1.00
     β[22]     -5.95      2.40     -5.94     -9.97     -2.07   4703.96      1.00
     β[23]      5.92      2.37      5.92      2.08      9.67   4153.01      1.00
     β[24]     -4.07      2.26     -4.07     -7.76     -0.37   5563.09      1.00
     β[25]     -5.09      3.02     -5.11    -10.29     -0.54   3885.94      1.00
     β[26]     -6.07      2.34     -6.07    -10.20     -2.40   5647.53      1.00
     β[27]     -2.60      2.13     -2.55     -6.46      0.69   4187.02      1.00
     β[28]    -13.06      3.09    -13.13    -18.11     -7.93   4795.54      1.00
     β[29]     -7.67      2.38     -7.71    -11.57     -3.82   4117.66      1.00
     β[30]      8.04      2.25      8.04      4.13     11.55   4600.00      1.00
     β[31]      2.45      2.26      2.47     -1.03      6.30   4007.59      1.00
     β[32]     -6.58      2.23     -6.59    -10.21     -2.82   3832.25      1.00
     β[33]     -0.80      2.37     -0.81     -4.99      2.80   4361.64      1.00
     β[34]     -0.83      2.16     -0.83     -4.16      2.91   4025.73      1.00
     β[35]    -10.61      2.40    -10.61    -14.60     -6.78   3563.09      1.00
     β[36]     -5.19      2.17     -5.18     -8.89     -1.93   3299.45      1.00
     β[37]     -0.25      2.42     -0.19     -4.01      3.83   3704.36      1.00
     β[38]     -2.73      2.35     -2.70     -6.32      1.24   4342.30      1.00
     β[39]     -3.52      1.84     -3.51     -6.77     -0.84   3976.06      1.00
     β[40]     -1.34      2.30     -1.37     -5.12      2.42   3742.77      1.00
     β[41]    -11.93      2.96    -11.83    -16.93     -7.23   4682.23      1.00
     β[42]     -5.60      1.96     -5.62     -8.69     -2.24   4732.31      1.00
     β[43]     -4.12      2.35     -4.08     -7.66     -0.08   4616.45      1.00
     β[44]     -5.79      2.20     -5.80     -9.27     -2.09   4924.40      1.00
     β[45]    -19.96      3.06    -19.91    -24.63    -14.61   3679.62      1.00
     β[46]     -7.46      3.85     -7.50    -13.75     -1.59   4155.61      1.00
     β[47]     -5.62      2.99     -5.66    -10.20     -0.32   4985.98      1.00
     β[48]      0.74      2.23      0.77     -2.89      4.42   3883.18      1.00
     β[49]     -0.05      1.96     -0.03     -3.14      3.21   5014.63      1.00
     β[50]     -6.66      2.21     -6.70    -10.32     -3.04   3677.17      1.00
     β[51]     -4.38      2.23     -4.35     -8.12     -0.87   4395.64      1.00
     β[52]     -6.26      2.33     -6.21    -10.11     -2.64   4683.23      1.00
     β[53]    -11.29      3.00    -11.30    -16.50     -6.57   4172.36      1.00
     β[54]      0.53      2.20      0.49     -2.74      4.33   3909.65      1.00
     β[55]     -4.76      2.25     -4.77     -8.53     -1.01   4825.57      1.00
     β[56]     -6.44      3.85     -6.41    -12.56      0.06   4113.27      1.00
     β[57]     -2.80      2.27     -2.74     -6.67      0.70   3687.04      1.00
     β[58]     -7.49      2.26     -7.51    -10.84     -3.50   4276.53      1.00
     β[59]     -4.83      2.15     -4.82     -8.31     -1.30   5575.94      1.00
     β[60]     -1.40      2.24     -1.41     -4.98      2.35   4515.19      1.00
     β[61]     -3.11      2.28     -3.11     -6.95      0.55   4081.08      1.00
     β[62]     -8.59      2.31     -8.61    -12.35     -4.93   3377.48      1.00
     β[63]     -0.72      2.21     -0.73     -4.17      3.04   4547.00      1.00
     β[64]     -1.68      2.26     -1.63     -5.77      1.66   3385.73      1.00
     β[65]     -3.64      2.22     -3.64     -7.07      0.12   5571.42      1.00
     β[66]     -3.56      2.21     -3.63     -7.35     -0.16   4991.63      1.00
     β[67]     -1.55      1.90     -1.54     -4.73      1.41   4627.93      1.00
     β[68]     -7.49      2.09     -7.48    -11.06     -4.34   4079.94      1.00
     β[69]     -1.36      2.33     -1.39     -5.11      2.52   4569.82      1.00
     β[70]     -0.29      2.23     -0.25     -4.25      2.95   4242.55      1.00
     β[71]    -10.03      1.98    -10.06    -13.32     -6.94   4442.15      1.00
     β[72]     -7.53      2.24     -7.46    -11.36     -3.98   5482.66      1.00
     β[73]     -3.37      2.34     -3.39     -7.11      0.40   5019.89      1.00
     β[74]      2.41      2.23      2.45     -1.18      6.12   4508.04      1.00
     β[75]     -1.25      2.26     -1.33     -4.92      2.55   5646.07      1.00
     β[76]     -4.17      2.23     -4.17     -7.97     -0.66   5675.45      1.00
     β[77]     -4.32      2.26     -4.33     -8.17     -0.89   4923.77      1.00
     β[78]     -0.77      2.26     -0.75     -4.48      2.97   4750.75      1.00
     β[79]     -3.90      2.27     -3.89     -7.59     -0.15   4487.57      1.00
     β[80]    -12.61      2.64    -12.60    -17.07     -8.36   4246.54      1.00
     β[81]     -1.52      2.29     -1.45     -5.58      2.05   4651.25      1.00
     β[82]     -4.46      2.18     -4.50     -7.95     -0.89   4468.52      1.00
     β[83]     -1.38      2.24     -1.39     -4.98      2.40   5150.16      1.00
     β[84]     10.63      2.32     10.60      6.78     14.33   4505.47      1.00
     β[85]      0.58      2.27      0.59     -3.09      4.36   3539.17      1.00
     β[86]    -10.39      2.28    -10.40    -14.11     -6.87   4488.76      1.00
     β[87]     -3.61      2.29     -3.60     -7.36      0.10   5316.67      1.00
     β[88]     -0.30      2.22     -0.32     -3.80      3.35   3819.55      1.00
     β[89]     -1.93      2.32     -1.92     -5.96      1.63   4179.68      1.00
     β[90]     -2.15      2.37     -2.15     -5.91      1.70   3341.19      1.00
     β[91]     -3.70      2.26     -3.69     -7.61     -0.12   3904.55      1.00
     β[92]      2.18      2.20      2.17     -1.22      6.00   4908.78      1.00
     β[93]     -3.33      2.28     -3.41     -7.01      0.50   3571.76      1.00
     β[94]     -2.74      2.32     -2.75     -6.52      0.98   4606.86      1.00
     β[95]     -4.31      2.28     -4.40     -7.87     -0.61   4132.73      1.00
     β[96]     -0.62      2.29     -0.61     -4.24      3.25   3616.64      1.00
     β[97]     -7.27      2.28     -7.18    -11.10     -3.62   4140.66      1.00
     β[98]     -5.21      2.29     -5.23     -8.89     -1.43   5202.64      1.00
     β[99]     -3.25      1.91     -3.24     -6.55     -0.23   4539.26      1.00
    β[100]     -2.23      2.32     -2.18     -6.27      1.43   5201.07      1.00
    β[101]     -4.05      2.14     -4.01     -7.60     -0.70   3869.35      1.00
    β[102]     -3.24      2.97     -3.28     -7.89      1.76   3127.43      1.00
    β[103]      0.79      2.43      0.81     -2.78      5.17   6492.31      1.00
    β[104]     -4.63      2.39     -4.58     -8.28     -0.38   4463.70      1.00
    β[105]     -3.55      2.32     -3.48     -7.44      0.17   3883.25      1.00
    β[106]     -3.12      2.38     -3.04     -7.04      0.68   4317.92      1.00
    β[107]     -0.32      2.26     -0.27     -3.93      3.38   4878.10      1.00
    β[108]    -19.46      2.24    -19.44    -23.17    -15.80   3707.86      1.00
    β[109]     -5.55      2.08     -5.50     -8.91     -2.23   5596.44      1.00
    β[110]     -3.00      2.37     -2.99     -6.79      0.91   4362.65      1.00
    β[111]     -1.77      2.96     -1.76     -6.89      2.89   4287.54      1.00
    β[112]     -5.76      4.03     -5.67    -12.41      0.60   4568.87      1.00
    β[113]      0.30      2.26      0.29     -3.39      4.05   4201.45      1.00
    β[114]     -8.96      2.15     -8.93    -12.54     -5.55   4902.72      1.00
    β[115]     -6.70      3.75     -6.70    -12.80     -0.58   3675.94      1.00
    β[116]     -4.58      2.27     -4.57     -7.94     -0.57   4612.40      1.00
    β[117]     -0.47      2.21     -0.45     -4.30      2.92   4604.65      1.00
    β[118]     -0.88      2.17     -0.89     -4.45      2.67   3857.26      1.00
    β[119]      3.37      3.06      3.34     -1.63      8.40   4641.89      1.00
    β[120]     -7.61      2.76     -7.61    -12.08     -3.07   4584.45      1.00
    β[121]     -7.80      2.18     -7.84    -11.53     -4.31   4143.18      1.00
    β[122]     -4.47      2.28     -4.42     -8.62     -1.11   5003.13      1.00
    β[123]     -1.87      3.11     -1.85     -6.48      3.77   3837.84      1.00
    β[124]     -7.47      2.30     -7.42    -11.15     -3.56   3488.98      1.00
    β[125]     -2.96      2.30     -2.95     -6.51      0.86   4045.44      1.00
    β[126]     -4.62      2.32     -4.67     -8.37     -0.90   4501.75      1.00
    β[127]     -7.88      2.22     -7.87    -11.55     -4.17   4004.76      1.00
    β[128]     -4.28      2.31     -4.32     -7.78     -0.44   3548.34      1.00
    β[129]     -5.02      2.21     -5.08     -8.72     -1.32   3864.48      1.00
    β[130]     -9.61      2.21     -9.60    -13.21     -6.13   4645.20      1.00
    β[131]      5.26      2.23      5.33      1.71      9.07   3949.80      1.00
    β[132]     -5.85      2.23     -5.85     -9.48     -2.19   4969.66      1.00
    β[133]      0.48      3.14      0.54     -4.68      5.72   4701.96      1.00
    β[134]     -1.64      2.16     -1.65     -5.28      1.70   4533.16      1.00
    β[135]      1.34      2.32      1.36     -2.49      5.12   4581.77      1.00
    β[136]     -5.27      2.41     -5.27     -9.53     -1.62   4665.97      1.00
    β[137]     -5.41      2.89     -5.32    -10.61     -1.06   4848.17      1.00
    β[138]     -0.34      2.25     -0.35     -3.84      3.53   6135.90      1.00
    β[139]     -8.16      2.34     -8.12    -11.84     -4.21   4310.14      1.00
    β[140]     -3.98      2.25     -3.96     -7.82     -0.19   5176.38      1.00
    β[141]     -2.42      2.22     -2.46     -5.94      1.29   4938.04      1.00
    β[142]     -7.50      2.23     -7.48    -10.93     -3.83   5237.66      1.00
    β[143]     -1.22      1.95     -1.24     -4.48      1.89   4102.71      1.00
    β[144]     -6.33      3.78     -6.32    -12.02      0.48   3855.15      1.00
    β[145]      0.36      2.30      0.31     -3.04      4.50   5106.15      1.00
    β[146]      5.67      2.24      5.69      2.32      9.50   4880.28      1.00
    β[147]     -1.28      2.33     -1.32     -4.71      2.91   4092.35      1.00
    β[148]     -2.58      2.34     -2.60     -6.47      1.23   3386.65      1.00
    β[149]    -12.39      2.23    -12.40    -16.17     -8.91   5298.42      1.00
    β[150]     -6.77      3.07     -6.72    -11.85     -1.95   4858.98      1.00
    β[151]     -7.95      2.38     -7.92    -11.65     -3.74   4461.45      1.00
    β[152]     -2.49      2.33     -2.50     -6.22      1.21   3764.36      1.00
    β[153]     -1.30      2.37     -1.29     -5.04      2.67   6336.56      1.00
    β[154]     -3.86      2.20     -3.87     -7.43     -0.39   3279.35      1.00
    β[155]     -2.65      2.41     -2.66     -7.01      0.97   4424.41      1.00
    β[156]     -3.71      2.22     -3.71     -7.33     -0.15   4825.16      1.00
    β[157]     -8.63      3.94     -8.60    -14.57     -1.90   3967.71      1.00
    β[158]     -3.99      2.31     -3.98     -7.49     -0.07   4349.28      1.00
    β[159]     -4.94      1.98     -4.90     -8.41     -1.91   4700.91      1.00
    β[160]     -6.93      2.43     -6.99    -10.98     -2.99   4785.30      1.00
    β[161]     -3.47      2.13     -3.47     -7.00     -0.16   3972.92      1.00
    β[162]    -15.86      3.11    -15.80    -21.10    -10.84   4183.32      1.00
    β[163]     -0.94      1.87     -0.95     -4.08      2.08   4370.92      1.00
    β[164]     -2.01      2.22     -2.00     -5.44      1.98   5639.13      1.00
    β[165]     -2.00      2.22     -1.98     -5.55      1.68   3479.20      1.00
    β[166]     -4.23      2.27     -4.25     -7.71     -0.32   4338.94      1.00
    β[167]     -1.70      3.12     -1.79     -6.56      3.57   4373.63      1.00
    β[168]     -1.91      2.28     -1.90     -5.93      1.68   4487.30      1.00
    β[169]    -17.05      2.31    -17.07    -20.95    -13.41   4857.17      1.00
    β[170]     -1.69      2.27     -1.68     -5.50      1.95   4521.26      1.00
    β[171]     -2.54      2.31     -2.57     -6.12      1.41   4106.66      1.00
    β[172]     -1.94      1.91     -1.91     -4.93      1.32   3996.98      1.00
    β[173]     -4.88      2.92     -4.92     -9.76     -0.26   4223.72      1.00
    β[174]     -7.95      2.24     -7.97    -11.47     -4.07   4681.34      1.00
    β[175]     -1.66      2.25     -1.69     -4.95      2.32   3655.75      1.00
       μ_α   2774.79     54.96   2775.49   2686.53   2863.70   3474.50      1.00
       μ_β     -4.17      0.41     -4.17     -4.85     -3.51   2474.19      1.00
         σ    136.71      2.74    136.64    132.14    141.13   1906.90      1.00
       σ_α    723.12     31.96    722.08    672.81    774.22   3737.96      1.00
       σ_β      4.99      0.36      4.96      4.39      5.55   2482.51      1.00

Number of divergences: 0


import arviz as az

data = az.from_numpyro(mcmc)
az.plot_trace(data, compact=True, figsize=(15, 25));

Let’s see how this is done:

The model is described by the following equations:

where is week variable.

There are 176 patients, each has multiple weekly observations. There are 176 and 176 (random effects for intercept and for slope). There are two “fixed effects”, and . And there is one standard deviation for the residuals, two standard deviations and for the random effects.

The idea of Bayesian inference is that we have this model, and the data, and we want to infer the posterior distribution of the parameters given the data. The posterior distribution is proportional to the likelihood times the prior distribution.

We can see in this case, the random effects model has similar results as the Bayesian model.

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, Predictive

train = pd.read_csv(
    "https://gist.githubusercontent.com/ucals/"
    "2cf9d101992cb1b78c2cdd6e3bac6a4b/raw/"
    "43034c39052dcf97d4b894d2ec1bc3f90f3623d9/"
    "osic_pulmonary_fibrosis.csv"
)


from jax import random

def model(patient_code, Weeks, FVC_obs=None):
    μ_α = numpyro.sample("μ_α", dist.Normal(0.0, 500.0))
    σ_α = numpyro.sample("σ_α", dist.HalfNormal(100.0))
    μ_β = numpyro.sample("μ_β", dist.Normal(0.0, 3.0))
    σ_β = numpyro.sample("σ_β", dist.HalfNormal(3.0))

    n_patients = len(np.unique(patient_code))

    with numpyro.plate("plate_i", n_patients):
        α = numpyro.sample("α", dist.Normal(μ_α, σ_α))
        β = numpyro.sample("β", dist.Normal(μ_β, σ_β))

    σ = numpyro.sample("σ", dist.HalfNormal(100.0))
    FVC_est = α[patient_code] + β[patient_code] * Weeks

    with numpyro.plate("data", len(patient_code)):
        numpyro.sample("obs", dist.Normal(FVC_est, σ), obs=FVC_obs)
       
from sklearn.preprocessing import LabelEncoder

patient_encoder = LabelEncoder()
train["patient_code"] = patient_encoder.fit_transform(train["Patient"].values)

FVC_obs = train["FVC"].values
Weeks = train["Weeks"].values
patient_code = train["patient_code"].values
        
        
numpyro.render_model(
    model=model,
    model_args=(patient_code, Weeks, FVC_obs),
    render_distributions=True,
    render_params=True,
)

Uplift Modeling in Python

Xiang Ao — Fri, 10 Oct 2025 04:00:00 GMT

Uplift modeling is a technique used to estimate the incremental impact of a treatment or intervention on an outcome, often used in marketing. It helps identify which individuals are more likely to respond positively to a treatment, allowing for more targeted interventions.

To me this is basically conditional average treatment effect (CATE).

I am copying this code from https://github.com/uber/causalml/blob/master/docs/examples/uplift_trees_with_synthetic_data.ipynb. Just to see how it works and also test out python code in quarto.

import numpy as np
import pandas as pd

from causalml.dataset import make_uplift_classification
from causalml.inference.tree import UpliftRandomForestClassifier
from causalml.metrics import plot_gain

from sklearn.model_selection import train_test_split
import importlib
print(importlib.metadata.version('causalml') )

df, x_names = make_uplift_classification()
df.head()

# Look at the conversion rate and sample size in each group
df.pivot_table(values='conversion',
               index='treatment_group_key',
               aggfunc=[np.mean, np.size],
               margins=True)
               
# Split data to training and testing samples for model validation (next section)
df_train, df_test = train_test_split(df, test_size=0.2, random_state=111)     

from causalml.inference.tree import UpliftTreeClassifier

clf = UpliftTreeClassifier(control_name='control')
clf.fit(df_train[x_names].values,
         treatment=df_train['treatment_group_key'].values,
         y=df_train['conversion'].values)
p = clf.predict(df_test[x_names].values)

df_res = pd.DataFrame(p, columns=clf.classes_)
df_res.head()

uplift_model = UpliftRandomForestClassifier(control_name='control')

uplift_model.fit(df_train[x_names].values,
                 treatment=df_train['treatment_group_key'].values,
                 y=df_train['conversion'].values)

df_res = uplift_model.predict(df_test[x_names].values, full_output=True)
print(df_res.shape)
df_res.head()          

y_pred = uplift_model.predict(df_test[x_names].values)

# Put the predictions to a DataFrame for a neater presentation
# The output of `predict()` is a numpy array with the shape of [n_sample, n_treatment] excluding the
# predictions for the control group.
result = pd.DataFrame(y_pred,
                      columns=uplift_model.classes_[1:])
result.head()

# If all deltas are negative, assing to control; otherwise assign to the treatment
# with the highest delta
best_treatment = np.where((result < 0).all(axis=1),
                           'control',
                           result.idxmax(axis=1))

# Create indicator variables for whether a unit happened to have the
# recommended treatment or was in the control group
actual_is_best = np.where(df_test['treatment_group_key'] == best_treatment, 1, 0)
actual_is_control = np.where(df_test['treatment_group_key'] == 'control', 1, 0)

synthetic = (actual_is_best == 1) | (actual_is_control == 1)
synth = result[synthetic]

auuc_metrics = (synth.assign(is_treated = 1 - actual_is_control[synthetic],
                             conversion = df_test.loc[synthetic, 'conversion'].values,
                             uplift_tree = synth.max(axis=1))
                     .drop(columns=list(uplift_model.classes_[1:])))
 
plot_gain(auuc_metrics, outcome_col='conversion', treatment_col='is_treated')

0.17.0
(800, 9)

Causal Simulation

Xiang Ao — Wed, 01 Oct 2025 04:00:00 GMT

Evans, R. J., & Didelez, V. (2024). “Parameterizing and simulating from causal models” is about simulating from a causal model. I found that interesting.

The usual causal inference is to have observational data, then assume causal structure, then estimate the parameters. This paper is about the opposite: Suppose we have a causal model, what kind of distribution would generate it?

Suppose we are interested in joint distribution , is covariates, is treatment, is outcome.

Basically the joint distribution of can be factorized into the marginal distribution of , the conditional distribution of given and , and the copula . The copula is a function that captures the dependence between and given . is the “past”, which can be specified. is the “marginal structure model” which can be specified. The copula model can also be specified.

About copula: copula is a function can link joint distribution to marginal distributions.

where and are the marginal distributions of and , respectively, and is the copula function that captures the dependence between and . The copula function is a multivariate distribution with uniform marginals. It can be used to generate joint distributions from marginal distributions.

Here , and can all be vectors.

examples

What do we use this for? We can use this to simulate data from a causal model, and then fit different models to the simulated data. This allows to compare the performance of different models.

example 1

Suppose we have a causal graph like this:

We need to specify , where is the marginal distribution of , and is the conditional distribution of given and . This is the “marginal structural model” (Robins, 2000). Also the “g-formula”.

We also need to specify , the “past”.

Finally we need to specify the copula , which captures the dependence between and given . Depending on different situations, different copula can be used.

In the “causl” package example, we specify: and Gaussian copula with correlation .

library(causl)
# formulae corresponding to covariates, treatments, outcomes and the dependence
forms <- list(X ~ 1, A ~ X, Y ~ A, ~ 1)
# vector of model families (3=gamma/exponential, 1=normal/Gaussian)
fam <- c(1, 1, 1, 1)
# list of parameters, including 'beta' (regression params) and 'phi' dispersion
pars <- list(X = list(beta=0, phi=1),   
             A = list(beta=c(0,0.5), phi=1),
             Y = list(beta=c(-0.5,0.5), phi=1),
             cop = list(beta=1))

## now create a `causl_model` object
cm <- causl_model(formulas=forms, family=fam, pars=pars,method="inversion")

# now simulate 1000 observations
set.seed(123456)
data <- rfrugal(n=1000, causl_model=cm)
head(data)

            X           A          Y
1  0.83373317  1.10709111  1.1416713
2 -0.27604777 -0.27218853 -1.0495996
3 -0.35500184  0.77575209  1.3179678
4  0.08748742 -0.05959823 -0.6106023
5  2.25225573  0.61909216  2.0789188
6  0.83446013  0.52019822  0.3229069

In this example, first we specify the structure of the model using a list of formulas. The first formula is for the covariates, the second for the treatments, the third for the outcomes, and the fourth for the dependence structure (copula). Note that outcome only depends on in this interventional distribution.

Then we specify the families of the random variables. Here we use exponential distribution for , normal distribution for and , and Gaussian copula for the dependence structure.

Finally we specify the parameters of the model, including regression coefficients and dispersion parameters. For example, has two coefficients, intercept -.5 and slope .5.

The causl_model object is then created using these components.

In this simple example, if we run a regression with , in the model, on the simulated data, we would get the correct effect back:

lm1 <- lm(Y ~ A*X, data=data)
summary(lm1)


Call:
lm(formula = Y ~ A * X, data = data)

Residuals:
    Min      1Q  Median      3Q     Max 
-2.6326 -0.6212 -0.0105  0.6051  3.3132 

Coefficients:
             Estimate Std. Error t value Pr(>|t|)    
(Intercept) -0.491922   0.030413 -16.175   <2e-16 ***
A            0.493471   0.027820  17.738   <2e-16 ***
X            0.468038   0.031878  14.682   <2e-16 ***
A:X         -0.009247   0.022360  -0.414    0.679    
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

Residual standard error: 0.8953 on 996 degrees of freedom
Multiple R-squared:  0.4875,    Adjusted R-squared:  0.4859 
F-statistic: 315.8 on 3 and 996 DF,  p-value: < 2.2e-16

We can also use MLE:

out <- fit_causl(data, formulas = list(X ~ 1, Y ~ A, ~ 1), family = c(1, 1, 1))
out

log-likelihood:  -2716.696 
X ~ 1
              est.   s.e. sandwich
(intercept) 0.0108 0.0314   0.0314
  residual s.e.:  0.983 0.044 0.0447 

Y ~ A
              est.   s.e. sandwich
(intercept) -0.491 0.0319   0.0319
A            0.494 0.0278   0.0295
  residual s.e.:  1.01 0.0472 0.0475 

copula parameters:
cop ~ 1
             est.  s.e. sandwich
(intercept) 0.998 0.069   0.0734

Plasmode simulation

In reality, we have a data set and we have a causal model. We want to simulate data from the causal model, but we want the simulated data to have the same distribution as the original data. This is called “plasmode simulation”. For example, we don’t want to use the real outcome variable. In stead, we simulate X’s from the original data, Then simulate Y and A from the causal model. We can then test which method works better in estimating the causal effect of on .

# A tibble: 4,802 × 58
     x_1 x_2     x_3   x_4   x_5   x_6   x_7   x_8   x_9  x_10  x_11  x_12  x_13
               
 1    29 C         1     7    60    85     0     0     1     0     0     1     0
 2    27 C         0     0    64   178     0     0     0     0     0     2     1
 3    27 C         0     0    60   102     0     0     0     0     0     1     0
 4    37 C         0     0    65   174     0     0     0     0     0     1     0
 5    24 C        20    14    63   129     0     0     0     0     0     1     0
 6    27 C        40    15    63   135     0     0     0     0     0     1     1
 7    26 C        20     8    69   140     0     0     0     1     0     0     0
 8    33 C         0     0    60   110     0     0     0     1     0     0     1
 9    28 C         3     5    61   160     0     0     0     0     0     3     1
10    31 C         0     0    63   114     0     1     0     0     0     0     0
# ℹ 4,792 more rows
# ℹ 45 more variables: x_14 , x_15 , x_16 , x_17 ,
#   x_18 , x_19 , x_20 , x_21 , x_22 , x_23 ,
#   x_24 , x_25 , x_26 , x_27 , x_28 , x_29 ,
#   x_30 , x_31 , x_32 , x_33 , x_34 , x_35 ,
#   x_36 , x_37 , x_38 , x_39 , x_40 , x_41 ,
#   x_42 , x_43 , x_44 , x_45 , x_46 , x_47 , …

# Model for the causal effect of smoking on birthweight
forms <- list(list(),
              list(A ~ x_1 + x_3 + x_4),
              list(Y ~ A),
              list(~ 1))
# fams <- list(integer(0), 5, 1, 1)
fams <- list(integer(0), "binomial", "gaussian", 1)
pars <- list(A = list(beta=c(-1.5,0.03,0.02,0.05)),
             Y = list(beta=c(3200, -500), phi=400^2),
             cop = list(beta=-1))

cm2 <- causl_model(formulas=forms, family=fams, pars=pars, dat = dat)
 
set.seed(123456)
data2 <- rfrugal(causl_model=cm2)
head(data2)

  x_1 x_2 x_3 x_4 x_5 x_6 x_7 x_8 x_9 x_10 x_11 x_12 x_13 x_14 x_15 x_16 x_17
1  29   C   1   7  60  85   0   0   1    0    0    1    0    0    2    0    0
2  27   C   0   0  64 178   0   0   0    0    0    2    1    0    2    0    0
3  27   C   0   0  60 102   0   0   0    0    0    1    0    0    0    0    0
4  37   C   0   0  65 174   0   0   0    0    0    1    0    0    1    0    0
5  24   C  20  14  63 129   0   0   0    0    0    1    0    0    0    0    0
6  27   C  40  15  63 135   0   0   0    0    0    1    1    2    2    0    0
  x_18 x_19 x_20 x_21 x_22 x_23 x_24 x_25 x_26 x_27 x_28 x_29 x_30 x_31 x_32
1   12   35   10    J    1   43    B   17   64  105   30   11    0    0    0
2   12   75   12    J    1   50    E   12   90  183   40    8    0    0    0
3   12   35   10    J    1   57    E   14   70  121   27   12    0    0    0
4   11   35   12    J    1   43    E   10   80  185   39   10    0    0    0
5    9   35   15    J    1   33    E   14   70  160   28    7    0    0    0
6    9   45   10    J    1   40    E   15   64  160   31    9    0    0    0
  x_33 x_34 x_35 x_36 x_37 x_38 x_39 x_40 x_41 x_42 x_43 x_44 x_45 x_46 x_47
1    0   79   28  340   48    1    9    9    3    0    3   54   16    0    0
2    0   69   30  430   73    0    9    9    3    0    5   51   16    0    0
3    0   84   29  450   76    0    3    6    2    0    4   41   13    0    0
4    0   82   30  440   78    1    8    9    2    0    8   70   16    0    0
5    1   85   35  360   71    1    5    8    0    2    3   65   18    0    0
6    1   81   28  515   56    1    7    9    2    0    1   63   20    0    0
  x_48 x_49 x_50 x_51 x_52 x_53 x_54 x_55 x_56 x_57 x_58 A        Y
1    0    0    0    0    0    0    0    0    0   45   39 1 2125.677
2    0    0    0    0    0    0    0    0    0   46   42 1 2688.990
3    0    1    0    0    0    0    0    0    0   45   40 0 3442.924
4    0    0    0    0    0    0    0    0    0   47   40 0 3212.584
5    0    2    0    0    0    0    0    0    0   47   43 0 2992.775
6    0    0    0    0    0    0    0    0    0   45   44 0 2738.893

# we can also use rfrugalParam(), which is older version of rfrugal().
#datAY <- rfrugalParam(formulas=forms, family=fams, pars=pars, dat=dat)

In this specification, we have a binary treatment and a continuous outcome . The covariates are , and . The copula is a Gaussian copula with correlation -1. We can then simulate the data.

a more complicated model

Let’s simulate a more complicated model:

library(ggplot2)
library(ggdag)

g <- dagify(
  A1 ~ L1,
  L2 ~ L1,
  L2 ~ A1,
  A2 ~ A1,
  A2 ~ L2,
  Y ~ A2,
  Y ~ L2,
  Y ~ A1,
  exposure = "A2",
  outcome = "Y",
  coords = list(x = c(L1 = 1,  L2 = 2, Y = 3, A1 = 1.5, A2 = 2.5),
                y = c(L1 = 1,  L2 = 1, Y = 1, A1 = 1.5, A2 = 1.5))
)

ggdag(g) + 
  theme_dag()

Let’s see how to write a causal model:

# formulae corresponding to covariates, treatments, outcomes and the dependence
forms <- list(list(L1  ~ 1, L2 ~ L1*A1),   # covariates
              list(A1 ~ L1, A2 ~ L2*A1), # treatments
              Y ~ A1*A2*L2,             # outcome
              ~ 1)

# vector of model families (3=gamma/exponential, 1=normal/Gaussian)

fams <- list(c(1, 1), c(5,5), 1, 1)
pars <- list(L1 = list(beta=0, phi=1),
             L2 = list(beta=c(0.3,0.5,-0.2,-0.1), phi=1),
             A1 = list(beta=c(-0.3,0.4), phi=1),
             A2 = list(beta=c(0.5,0.3,0.1,0), phi=1),
             Y = list(beta=c(0,1,2,0.5,1,0.2,0,0.8), phi=1), 
             cop = list(beta = 0.5))

# we can use rfrugalParam() to simulate data from the model, or rfrugal().
#dat <- rfrugalParam(n=1e4, formulas=forms, family = fams, pars=pars)
#head(dat)

cm3 <- causl_model(formulas=forms, family=fams, pars=pars)
 
set.seed(123456)
data3 <- rfrugal(n=1e4,causl_model=cm3)
head(data3)

           L1          L2 A1 A2         Y
1  0.83373317  1.91691936  1  1 8.1680612
2 -0.27604777  0.17008294  1  1 5.2276976
3 -0.35500184  0.16907708  1  1 2.8194835
4  0.08748742  0.19688985  0  1 2.6789661
5  2.25225573 -0.29257431  0  0 0.3180188
6  0.83446013  0.02661196  0  1 1.5694426

If we have the correct specification for the outcome model, we’ll get it right by linear model:

m1 <- glm(Y ~ A1*A2*L2, data = data3)
summary(m1)


Call:
glm(formula = Y ~ A1 * A2 * L2, data = data3)

Coefficients:
            Estimate Std. Error t value Pr(>|t|)    
(Intercept) -0.11524    0.02110  -5.461 4.85e-08 ***
A1           1.11589    0.03245  34.391  < 2e-16 ***
A2           2.02332    0.02682  75.454  < 2e-16 ***
L2           0.80424    0.01929  41.682  < 2e-16 ***
A1:A2        0.98345    0.04090  24.044  < 2e-16 ***
A1:L2        0.16028    0.02976   5.386 7.37e-08 ***
A2:L2       -0.03446    0.02401  -1.435    0.151    
A1:A2:L2     0.83541    0.03693  22.620  < 2e-16 ***
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

(Dispersion parameter for gaussian family taken to be 0.9100255)

    Null deviance: 53902  on 9999  degrees of freedom
Residual deviance:  9093  on 9992  degrees of freedom
AIC: 27446

Number of Fisher Scoring iterations: 2

library(marginaleffects)
avg_comparisons(
    model = m1,
    variables = "A2")


 Estimate Std. Error   z Pr(>|z|)   S 2.5 % 97.5 %
     2.51     0.0202 124   <0.001 Inf  2.47   2.55

Term: A2
Type: response
Comparison: 1 - 0

So the ATE is 2.5.

Let’s use “lmtp” package to estimate the treatment effect of on .

library(lmtp)

d1 <- function(data, trt) {
  rep(1, nrow(data))
}
A <- "A2"
Y <- "Y"
W <- c("L1", "L2", "A1")

set.seed(34465)

treat <- lmtp_tmle(
  data = data3, 
  trt = "A2", 
  outcome = "Y", 
  baseline = W, 
  outcome_type = "continuous", 
  shift = d1, 
  folds = 1, 
  learners_trt = "SL.glm", 
  learners_outcome = "SL.glm"
)

print(treat)

LMTP Estimator: TMLE

Trt. Policy: (d1)

── Population intervention estimate ──

  Estimate: 3.046
Std. error: 0.022

95% Conf. int.: 3.003, 3.089

d2 <- function(data, trt) {
  rep(0, nrow(data))
}
control <- lmtp_tmle(
  data = data3, 
  trt = "A2", 
  outcome = "Y", 
  baseline = W, 
  outcome_type = "continuous", 
  shift = d2, 
  folds = 1, 
  learners_trt = "SL.glm", 
  learners_outcome = "SL.glm"
)

print(control)

LMTP Estimator: TMLE

Trt. Policy: (d2)

── Population intervention estimate ──

  Estimate: 0.542
Std. error: 0.021

95% Conf. int.: 0.501, 0.582

lmtp_contrast(treat, ref = control)

LMTP Contrast: additive

Null hypothesis: theta == 0

shift ref estimate std.error conf.low conf.high p.value 1 3.05 0.542 2.5 0.0227 2.46 2.55 <0.001

Longitudinal Modified Treatment Policies (LMTP)

Xiang Ao — Sat, 23 Aug 2025 04:00:00 GMT

I read a few papers with longitudinal modified treatment policy (LMTP), and found it interesting. It has been used in epidemiology and biostatistics, but I have not seen it in applied econometrics yet.

Here I am mostly following Nicholas Williams: https://beyondtheate.com/

We are usually interested in ATE, the average treatment effect. However, there could be more complicated situations that the treatment is continuous, or the treatment is multivalued, or the treatment is time-varying. The static interventions have problems. For example, the hypothetical interventions that treatment applies to everyone might be inconceivable. Or such intervention could make positivity assumption fail.

Suppose we have such a DAG:

library(ggplot2)
library(ggdag)

g <- dagify(
  A1 ~ L1,
  L2 ~ L1,
  L2 ~ A1,
  A2 ~ A1,
  A2 ~ L2,
  Y ~ A2,
  Y ~ L2,
  Y ~ A1,
  exposure = "A2",
  outcome = "Y",
  coords = list(x = c(L1 = 1,  L2 = 2, Y = 3, A1 = 1.5, A2 = 2.5),
                y = c(L1 = 1,  L2 = 1, Y = 1, A1 = 1.5, A2 = 1.5))
)

ggdag(g) + 
  theme_dag()

We have multiple time points, and the treatment is time-varying. We are interested in not only the ATE of on , but also some other hypothetical interventions.

Definitions and assumptions

Assumptions

Positivity. Basically if there is a unit with and , then there is a unit with and .
Sequential unconfoundedness. There is no unmeasured confounders.

dynamic treatment regime

Some notations commonly used in this literature: We observe , where is the observed confounder, is the treatment, and is the outcome. The treatment can be time-varying. We can also have multiple treatments at different time points. For example, we can have , , etc. In this graph, baseline affects , which in turn affects , which then affects , and finally affects the outcome .

History is the history of data up to time , right before . for example, , and . is the hypothetical intervention function, or shift function, which is a function of the history and the treatment . For example, is a user-given function to map , , and to a potential treatment value. The function can be deterministic, or it can be stochastic. Then we can replace with . Then after that is called the natural value of treatment.

This is very general, comparing to the static treatment regime.

Suppose treatment is a function of the history of treatment and confounders, for example, . can be set to a fixed value, say 1 or 0, or some value . This function can be anything, it can be taking a deterministic value, or it can be a function that takes the natural treatment value as input. In the package “lmtp”, this is called a shift function, or hypothetical intervention. For example, can be set to 1 if , or can be set to double the natural value of . Many possibilities.

In comparison, for ATE, we only need to set to 0 or 1.

Under this LMTP, the causal parameter is

is the potential outcome under the hypothetical intervention . At time 1, .

modified treatment policy

Let’s look at a simulated data set to see how exactly we can estimate it.

This simulation is from Susmann et al. (2024) “Longitudinal Generalizations of the Average Treatment Effect on the Treated for Multi-valued and Continuous Treatments”. I modified slightly to fit the DAG above.

library(tidyverse)
library(tidyr)

mtp <- function(data, trt) {
  a <- data[[trt]]
  a * 0 + 1
}
simulate_data <- function(seed, N, tau, sigma = 0.5) {
  set.seed(seed)
  
  data <- tibble(id = 1:N)
  
  for(t in 1:tau) {
    Lt <- paste0("L_", t)
    Ltd <- paste0("L_", t, "d")
    
    At <- paste0("A_", t)
    Atd <- paste0("A_", t, "d")
    
    Lt1 <- paste0("L_", t - 1)
    Lt1d <- paste0("L_", t - 1, "d")
    
    At1 <- paste0("A_", t - 1)
    At1d <- paste0("A_", t - 1, "d")
    
    if(t == 1) {
      data[[Lt]] <- runif(N, 0, 1)
      data[[Ltd]] <- data[[Lt]]
      
      data[[At]] <- rbinom(N, size = 1, prob = 0.5)
    } 
    else {
      data[[Lt]] <- rnorm(N, mean = 0.25 * data[[Lt1]], 0.5)
      data[[Ltd]] <- rnorm(N, mean = 0.25 * data[[Lt1d]], 0.5)
      
      data[[At]] <- rbinom(N, size = 1, prob = plogis(0.5 - 0.2 * data[[At1]] + 0.1 * data[[Lt1]]))
    }
    
    data[[Atd]] <- mtp(data, At)
  }
  data$Y  <- rnorm(N, data[[At]] + data[[Lt]], sigma)
  data$Yd <- rnorm(N, data[[Atd]] + data[[Ltd]], sigma)
  data
}

simulated_data1 <- simulate_data(seed = 123, N = 10000, tau = 2)

mean(simulated_data1$Yd)

[1] 1.128593

mtp <- function(data, trt) {
  a <- data[[trt]]
  a * 0 + 0
}

simulated_data2 <- simulate_data(seed = 123, N = 1000, tau = 2)

mean(simulated_data2$Yd)

[1] 0.114403

mean(simulated_data1$Yd) - mean(simulated_data2$Yd)

[1] 1.01419

Note in this simulation the variables ending with “d” are the variables under hypothetical intervention, or modified treatment policy. is from , and is from . The treatment is from a Bernoulli distribution with probability 0.5, and the treatment is from a Bernoulli distribution with probability . The outcome is from a normal distribution with mean . In the simulated data, the modified treatment policy is to set the treatment to 0, then 1. The difference would be SATE.

In this case, we can just do a linear regression to get the effect of A2 on Y, knowing the exact DAG.

m1 <- glm(Y ~ L_2 + A_1   + A_2, data = simulated_data1)
summary(m1)


Call:
glm(formula = Y ~ L_2 + A_1 + A_2, data = simulated_data1)

Coefficients:
              Estimate Std. Error t value Pr(>|t|)    
(Intercept) -3.453e-05  9.615e-03  -0.004    0.997    
L_2          9.998e-01  9.889e-03 101.103   <2e-16 ***
A_1         -6.042e-03  1.001e-02  -0.604    0.546    
A_2          1.000e+00  1.027e-02  97.380   <2e-16 ***
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

(Dispersion parameter for gaussian family taken to be 0.2500364)

    Null deviance: 7400.0  on 9999  degrees of freedom
Residual deviance: 2499.4  on 9996  degrees of freedom
AIC: 14523

Number of Fisher Scoring iterations: 2

Let’s try a different MTP: set half of the time to 0, the other half remain unchanged.

This is, say, to set half of smokers to non-smokers, and the other half remain smokers.

# mtp <- function(data, trt) {
#   a <- data[[trt]]
#   epsilon <- rbinom(nrow(data), size = 1, prob = 0.5)
#   ifelse(epsilon <.5 & a == 1, 0, a)
# }
d <- function(a) {
  epsilon <- runif(length(a))
  ifelse(epsilon < 0.5 & a == 1, 0, a)
}
simulated_data1$m3_d <- simulated_data1$Y 

m2 <- glm(m3_d ~ L_1 + A_1 + L_2 + A_2, data = simulated_data1)
summary(m2)


Call:
glm(formula = m3_d ~ L_1 + A_1 + L_2 + A_2, data = simulated_data1)

Coefficients:
             Estimate Std. Error t value Pr(>|t|)    
(Intercept) -0.005815   0.012842  -0.453    0.651    
L_1          0.011964   0.017621   0.679    0.497    
A_1         -0.005956   0.010012  -0.595    0.552    
L_2          0.998831   0.009987 100.015   <2e-16 ***
A_2          0.999873   0.010273  97.334   <2e-16 ***
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

(Dispersion parameter for gaussian family taken to be 0.2500499)

    Null deviance: 7400.0  on 9999  degrees of freedom
Residual deviance: 2499.2  on 9995  degrees of freedom
AIC: 14525

Number of Fisher Scoring iterations: 2

simulated_data1$m2_d <- predict(m2, mutate(simulated_data1, A_2 = d(A_2)))

m1 <- glm(m2_d ~ L_1 + A_1, data = simulated_data1)
summary(m1)


Call:
glm(formula = m2_d ~ L_1 + A_1, data = simulated_data1)

Coefficients:
             Estimate Std. Error t value Pr(>|t|)    
(Intercept)  0.277598   0.015072  18.418   <2e-16 ***
L_1          0.289658   0.023474  12.339   <2e-16 ***
A_1         -0.008247   0.013462  -0.613     0.54    
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

(Dispersion parameter for gaussian family taken to be 0.4528119)

    Null deviance: 4596.0  on 9999  degrees of freedom
Residual deviance: 4526.8  on 9997  degrees of freedom
AIC: 20461

Number of Fisher Scoring iterations: 2

m1_d <- predict(m1, mutate(simulated_data1, A_1 = d(A_1)))
mean(m1_d)

[1] 0.4197166

Note this recursive process is based on Diaz, et al. (2023) “Nonparametric Causal Effects Based on Longitudinal Modified Treatment Policies”.

Here is how exactly we estimate it: we start with the last time point. Regress it on previous treatment and confounders, and then get the predicted value with changed based on the MTP. Regress that predicted value on the previous treatment and confounders, get predicted values with changed based on MTP. Repeat until the first time point. The average of the predicted value at time 1 is the expected value under this MTP.

This is basically g-formula extended to longitudinal data.

In a general case, this is the generalized g-formula to estimate :

Set , let . For , recursively define:

Then .

We start from the last period. Regress , which is on and . Then get the predicted value with changed based on the MTP. Then regress that predicted value on and . Repeat until the first time point. The average of the predicted value at time 1 is the expected value under this MTP.

Estimators

The authors advocate two estimators, TMLE and SDR (sequentially doubly robust estimator). The procedures are the same, starting from the last time point, then apply TMLE or SDR, iterate to the first time point.

Using lmtp package

library(tidyverse)
library(lmtp)
d1 <- function(data, trt) {
  rep(1, nrow(data))
}
A <- "A_2"
Y <- "Y"
W <- c("L_1", "L_2", "A_1")

set.seed(34465)

treat <- lmtp_tmle(
  data = simulated_data1, 
  trt = "A_2", 
  outcome = "Y", 
  baseline = W, 
  outcome_type = "continuous", 
  shift = d1, 
  folds = 1, 
  learners_trt = "SL.glm", 
  learners_outcome = "SL.glm"
)

print(treat)

LMTP Estimator: TMLE

Trt. Policy: (d1)

── Population intervention estimate ──

  Estimate: 1.117
Std. error: 0.008

95% Conf. int.: 1.101, 1.133

d2 <- function(data, trt) {
  rep(0, nrow(data))
}
A <- "A_2"
Y <- "Y"
W <- c("L_1", "L_2", "A_1")

set.seed(34465)

control <- lmtp_tmle(
  data = simulated_data1, 
  trt = "A_2", 
  outcome = "Y", 
  baseline = W, 
  outcome_type = "continuous", 
  shift = d2, 
  folds = 1, 
  learners_trt = "SL.glm", 
  learners_outcome = "SL.glm"
)

print(control)

LMTP Estimator: TMLE

Trt. Policy: (d2)

── Population intervention estimate ──

  Estimate: 0.117
Std. error: 0.009

95% Conf. int.: 0.099, 0.136

lmtp_contrast(treat, ref = control)

LMTP Contrast: additive

Null hypothesis: theta == 0

shift ref estimate std.error conf.low conf.high p.value 1 1.12 0.117 1 0.0102 0.98 1.02 <0.001

another example

The data set bmi, from the DynTxRegime package, are simulated to reflect a two-stage RCT {A1, A2} that studied the effect of meal replacement (MR) shakes versus a calorie deficit (CD) diet on adolescent

shift function 1

Consider a shift function that assigns meal replacement to all observations at time 1, but only meal replacement at time 2 to those observations whose 4-month BMI is greater than 30.

library(DynTxRegime)
data(bmiData)
bmi <- bmiData
glimpse(bmi)

Rows: 210
Columns: 8
$ gender       0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1…
$ race         1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0…
$ parentBMI    31.59683, 30.17564, 30.27918, 27.49256, 26.42350, 29.30970…
$ baselineBMI  35.84005, 37.30396, 36.83889, 36.70679, 34.84207, 36.68640…
$ month4BMI    34.22717, 36.38014, 34.42168, 32.52011, 33.72922, 32.06622…
$ month12BMI   34.27263, 36.38401, 34.41447, 32.52397, 33.73546, 32.15977…
$ A1           "CD", "CD", "MR", "CD", "CD", "MR", "CD", "CD", "CD", "CD"…
$ A2           "MR", "MR", "CD", "CD", "CD", "MR", "MR", "CD", "MR", "MR"…

d_dtr <- function(data, trt) {
  if (trt == "A1") return(rep("MR", nrow(data)))
  
  ifelse(data$month4BMI > 30, "MR", "CD")
}

fit_dtr <- lmtp_sdr(
  data = bmi, 
  trt = c("A1", "A2"), 
  outcome = "month12BMI", 
  baseline = c("gender", "race", "parentBMI"), 
  time_vary = list("baselineBMI", "month4BMI"),
  shift = d_dtr, 
  outcome_type = "continuous",
  folds = 1,
  learners_trt = "SL.glm", 
  learners_outcome = c("SL.mean", "SL.glm", "SL.gam")
)

fit_dtr

      Estimate: 35.853
    Std. error: 0.362

shift function 2

Suppose we are interested in comparing the dynamic treatment regime to a static treatment regime where all patients receive meal replacement at both time points. Using the SDR estimator, estimate the effect of this static intervention.

fit_MR <- lmtp_sdr(
  data = bmi, 
  trt = c("A1", "A2"), 
  outcome = "month12BMI", 
  baseline = c("gender", "race", "parentBMI"), 
  time_vary = list("baselineBMI", "month4BMI"),
  shift = \(data, trt) rep("MR", nrow(data)), 
  outcome_type = "continuous",
  folds = 1,
  learners_trt = "SL.glm", 
  learners_outcome = c("SL.mean", "SL.glm", "SL.gam")
)

fit_MR

LMTP Estimator: SDR

Trt. Policy: (function(data, trt) rep("MR", nrow(data)))

── Population intervention estimate ──

  Estimate: 35.831
Std. error: 0.361

95% Conf. int.: 35.123, 36.539

Let’s also estimate the effect of an intervention where all patients receive a calorie deficit diet at both time points.

fit_CD <- lmtp_sdr(
  data = bmi, 
  trt = c("A1", "A2"), 
  outcome = "month12BMI", 
  baseline = c("gender", "race", "parentBMI"), 
  time_vary = list("baselineBMI", "month4BMI"),
  shift = \(data, trt) rep("CD", nrow(data)), 
  outcome_type = "continuous",
  folds = 1,
  learners_trt = "SL.glm", 
  learners_outcome = c("SL.mean", "SL.glm", "SL.gam")
)

fit_CD

LMTP Estimator: SDR

Trt. Policy: (function(data, trt) rep("CD", nrow(data)))

── Population intervention estimate ──

  Estimate: 35.053
Std. error: 0.301

95% Conf. int.: 34.463, 35.643

Finally, we can compare the three treatment regimes using the lmtp_contrast function.

lmtp_contrast(fit_dtr, fit_MR, ref = fit_CD)


  shift  ref estimate std.error conf.low conf.high p.value
1  35.9 35.1    0.801     0.335    0.144      1.46  0.0169
2  35.8 35.1    0.778     0.335    0.122      1.43  0.0201