causal-toolkit-yamashita 0.1.2

A Python package for causal inference methods including ATE estimation, propensity score methods, and meta-learners

causal-toolkit-yamashita

A Python package for causal inference methods, including:

• Randomized experiment analysis — average treatment effect (ATE) with confidence intervals and p-values
• Propensity score methods — inverse probability weighting (IPW) and doubly robust estimation
• Meta-learners — S-learner, T-learner, X-learner, and double machine learning for conditional average treatment effects (CATE)

Installation

Install the released package from PyPI:

pip install causal-toolkit-yamashita

Or install from source in editable mode for development:

git clone https://github.com/yamashann/causal-toolkit-yamashita.git
cd causal-toolkit-yamashita
uv pip install -e .

Quick Start

After installing (see above), save the following as quickstart.py and run it with python quickstart.py:

import numpy as np
import pandas as pd

from causal_toolkit_yamashita import calculate_ate_ci, doubly_robust, t_learner_discrete

rng = np.random.default_rng(0)
n = 1000

# --- Randomized experiment: estimate the average treatment effect (ATE) ---
# The DataFrame must have a treatment column "T" (0/1) and an outcome column "Y".
T = rng.integers(0, 2, size=n)
Y = 3.0 * T + rng.normal(size=n)                 # true effect = 3.0
rct = pd.DataFrame({"T": T, "Y": Y})

ate, lo, hi = calculate_ate_ci(rct)
print(f"RCT ATE        = {ate:.2f}  (95% CI [{lo:.2f}, {hi:.2f}])")

# --- Observational data where covariates confound the treatment ---
x1, x2 = rng.normal(size=n), rng.normal(size=n)
treat = rng.binomial(1, 1 / (1 + np.exp(-(0.8 * x1 + 0.5 * x2))))
outcome = 2.0 * treat + 1.5 * x1 + x2 + rng.normal(size=n)   # true effect = 2.0
obs = pd.DataFrame({"x1": x1, "x2": x2, "treat": treat, "outcome": outcome})

# doubly_robust(df, formula, T, Y) — `formula` is a patsy formula of the covariates
print(f"Doubly-robust  = {doubly_robust(obs, 'x1 + x2', 'treat', 'outcome'):.2f}")

# --- Per-row heterogeneous effects (CATE) via a meta-learner ---
train, test = obs.iloc[:800], obs.iloc[800:]
cate = t_learner_discrete(train, test, X=["x1", "x2"], T="treat", y="outcome")
print(f"Mean CATE      = {cate['cate'].mean():.2f}  ({len(cate)} test rows)")

Expected output (estimates land near the true effects baked into the data):

RCT ATE        = 3.12  (95% CI [3.00, 3.25])
Doubly-robust  = 2.08
Mean CATE      = 2.16  (200 test rows)

All eight functions (calculate_ate_ci, calculate_ate_pvalue, ipw, doubly_robust, s_learner_discrete, t_learner_discrete, x_learner_discrete, double_ml_cate) are importable directly from causal_toolkit_yamashita; see the API table below for signatures.

API

rct — Randomized experiments

Function	Description
calculate_ate_ci(data, alpha=0.05)	Returns (ate, ci_lower, ci_upper)
calculate_ate_pvalue(data)	Returns (ate, t_stat, p_value)

propensity — Propensity score methods

Function	Description
ipw(df, ps_formula, T, Y)	Inverse probability weighted ATE
doubly_robust(df, formula, T, Y)	Doubly robust ATE estimate

meta_learners — CATE estimation

Function	Description
s_learner_discrete(train, test, X, T, y)	Single-model learner; returns DataFrame with cate column
t_learner_discrete(train, test, X, T, y)	Two-model learner
x_learner_discrete(train, test, X, T, y)	Cross-fitted X-learner
double_ml_cate(train, test, X, T, y)	Double machine learning CATE

Running tests

uv run pytest

License

MIT