Merge pull request #2 from kklein/ate

Minor adaptations to your PR

.github/workflows/ci.yml

@@ -1,5 +1,15 @@
 name: CI
-on: [push]
+on:
+  # We would like to trigger for CI for any pull request action -
+  # both from QuantCo's branches as well as forks.
+  pull_request:
+  # In addition to pull requests, we want to run CI for pushes
+  # to the main branch and tags.
+  push:
+    branches:
+      - "main"
+    tags:
+      - "*"
 
 # Automatically stop old builds on the same branch/PR
 concurrency:

CHANGELOG.rst

@@ -12,7 +12,21 @@ Changelog
 
 **New features**
 
-* Added :meth:`metalearners.metalearner.MetaLearner.init_params`.
+* Add :meth:`metalearners.metalearner.MetaLearner.init_params`.
+
+* Add :class:`metalearners.utils.FixedBinaryPropensity`.
+
+* Added ``_build_onnx`` to :class:`metalearners.MetaLearner` abstract class and implement it
+  for :class:`metalearners.TLearner`, :class:`metalearners.XLearner`, :class:`metalearners.RLearner`
+  and :class:`metalearners.DRLearner`.
+
+* Added ``_necessary_onnx_models`` to :class:`metalearners.MetaLearner`.
+
+* Added :meth:`metalearners.metalearner.DRLearner.average_treatment_effect` to
+  compute AIPW point estimate and standard error for _average
+  treatment effects (ATE)_
+  without requiring a full model fit (which is required for CATE
+  estimation). A new notebook contains examples.
 
 * Added :meth:`metalearners.metalearner.DRLearner.treatment_effect` to compute AIPW point estimate and standard error for _average treatment effects (ATE)_ without requiring a full model fit (which is required for CATE estimation). A new notebook contains examples.
 
@@ -21,7 +35,7 @@ Changelog
 
 **New features**
 
-* Added :meth:`metalearners.metalearner.MetaLearner.fit_all_nuisance` and
+* Add :meth:`metalearners.metalearner.MetaLearner.fit_all_nuisance` and
   :meth:`metalearners.metalearner.MetaLearner.fit_all_treatment`.
 
 * Add optional ``store_raw_results`` and ``store_results`` parameters to :class:`metalearners.grid_search.MetaLearnerGridSearch`.

docs/examples/example_estimating_ates.ipynb

@@ -14,7 +14,7 @@
     "\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. \n",
     "\n",
-    "In this example, we simulate some data with a confounded binary treatment 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`). We then show how this approach generalizes to a setting with multiple discrete treatments, and finally to a setting with discrete-valued outcomes.\n",
+    "In this example, we simulate some data with a confounded binary treatment and demonstrate the `average_treatment_effect` method of the `DRLearner` class, which estimates the ATE, and compare it to estimates from some other popular libraries (`econML` and `doubleML`). We then show how this approach generalizes to a setting with multiple discrete treatments, and finally to a setting with discrete-valued outcomes.\n",
     "\n",
     "Example\n",
     "-------"
@@ -236,10 +236,11 @@
     "    y=df[outcome_column],\n",
     "    w=df[treatment_column],\n",
     ")\n",
-    "metalearners_est = metalearners_dr.treatment_effect( # still need to pass data objects since DRLearner does not retain any data\n",
+    "metalearners_est = metalearners_dr.average_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",
+    "    is_oos=False,\n",
     ")\n",
     "metalearners_est"
    ]
@@ -248,7 +249,7 @@
    "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"
+    "Manual computation with pseudo outcome method produces the same estimate (`average_treatment_effect` does a generalisation of this under the hood) yields the same estimate"
    ]
   },
   {
@@ -713,10 +714,11 @@
     "    y=df_multi[outcome_column],\n",
     "    w=df_multi[treatment_column],\n",
     ")\n",
-    "metalearners_est2 = metalearners_dr_2.treatment_effect( # still need to pass data objects since DRLearner does not retain any data\n",
+    "metalearners_est2 = metalearners_dr_2.average_treatment_effect( # still need to pass data objects since DRLearner does not retain any data\n",
     "    X=df_multi[feature_columns],\n",
     "    y=df_multi[outcome_column],\n",
     "    w=df_multi[treatment_column],\n",
+    "    is_oos=False,\n",
     ")\n",
     "metalearners_est2"
    ]