From Formula To Function A Systematic Way To Turn

You’ve stared at a formula and felt your brain lock up. Not because the math is hard, but because the formula is a picture of the destination, not a roadmap for how to get there. It tells you what the answer is, but not how to compute it step by step.

Think of it this way: a formula like ATE_Hajek = sum(T*Y/e)/sum(T/e) - sum((1-T)*Y/(1-e))/sum((1-T)/(1-e)) is like looking at a blueprint of a house. It shows you the walls and the roof, but it doesn’t tell you how to pour the foundation or hang the drywall. This is the hardest part of causal inference — and it’s almost never taught.

Promise: By the end of this tutorial, you’ll have a repeatable 5-step method — what I call the Step 0 Formula Extension — to turn any formula into runnable Python. No more staring. No more freezing.

What We’re Building: The Hajek ATE Estimator

Before we write a single line of code, let’s make sure we understand what we’re actually building. In plain English, the ATE (Average Treatment Effect) answers this question: “What would be the average difference in outcomes if everyone received the treatment versus if no one did?” It’s the holy grail of causal inference.

IPW (Inverse Probability Weighting) is a method to estimate the ATE from observational data where treatment isn’t randomly assigned. The trick: we reweight each observation by the inverse of its probability of receiving the treatment it actually got. This creates a ‘pseudo-population’ where treatment is independent of covariates — like a randomized experiment.

Now here’s the interesting part: the Hajek form is a more stable version of simple IPW. Instead of just summing the weighted outcomes, it normalizes the weights by dividing by their sum. This makes the estimate more robust when weights are very large or very small.

Let’s define the variables we’re working with:

• Y: outcome (float array, length n)
• T: binary treatment (0 or 1, length n)
• X: covariates (matrix, n x k)
• e: propensity score (probability of receiving treatment given X, length n)

The formula looks like this:

# ATE_Hajek = (sum(T * Y / e) / sum(T / e)) - (sum((1-T) * Y / (1-e)) / sum((1-T) / (1-e)))

This is the exact formula we’ll translate into code. But first, we need to stop and plan.

Step 0a: Copy the Formula Verbatim Into a Comment Block

Here’s the first and most overlooked step: get the formula into your code as a comment before you write a single line of logic. Why verbatim? Because you’ll refer back to it constantly. Typing it out forces you to see every symbol.

# ============================================================
# Step 0a: Formula — verbatim from paper
# ATE_Hajek = (sum(T * Y / e) / sum(T / e)) - (sum((1-T) * Y / (1-e)) / sum((1-T) / (1-e)))
# ============================================================

This is the ‘ground truth’ you’ll check your code against. Every time you get lost, you come back here.

Step 0b: Map Every Symbol to a Named Python Variable

Formulas use single letters. Code needs descriptive names with types. This mapping is where many errors start.

# Step 0b: Symbol-to-variable mapping
# T        → treatment     (np.ndarray, shape (n,), dtype int)
# Y        → outcome       (np.ndarray, shape (n,), dtype float)
# X        → covariates    (np.ndarray, shape (n, k), dtype float)
# e        → propensity    (np.ndarray, shape (n,), dtype float)

Why do numpy types matter? Shape and dtype errors are the most common bugs in this translation. If you try to divide a (n, 1) array by a (n,) array, things break silently. Specifying types upfront prevents that.

Step 0c: Name All Intermediate Quantities as Variables

This is the most important sub-step. Formulas have implicit intermediate results that must be named before they can become code. Naming them forces you to think about the execution order.

# Step 0c: Intermediate quantities (names, shapes, types)
# treated_weights          → T / e                           (n,) float
# control_weights          → (1-T) / (1-e)                   (n,) float
# sum_treated_weighted_Y   → sum(T * Y / e)                  scalar float
# sum_treated_weights      → sum(T / e)                      scalar float
# treated_mean             → sum_treated_weighted_Y / sum_treated_weights   scalar float
# sum_control_weighted_Y   → sum((1-T) * Y / (1-e))          scalar float
# sum_control_weights      → sum((1-T) / (1-e))              scalar float
# control_mean             → sum_control_weighted_Y / sum_control_weights   scalar float
# ate_hajek                → treated_mean - control_mean     scalar float

Why naming matters: it makes the computation debuggable. You can print each intermediate and check it against your mental model. Without names, you’re debugging a spaghetti line of parentheses.

Step 0d: Write Pseudocode, One Comment per Intermediate Quantity

Now write the execution plan in plain English comments, one line per intermediate. This is the ‘algorithm’ before the code.

# Step 0d: Pseudocode — execution plan in plain English
# 1. Fit logistic regression to get propensity scores e = P(T=1|X)
# 2. Clip propensity scores to avoid extreme weights (e = clip(e, 1e-6, 1-1e-6))
# 3. Compute treated weights = T / e
# 4. Compute control weights = (1 - T) / (1 - e)
# 5. Compute weighted sum of outcomes for treated: sum(T * Y / e)
# 6. Compute sum of treated weights: sum(T / e)
# 7. Compute weighted mean for treated: step5 / step6
# 8. Repeat steps 5-7 for control group
# 9. ATE = treated_mean - control_mean

This pseudocode is the ‘recipe’ your code will follow. It’s the bridge between the formula and the execution.

Step 0e: State the Simulation-Verification Plan Before Writing Any Code

Before you write a line of executable code, plan how you’ll test it. Simulate data with known ground truth and assert recovery within a tolerance.

# Step 0e: Test plan — verification before implementation
# Simulation parameters:
#   n = 2000  (observations)
#   true_ATE = 2.0
#   DGP: generate Y from a data-generating process that includes treatment effect + noise
# Assertion:
#   |estimated_ATE - 2.0| < 0.3
# This is not optional — it's the only way to know your translation is correct.

Why simulate first? Because you need to know the true answer to verify your code is correct. Without ground truth, you’re coding blind.

Putting It All Together: From Pseudocode to Python

Now write the actual code, line by line, following the pseudocode. We’ll walk through each step and interpret every number in plain English.

import numpy as np
from sklearn.linear_model import LogisticRegression
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# Step 0a: Formula — verbatim from paper
# ATE_Hajek = (sum(T * Y / e) / sum(T / e)) - (sum((1-T) * Y / (1-e)) / sum((1-T) / (1-e)))
# ============================================================

# Generate synthetic data with known true ATE
np.random.seed(42)
n = 2000

# Two correlated covariates
X1 = np.random.normal(0, 1, n)
X2 = 0.5 * X1 + np.random.normal(0, 0.5, n)
X = np.column_stack([X1, X2])

# True propensity: logistic function of X
beta_t = np.array([0.5, -0.3])
linear_pred = X @ beta_t
true_e = 1 / (1 + np.exp(-linear_pred))

# Assign treatment
T = np.random.binomial(1, true_e)

# Generate outcome with treatment effect = 2.0
baseline = 1.0 + 0.5 * X1 + 0.3 * X2
treatment_effect = 2.0 + 0.2 * X1
true_ate = np.mean(treatment_effect)  # population ATE ≈ 2.0
Y = baseline + treatment_effect * T + np.random.normal(0, 0.5, n)

print(f"True ATE (ground truth): {true_ate:.4f}")
print(f"Sample size: {n}")
print(f"Proportion treated: {T.mean():.3f}")
print()

# ============================================================
# Step 0b: Map symbols to variables
# T -> treatment, Y -> outcome, X -> covariates, e -> propensity
# ============================================================

# Fit propensity model (step 1 of pseudocode)
model = LogisticRegression()
model.fit(X, T)
e_hat = model.predict_proba(X)[:, 1]

# Clip propensity scores (step 2 of pseudocode)
# This prevents extreme weights that cause high variance
epsilon = 1e-6
e_hat = np.clip(e_hat, epsilon, 1 - epsilon)

print(f"Propensity score range: [{e_hat.min():.4f}, {e_hat.max():.4f}]")
print()

# ============================================================
# Step 0c: Name intermediate quantities
# ============================================================

# Treated group calculations
treated_weights = T / e_hat  # Step 3
print(f"Treated weights: mean = {treated_weights.mean():.4f}, std = {treated_weights.std():.4f}")
print(f"  ▸ This means treated units are upweighted by {treated_weights.mean():.2f}x on average.")

sum_treated_weighted_Y = np.sum(T * Y / e_hat)  # Step 5
sum_treated_weights = np.sum(T / e_hat)  # Step 6
treated_mean = sum_treated_weighted_Y / sum_treated_weights  # Step 7

print(f"  ▸ Weighted sum of treated outcomes: {sum_treated_weighted_Y:.2f}")
print(f"  ▸ Sum of treated weights: {sum_treated_weights:.2f}")
print(f"  ▸ Weighted mean outcome for treated: {treated_mean:.4f}")
print()

# Control group calculations (Step 8)
control_weights = (1 - T) / (1 - e_hat)
print(f"Control weights: mean = {control_weights.mean():.4f}, std = {control_weights.std():.4f}")
print(f"  ▸ Control units are upweighted by {control_weights.mean():.2f}x on average.")

sum_control_weighted_Y = np.sum((1 - T) * Y / (1 - e_hat))
sum_control_weights = np.sum((1 - T) / (1 - e_hat))
control_mean = sum_control_weighted_Y / sum_control_weights

print(f"  ▸ Weighted sum of control outcomes: {sum_control_weighted_Y:.2f}")
print(f"  ▸ Sum of control weights: {sum_control_weights:.2f}")
print(f"  ▸ Weighted mean outcome for control: {control_mean:.4f}")
print()

# ============================================================
# Step 0d: Followed pseudocode — now compute ATE (Step 9)
# ============================================================

ate_hajek = treated_mean - control_mean

print("=" * 55)
print(f"Estimated ATE (Hajek): {ate_hajek:.4f}")
print(f"True ATE:             {true_ate:.4f}")
print(f"Estimation error:     {ate_hajek - true_ate:.4f}")

Let’s check what happened. The estimated ATE is close to 2.0. The estimation error is small — within about 0.1 of the true value. That’s not bad for a single simulation with 2000 observations.

Notice how we printed each intermediate and gave it a plain-English interpretation. This is the key to debugging and understanding. If the final ATE were wrong, you could look at every intermediate to find where things broke.

Running the Verification: Does Our Code Recover the Truth?

Now the moment of truth. We run our simulation with the assertion we planned in Step 0e.

# Step 0e verification: run simulation and check assertion
np.random.seed(42)
n = 2000

# Regenerate data (same as above, self-contained for clarity)
X1 = np.random.normal(0, 1, n)
X2 = 0.5 * X1 + np.random.normal(0, 0.5, n)
X = np.column_stack([X1, X2])
beta_t = np.array([0.5, -0.3])
true_e = 1 / (1 + np.exp(-X @ beta_t))
T = np.random.binomial(1, true_e)
baseline = 1.0 + 0.5 * X1 + 0.3 * X2
treatment_effect = 2.0 + 0.2 * X1
true_ate = np.mean(treatment_effect)
Y = baseline + treatment_effect * T + np.random.normal(0, 0.5, n)

# Full Hajek estimator (copy-pasted from above for self-containment)
model = LogisticRegression()
model.fit(X, T)
e_hat = model.predict_proba(X)[:, 1]
e_hat = np.clip(e_hat, 1e-6, 1 - 1e-6)

sum_treated_weighted_Y = np.sum(T * Y / e_hat)
sum_treated_weights = np.sum(T / e_hat)
treated_mean = sum_treated_weighted_Y / sum_treated_weights

sum_control_weighted_Y = np.sum((1 - T) * Y / (1 - e_hat))
sum_control_weights = np.sum((1 - T) / (1 - e_hat))
control_mean = sum_control_weighted_Y / sum_control_weights

ate_hajek = treated_mean - control_mean

# Assertion: recover true ATE within tolerance
tolerance = 0.3
assert abs(ate_hajek - true_ate) < tolerance, f"ATE {ate_hajek:.4f} too far from true {true_ate:.4f}"

print(f"Estimated ATE: {ate_hajek:.4f}")
print(f"True ATE:      {true_ate:.4f}")
print(f"Error:         {ate_hajek - true_ate:.4f}")
print(f"Within tolerance (±{tolerance}): ✅ PASS")

The assertion passes. The code works. If it had failed, what would we do? Debugging strategy: check the intermediates. Is the propensity model fitting correctly? Is the clipping clipping too aggressively? Is the DGP generating data that matches our assumptions? Each intermediate we named is a potential debugging hook.

Why This Method Works: The Psychology of Translation

Why is the Step 0 Formula Extension so effective, especially for readers with weak math backgrounds? It’s not just about code — it’s about reducing cognitive load.

Here’s what’s actually going on in your brain: Formulas are compressed information. They pack multiple operations into a single line. Code requires decompressed, sequential steps — one operation at a time, from top to bottom. The gap between compression and decompression is where you get stuck.

• The verbatim copy anchors your brain to the original. You can always look back and check.
• The symbol mapping reduces ambiguity. No more ‘what does e mean again?’
• Naming intermediates forces you to think about the computation graph — what depends on what.
• Pseudocode is a ‘mental scaffold’ that makes the code feel familiar before you write it.
• The simulation plan gives you a ‘sanity check’ — without it, you’re coding blind.

This method works for any formula. Not just causal inference. Try it on a loss function, a gradient update, or a statistical test. The process is the same.

What You’ve Learned — And Where This Method Applies Next

Let’s recap the 5 sub-steps of Step 0 Formula Extension:

1. Step 0a: Copy the formula verbatim into a comment block.
2. Step 0b: Map every symbol to a named Python variable with type and shape.
3. Step 0c: Name all intermediate quantities as variables.
4. Step 0d: Write pseudocode, one comment per intermediate quantity.
5. Step 0e: Plan your simulation-verification before writing any code.

This method is a systematic way to turn the compressed language of math into the sequential language of code. It works because it forces you to think about execution order, intermediate states, and verification — the three things that make code correct.

In the next part of the series, we’ll apply this same method to a more complex formula: the doubly robust ATE estimator (AIPW), which combines two models — a propensity model and an outcome model. The Step 0 method will make that translation just as smooth.

Explore the next tutorial in the series to see this method applied to a formula with two nuisance models.

Check Your Understanding

Remember: What are the 5 sub-steps of the Step 0 Formula Extension? List them in order.

Understand: In your own words, explain why formulas are ‘compressed information’ and code is ‘decompressed information.’

Apply: Take this formula for a simple weighted average: weighted_mean = sum(w * x) / sum(w). Walk through Steps 0a through 0e to plan a Python function that computes it.

Analyze: In the Hajek estimator, what would happen if we skipped the clipping step? Which intermediate quantities would be affected, and how would the final ATE change?

Evaluate: Compare the Step 0 method to ‘just writing the code directly.’ What are the trade-offs? When might you skip the pseudocode step?

Create: Design a Step 0 plan for a formula of your choice — maybe a gradient descent update, a correlation coefficient, or a log-likelihood function. Write out Steps 0a through 0e in full.