Interventional what-if analysis using the PyMC do-operator#

import warnings

import arviz as az
import numpy as np
import pandas as pd
import pymc as pm

warnings.filterwarnings("ignore")

%config InlineBackend.figure_format = "retina"
az.style.use("arviz-variat")
rng = np.random.default_rng(42)
SEED = 8927

Introduction#

In the realm of data science and analytics, understanding the causal relationships between variables is paramount. While traditional statistical methods have provided insights into these relationships, the advent of probabilistic programming has ushered in a new era of causal analysis. In this article, we will explore the power of interventional what-if analysis using the PyMC framework, with a special focus on the do-operator.

The do-operator lets us answer interventional “what-if” questions: What would happen to the outcome if we forced a variable to a specific value? By leveraging pm.do, we can programmatically simulate these scenarios, giving us a deeper understanding of the causal relationships between predictors and target variables.

Through a step-by-step guide, we will delve into the process of building a PyMC model skeleton, generating data using the do-operator, and validating the relationships captured by the model. Furthermore, we will explore how the do-operator can simulate different what-if scenarios, akin to programmatic A/B testing.

Interventions vs. counterfactuals: where does this notebook sit on the causal ladder?

Pearl’s causal ladder distinguishes three levels of causal reasoning:

L1 — Association (seeing): What do we observe? \(P(Y \mid X)\)
L2 — Intervention (doing): What happens if we force a variable to a value? \(P(Y \mid \operatorname{do}(X = x))\)
L3 — Counterfactual (imagining): What would have happened to this specific unit under a different action, given what we already observed? \(Y_x \mid X = x', Y = y'\)

pm.do directly supports L2 queries. This notebook demonstrates interventional what-if analysis — we set a variable to a new value and predict the outcome. True unit-level counterfactuals (L3) require an additional abduction step: inferring unit-specific exogenous terms (\(U\)) from observed evidence before predicting in the intervened world. That step is not demonstrated here.

For readers interested in L3 counterfactual reasoning, see Pearl, Glymour & Jewell (2016), Causal Inference in Statistics: A Primer, Sections 4.2.3–4.2.4.

Steps

The basic workflow for interventional what-if analysis using the do-operator in PyMC can be summarized in the following steps:

Step 1. Build a PyMC model skeleton
Step 2. Use model skeleton and generate data using do-operator to infuse relationship between predictors and target variable
Step 3. Use observe-operator to assign generated data on the model skeleton
Step 4. Create samples and validate that the infused relationship between predictors and target variable are captured by the model samples
Step 5. Use do-operator to generate target variable with different what-if scenarios (programmatic A/B testing)

Step 1. Build a pymc model skeleton#

For this demo, we are building a very simple Linear Regression model.

Predictor — ‘a’, ‘b’, ‘c’
Target Variable — ‘y’
Coefficients —

‘beta_ay’ -> coefficient of |a|

‘beta_by’ -> coefficient of |b|

‘beta_cy’ -> coefficient of |c|

with pm.Model(coords={"i": [0]}) as model_generative:
    # priors
    beta_y0 = pm.Normal("beta_y0")
    beta_ay = pm.Normal("beta_ay")
    beta_by = pm.Normal("beta_by")
    beta_cy = pm.Normal("beta_cy")
    # observation noise on Y
    sigma_y = pm.HalfNormal("sigma_y")
    # core nodes and causal relationships
    a = pm.Normal("a", mu=0, sigma=1, dims="i")
    b = pm.Normal("b", mu=0, sigma=1, dims="i")
    c = pm.Normal("c", mu=0, sigma=1, dims="i")
    y_mu = pm.Deterministic(
        "y_mu", beta_y0 + (beta_ay * a) + (beta_by * b) + (beta_cy * c), dims="i"
    )
    y = pm.Normal("y", mu=y_mu, sigma=sigma_y, dims="i")


pm.model_to_graphviz(model_generative)

../_images/1e77af97418e914d8790b118abd7bb3d132f148ee0a5dc2de827c44f1132a7ac.svg

Step 2. Use model skeleton and generate data using do-operator to infuse relationship between predictors and target variable. We will use this generated data for modelling later.#

Let’s first define the predictors relationship with target variable.

true_values = {
    "beta_ay": 1.5,
    "beta_by": 0.7,
    "beta_cy": 0.3,
    "sigma_y": 0.2,
    "beta_y0": 0.0,
}

Basically what we are saying here is, we are intentionally defining the coefficient values, which we expect predictive model to predict later on.

Now the magic begins. We will use do-operator to use this dictionary and sample data variables. How do we do this ? Simple by passing two arguments to pymc do-operator. First, the model skeleton object. And second, the coefficient dictionary.

model_simulate = pm.do(model_generative, true_values)

This will create a new model object with the coefficent variables values infused.

model_simulate.to_graphviz()

../_images/57291467ee2943e19c5b1e0ab077027d22eca557e7014d635037059f72aa4b58.svg

The gray shades on the coefficient variables depicts the tale. Check the previous model graph, it was all white.

Now, all we have to do is generate samples, the known pymc way.

Lets generate 100 samples.

N = 100

with model_simulate:
    simulate = pm.sample_prior_predictive(draws=N)

Sampling: [a, b, c, y]

We know that this generates an ArviZ object, and since we have called sample_prior_predictive, hence the object will only contain priors.

Extract the sampled prior data into a pandas dataframe.

observed = {
    "a": simulate.prior["a"].values.flatten(),
    "b": simulate.prior["b"].values.flatten(),
    "c": simulate.prior["c"].values.flatten(),
    "y": simulate.prior["y"].values.flatten(),
}

df = pd.DataFrame(observed)
print(df.shape)
df.head(5)

(100, 4)

	a	b	c	y
0	0.510869	-1.311934	-0.362795	-0.306949
1	-0.734647	-0.539776	0.019862	-1.162661
2	-0.358604	0.343649	-0.093558	-0.278148
3	-0.251229	-1.154599	-0.542470	-1.489163
4	-1.316097	0.114966	-0.264377	-2.143315

Ok, so now we are all set with a sample data.

Step 3. Use observe-operator to assign generated data on the model skeleton#

Now, this is another cool feature of pymc newly introduced observe method. Observe method, takes in a model skeleton and the dictionary with the data for the variables we want to infuse into the model.

data_dict = {"a": df["a"], "b": df["b"], "c": df["c"], "y": df["y"]}
model_inference = pm.observe(model_generative, data_dict)
model_inference.set_dim("i", N, coord_values=np.arange(N))
pm.model_to_graphviz(model_inference)

../_images/a552c74da28dab860e6f35b8f9dbc26bd4fa2c1a00a8ac3e14b240ebce3231b4.svg

See the gray matter again. This time we have observed data infused into the model, and we have to sample for the coefficient and other parameters.

So, lets sample.

Step 4. Create samples and validate that the infused relationship between predictors and target variable are captured by the model samples#

with model_inference:
    idata = pm.sample(random_seed=SEED)

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta_cy, beta_by, beta_ay, beta_y0, sigma_y]

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 7 seconds.

Now, lets validate if model captured the infused coefficient values in the data.

pc = az.plot_dist(
    idata,
    var_names=list(true_values.keys()),
    col_wrap=3,
)
az.add_lines(pc, true_values);

../_images/e781d0b94a45fccfc86b9329db342db15650fbe46275dc7397b175d39ba76e4b.png

Pretty nice fit!

Now we can use the do-operator for interventional what-if analysis: generating predictions for the target variable under different hypothetical settings of the predictors.

What if all values of b were 0?

What if all values of b were doubled?

Here’s how.

Step 5. Use `do`-operator to generate target variable with different what-if scenarios#

Since we want to experiment with b, let’s first assign observed values to a and c. Not to y, because that’s what we want to sample.

model_counterfactual = pm.do(model_inference, {"a": df["a"], "c": df["c"]})

Now, let’s begin the fun part — generating what-if predictions.

Scenario 1 :- What if all values for ‘b’ were 0 ?#

model_b0 = pm.do(model_counterfactual, {"b": np.zeros(N, dtype="int32")}, prune_vars=True)
model_b1 = pm.do(model_counterfactual, {"b": df["b"]}, prune_vars=True)

Just sample.

# Sample when 'b' was 0: P(y | (a,c), do(b=0))
idata_b0 = pm.sample_posterior_predictive(
    idata,
    model=model_b0,
    predictions=True,
    var_names=["y_mu"],
    random_seed=SEED,
)
# Sample when 'b' was as observed: P(y | (a,c), do(b=observed))
idata_b1 = pm.sample_posterior_predictive(
    idata,
    model=model_b1,
    predictions=True,
    var_names=["y_mu"],
    random_seed=SEED,
)

Sampling: []

Sampling: []

Some basic Python and here we have the what-if predictions.

df["b_scenario_1"] = 0
df["y_scenario_1"] = (
    idata_b0.predictions.y_mu.mean(("chain", "draw")).values.reshape(1, -1).flatten()
)
df.head(5)

	a	b	c	y	y_scenario_1
0	0.510869	-1.311934	-0.362795	-0.306949	0.681877
1	-0.734647	-0.539776	0.019862	-1.162661	-1.122175
2	-0.358604	0.343649	-0.093558	-0.278148	-0.576914
3	-0.251229	-1.154599	-0.542470	-1.489163	-0.537017
4	-1.316097	0.114966	-0.264377	-2.143315	-2.093052

Scenario 2: What if ‘b’ was 5 times as observed#

model_b0 = pm.do(model_counterfactual, {"b": 5 * df["b"]}, prune_vars=True)
model_b1 = pm.do(model_counterfactual, {"b": df["b"]}, prune_vars=True)

Sample.

# Sample when 'b' was 5 times b: P(y | (a,c), do(b=5*b))
idata_b0 = pm.sample_posterior_predictive(
    idata,
    model=model_b0,
    predictions=True,
    var_names=["y_mu"],
    random_seed=SEED,
)
# Sample when 'b' was as observed: P(y | (a,c), do(b=observed))
idata_b1 = pm.sample_posterior_predictive(
    idata,
    model=model_b1,
    predictions=True,
    var_names=["y_mu"],
    random_seed=SEED,
)

df["b_scenario_2"] = 5 * df["b"]
df["y_scenario_2"] = (
    idata_b0.predictions.y_mu.mean(("chain", "draw")).values.reshape(1, -1).flatten()
)
df.head(5)

Sampling: []

Sampling: []

	a	b	c	y	y_scenario_1	b_scenario_2	y_scenario_2
0	0.510869	-1.311934	-0.362795	-0.306949	0.681877	-6.559670	-3.849755
1	-0.734647	-0.539776	0.019862	-1.162661	-1.122175	-2.698878	-2.986647
2	-0.358604	0.343649	-0.093558	-0.278148	-0.576914	1.718247	0.610107
3	-0.251229	-1.154599	-0.542470	-1.489163	-0.537017	-5.772996	-4.525190
4	-1.316097	0.114966	-0.264377	-2.143315	-2.093052	0.574829	-1.695942

Now you’ve got the idea. It’s an open playground — intervene on whichever variable you like and see how the predicted outcome changes.

This opens the door for many possibilities in causal analytics: comparing treatment policies, sensitivity checks, and programmatic A/B testing.

Authors#

Authored by Shekhar Khandelwal in August 2023
Updated by Osvaldo Martin in February 2026
Updated by Osvaldo Martin in April 2026
Revised wording around interventions and counterfactuals, by Benjamin T. Vincent in March 2026

References#

Pearl, J., Glymour, M., & Jewell, N. P. (2016). Causal Inference in Statistics: A Primer. Wiley. Sections 4.2.3–4.2.4.
Khandelwal, S. (2023). Counterfactuals for Causal Analysis via PyMC Do-Operator. Medium.
PyMC Labs. Causal Analysis with PyMC: Answering “What If” with the New Do-Operator.

Watermark#

%load_ext watermark
%watermark -n -u -v -iv -w -p pytensor

Last updated: Mon, 04 May 2026

Python implementation: CPython
Python version       : 3.14.4
IPython version      : 9.12.0

pytensor: 2.38.0+158.gb3cfb34cb

arviz : 1.1.0
numpy : 2.4.4
pandas: 3.0.2
pymc  : 5.28.0+58.gf58491a3

Watermark: 2.6.0

License notice#

All the notebooks in this example gallery are provided under the MIT License which allows modification, and redistribution for any use provided the copyright and license notices are preserved.

Citing PyMC examples#

To cite this notebook, use the DOI provided by Zenodo for the pymc-examples repository.

Important

Many notebooks are adapted from other sources: blogs, books… In such cases you should cite the original source as well.

Also remember to cite the relevant libraries used by your code.

Here is an citation template in bibtex:

@incollection{citekey,
  author    = "<notebook authors, see above>",
  title     = "<notebook title>",
  editor    = "PyMC Team",
  booktitle = "PyMC examples",
  doi       = "10.5281/zenodo.5654871"
}

which once rendered could look like:

Categories

Tags

Interventional what-if analysis using the PyMC do-operator#

Introduction#

Step 1. Build a pymc model skeleton#

Step 2. Use model skeleton and generate data using do-operator to infuse relationship between predictors and target variable. We will use this generated data for modelling later.#

Step 3. Use observe-operator to assign generated data on the model skeleton#

Step 4. Create samples and validate that the infused relationship between predictors and target variable are captured by the model samples#

Step 5. Use `do`-operator to generate target variable with different what-if scenarios#

Scenario 1 :- What if all values for ‘b’ were 0 ?#

Scenario 2: What if ‘b’ was 5 times as observed#

Authors#

References#

Watermark#

License notice#

Citing PyMC examples#

Categories

Tags

Interventional what-if analysis using the PyMC do-operator#

Introduction#

Step 1. Build a pymc model skeleton#

Step 2. Use model skeleton and generate data using do-operator to infuse relationship between predictors and target variable. We will use this generated data for modelling later.#

Step 3. Use observe-operator to assign generated data on the model skeleton#

Step 4. Create samples and validate that the infused relationship between predictors and target variable are captured by the model samples#

Step 5. Use do-operator to generate target variable with different what-if scenarios#

Scenario 1 :- What if all values for ‘b’ were 0 ?#

Scenario 2: What if ‘b’ was 5 times as observed#

Authors#

References#

Watermark#

License notice#

Citing PyMC examples#

Step 5. Use `do`-operator to generate target variable with different what-if scenarios#