diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bd987b68..53378c20 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,15 @@
Changelog
=========
+0.9.1 (2024-07-xx)
+------------------
+
+**New features**
+
+* Added :meth:`metalearners.metalearner.DRLearner.treatment_effect` to compute AIPW point estimate and standard error without requiring a full model fit.
+
+
+
0.9.0 (2024-07-xx)
------------------
diff --git a/docs/examples/example_estimating_ates.ipynb b/docs/examples/example_estimating_ates.ipynb
new file mode 100644
index 00000000..e6539188
--- /dev/null
+++ b/docs/examples/example_estimating_ates.ipynb
@@ -0,0 +1,508 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "(example-ate)=\n",
+ "\n",
+ " Example: Estimating Average Treatment Effects\n",
+ "=============================\n",
+ "\n",
+ "Motivation\n",
+ "----------\n",
+ "\n",
+ "Estimating average treatment effects (ATEs) involves a subset of the tasks involved in estimating Conditional Average Treatment Effects (CATEs), so we can use methods that are designed for estimating CATEs to estimate ATEs. In this example, we simulate some data with confounding and demonstrate the `treatment_effect` method of the `DRLearner` class, which estimates the ATE, and compare it to estimates from some other popular libraries (`econML` and `doubleML`).\n",
+ "\n",
+ "Example\n",
+ "-------"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import pandas as pd\n",
+ "import numpy as np\n",
+ "import statsmodels.formula.api as smf"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# DGP with confounding"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# generate covariate matrix with mixture of continuous and categorical variables\n",
+ "np.random.seed(123)\n",
+ "\n",
+ "\n",
+ "def dgp(n, k, pscore_fn, tau_fn, outcome_fn, k_cat=1):\n",
+ " \"\"\"DGP for a confounded treatment assignment dgp\n",
+ "\n",
+ " Args:\n",
+ " n (int): sample size\n",
+ " k (int): number of continuous covariates\n",
+ " pscore_fn (lambda): propensity score function\n",
+ " tau_fn (lambda): treatment effect function. Can be scalar for constant effect.\n",
+ " outcome_fn (lambda): outcome DGP\n",
+ "\n",
+ " Returns:\n",
+ " _type_: _description_\n",
+ " \"\"\"\n",
+ " Sigma = np.random.uniform(-1, 1, (k, k))\n",
+ " Sigma = Sigma @ Sigma.T\n",
+ " Xnum = np.random.multivariate_normal(np.zeros(k), Sigma, n)\n",
+ " # generate categorical variables\n",
+ " Xcat = np.random.binomial(1, 0.5, (n, k_cat))\n",
+ " X = np.c_[Xnum, Xcat]\n",
+ " W = np.random.binomial(1, pscore_fn(X), n)\n",
+ " Y = outcome_fn(X, W, tau_fn)\n",
+ " df = pd.DataFrame(\n",
+ " np.c_[Y, W, X], columns=[\"Y\", \"W\"] + [f\"X{i}\" for i in range(k + 1)]\n",
+ " )\n",
+ " return df\n",
+ "\n",
+ "\n",
+ "pscore_fn = lambda x: 1 / (1 + np.exp(-x[:, 0] - x[:, 1] - x[:, 2] ** 2 + x[:, 3]))\n",
+ "\n",
+ "\n",
+ "# tau_fn = lambda x: 1 + 2 * x[:, 0] + 3 * x[:, 1] + 4 * x[:, 2] + 5 * x[:, 3]\n",
+ "def outcome_fn(x, w, taufn):\n",
+ " return (\n",
+ " taufn(x) * w\n",
+ " + x[:, 0]\n",
+ " + 2 * x[:, 1] ** 2\n",
+ " + 3 * x[:, 3] * x[:, 1]\n",
+ " + x[:, 2]\n",
+ " + x[:, 3]\n",
+ " + np.random.normal(0, 1, n)\n",
+ " )\n",
+ "\n",
+ "\n",
+ "n, k = 10_000, 3\n",
+ "df = dgp(n, k, pscore_fn, tau_fn=lambda x: 1, outcome_fn=outcome_fn)\n",
+ "outcome_column, treatment_column = \"Y\", \"W\"\n",
+ "feature_columns = [f\"X{i}\" for i in range(k + 1)]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Linear Regression "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "(1.624943884109307, 0.04532725031682251)"
+ ]
+ },
+ "execution_count": 3,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "naive_lm = smf.ols(f\"{outcome_column} ~ {treatment_column}\", df) .fit(cov_type=\"HC1\")\n",
+ "naive_est = naive_lm.params.iloc[1], naive_lm.bse.iloc[1]\n",
+ "naive_est"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "(1.1326349969274776, 0.02972906033475406)"
+ ]
+ },
+ "execution_count": 4,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "covaradjust_lm = smf.ols(f\"{outcome_column} ~ {treatment_column}+{'+'.join(feature_columns)}\",\n",
+ " df) .fit(cov_type=\"HC1\")\n",
+ "linreg_est = covaradjust_lm.params.iloc[1], covaradjust_lm.bse.iloc[1]\n",
+ "linreg_est"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Linear model is misspecified, so both the naive and conditional estimates are biased."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## `metalearners`: `DRLearner`\n",
+ "\n",
+ "Point estimates and standard errors for treatment effects for the AIPW estimator can be computed by aggregating the pseudo-outcome computed by the `DRLearner` class."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n"
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "from metalearners import DRLearner\n",
+ "from lightgbm import LGBMRegressor, LGBMClassifier\n",
+ "from sklearn.dummy import DummyRegressor"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "array([[1.01467296, 0.03699228]])"
+ ]
+ },
+ "execution_count": 6,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "metalearners_dr = DRLearner(\n",
+ " nuisance_model_factory=LGBMRegressor,\n",
+ " treatment_model_factory=DummyRegressor, # not actually used since we don't fit treatment model\n",
+ " propensity_model_factory=LGBMClassifier,\n",
+ " is_classification=False,\n",
+ " n_variants=2,\n",
+ " nuisance_model_params={\"verbose\": -1},\n",
+ " propensity_model_params={\"verbose\": -1},\n",
+ ")\n",
+ "\n",
+ "metalearners_dr.fit_all_nuisance(\n",
+ " X=df[feature_columns],\n",
+ " y=df[outcome_column],\n",
+ " w=df[treatment_column],\n",
+ ")\n",
+ "\n",
+ "metalearners_est = metalearners_dr.treatment_effect( # still need to pass data objects since DRLearner does not retain any data\n",
+ " X=df[feature_columns],\n",
+ " w=df[treatment_column],\n",
+ " y=df[outcome_column],\n",
+ ")\n",
+ "metalearners_est"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Manual computation with pseudo outcome method produces the same estimate (`treatment_effect` does a generalisation of this under the hood) yields the same estimate"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "est: 1.0146729600419584, se: 0.036994131087587\n"
+ ]
+ }
+ ],
+ "source": [
+ "gamma_i = metalearners_dr._pseudo_outcome(\n",
+ " X=df[feature_columns],\n",
+ " w=df[treatment_column],\n",
+ " y=df[outcome_column],\n",
+ " treatment_variant=1,\n",
+ " is_oos=False,\n",
+ ")\n",
+ "gamma_i.mean(), gamma_i.std()/np.sqrt(n)\n",
+ "est, se = gamma_i.mean(), gamma_i.std()/np.sqrt(n)\n",
+ "print(f\"est: {est}, se: {se}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## `doubleml`: `DoubleMLIRM`\n",
+ "\n",
+ "The `doubleML` library focuses on estimating average effects and has an 'interactive regression model (IRM)' class that estimates the ATE using the same pseudo-outcome method as the `DRLearner` class."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%capture\n",
+ "from doubleml import DoubleMLIRM, DoubleMLData\n",
+ "dml_data = DoubleMLData(\n",
+ " df,\n",
+ " x_cols=feature_columns,\n",
+ " y_col=outcome_column,\n",
+ " d_cols=treatment_column,\n",
+ ")\n",
+ "\n",
+ "aipw_mod = DoubleMLIRM(\n",
+ " dml_data,\n",
+ " ml_g = LGBMRegressor(),\n",
+ " ml_m = LGBMClassifier(),\n",
+ " n_folds=5,\n",
+ ")\n",
+ "\n",
+ "aipw_mod.fit()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[1.01366331 0.03880938]\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(doubleml_est := aipw_mod.summary.values[0, :2])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## `econML`: `LinearDRLearner`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from econml.dr import LinearDRLearner\n",
+ "import formulaic as fm"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Y ~ 0 + X0+X1+X2+X3\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(ff := f\"{outcome_column} ~ 0 + {'+'.join(feature_columns)}\")\n",
+ "y, X = fm.Formula(ff).get_model_matrix(df, output=\"numpy\")\n",
+ "W = df[treatment_column].values[:, np.newaxis]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%capture\n",
+ "econml_dr = LinearDRLearner(model_regression=LGBMRegressor(), model_propensity=LGBMClassifier())\n",
+ "econml_dr.fit(y, T=W, W=X)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[1.069 0.059]\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(econml_est := econml_dr.intercept__inference(1).summary_frame().iloc[0, :2].values)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## comparison\n",
+ "\n",
+ "All ml-based estimators yield comparable results."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " naive | \n",
+ " linreg | \n",
+ " metalearners | \n",
+ " doubleml | \n",
+ " econml | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | est | \n",
+ " 1.624944 | \n",
+ " 1.132635 | \n",
+ " 1.014673 | \n",
+ " 1.013663 | \n",
+ " 1.069 | \n",
+ "
\n",
+ " \n",
+ " | 0.036994 | \n",
+ " 0.045327 | \n",
+ " 0.029729 | \n",
+ " 0.036992 | \n",
+ " 0.038809 | \n",
+ " 0.059 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "