Add routine for estimating ATE to example

Author: cthoyt

src/y0/examples/utils.py

@@ -33,3 +33,21 @@ class Example:
     #: what might be interesting to use in the ID algorithm
     example_queries: Optional[list[Query]] = None
     generate_data: Optional[Callable[[int, Optional[dict[Variable, float]]], pd.DataFrame]] = None
+
+    def generate_ate(
+        self,
+        *,
+        num_samples: int,
+        treatment: Variable,
+        outcome: Variable,
+        treatment_0: float = 0.0,
+        treatment_1: float = 1.0,
+        **kwargs,
+    ) -> float:
+        """Calculate the ATE for a single treatment/outcome pair."""
+        if self.generate_data is None:
+            raise TypeError(f"no generation method provided in example: {self.name}")
+
+        data_1 = self.generate_data(num_samples, {treatment: treatment_1}, **kwargs)
+        data_0 = self.generate_data(num_samples, {treatment: treatment_0}, **kwargs)
+        return data_1.mean()[outcome.name] - data_0.mean()[outcome.name]