Prepare release 0.9.0 (#76)

* Update changelog.

* Return tuple of arrays instead of concatenated arrays.

CHANGELOG.rst

@@ -7,7 +7,7 @@
 Changelog
 =========
 
-0.9.0 (2024-07-xx)
+0.9.0 (2024-08-02)
 ------------------
 
 **New features**
@@ -16,19 +16,16 @@ Changelog
 
 * Add :class:`metalearners.utils.FixedBinaryPropensity`.
 
-* Added ``_build_onnx`` to :class:`metalearners.MetaLearner` abstract class and implement it
+* Add ``_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`.
+* Add ``_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.
+* Add :meth:`metalearners.metalearner.DRLearner.average_treatment_effect` to
+  compute the AIPW point estimate and standard error for
+  _average treatment effects (ATE)_ without requiring a full model fit.
 
-* 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.
 
 0.8.0 (2024-07-22)
 ------------------

docs/examples/example_estimating_ates.ipynb

@@ -310,7 +310,7 @@
     " np.c_[\n",
     "    naive_est,\n",
     "    linreg_est,\n",
-    "    metalearners_est.flatten(),\n",
+    "    np.hstack(metalearners_est),\n",
     "    doubleml_est,\n",
     "    econml_est,\n",
     "], index = ['est', 'se'],\n",

metalearners/drlearner.py

@@ -344,11 +344,11 @@ def average_treatment_effect(
         y: Vector,
         w: Vector,
         is_oos: bool,
-    ) -> np.ndarray:
+    ) -> tuple[np.ndarray, np.ndarray]:
         """Compute Average Treatment Effect (ATE) for each treatment variant using the
         Augmented IPW estimator (Robins et al 1994). Does not require fitting a second-
-        stage treatment model: it uses the pseudo-outcome alone and computes the average
-        and SE. Can be used following the
+        stage treatment model: it uses the pseudo-outcome alone and computes the point
+        estimate and standard error. Can be used following the
         :meth:`~metalearners.drlearner.DRLearner.fit_all_nuisance` method.
 
         Args:
@@ -358,7 +358,8 @@ def average_treatment_effect(
             is_oos (bool): indicator whether data is out of sample
 
         Returns:
-            np.ndarray: Treatment effect and standard error for each treatment variant.
+            np.ndarray: Treatment effect for each treatment variant.
+            np.ndarray: Standard error for each treatment variant.
         """
         if not self._nuisance_models_fit:
             raise ValueError(
@@ -375,7 +376,7 @@ def average_treatment_effect(
             )
         treatment_effect = gamma_matrix.mean(axis=0)
         standard_error = gamma_matrix.std(axis=0) / np.sqrt(len(X))
-        return np.c_[treatment_effect, standard_error]
+        return treatment_effect, standard_error
 
     def _pseudo_outcome(
         self,

tests/test_drlearner.py

@@ -135,7 +135,7 @@ def test_drlearner_onnx(
     np.testing.assert_allclose(ml.predict(X, True, "overall"), pred_onnx, atol=5e-4)
 
 
-def test_treatment_effect(
+def test_average_treatment_effect(
     numerical_experiment_dataset_continuous_outcome_binary_treatment_linear_te,
 ):
     X, _, W, Y, _, tau = (
@@ -150,5 +150,5 @@ def test_treatment_effect(
         n_folds=2,
     )
     ml.fit_all_nuisance(X, Y, W)
-    est = ml.average_treatment_effect(X, Y, W, is_oos=False)
-    np.testing.assert_almost_equal(est[:, 0], tau.mean(), decimal=1)
+    ate_estimate, _ = ml.average_treatment_effect(X, Y, W, is_oos=False)
+    np.testing.assert_almost_equal(ate_estimate, tau.mean(), decimal=1)