Quantco / metalearners

MetaLearners for CATE estimation
https://metalearners.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
34 stars 4 forks source link

Implement `predict_conditional_average_outcomes` for `RLearner`? #67

Closed FrancescMartiEscofetQC closed 3 months ago

FrancescMartiEscofetQC commented 4 months ago

All implemented MetaLearners allow the user to call predict_conditional_average_outcomes. At the beginning we thought this was not possible for the RLearner but I think the following formulas may work:

(For ease of notation I'll use $Y(k) := \mathbb{E}[Y_i(k)]$, $Y = \mathbb{E}[Y | X])$, $\tau(k) = \mathbb{E}[Y(k) - Y(0) | X]$ and $e(k) = \mathbb{P}[W = k | X]$)

We know this system of $K$ linear equations is true:

\begin{cases}
      Y(1) - Y(0) = \tau(1)\\
      Y(2) - Y(0) = \tau(2)\\
      \vdots \\
      Y(K) - Y(0) = \tau(K) \\
      e(0) Y(0) + e(1) Y(1) + \dots + e(K) Y(K) = Y
\end{cases} 

that we need to solve for $Y(0), Y(1), \dots, Y(K)$.

Isolating $Y(1), Y(2), \dots, Y(K)$ from each of the first $K-1$ equations and plugging it into the last we get:

e(0) Y(0) + e(1) (\tau(1) + Y(0)) + \dots + e(K) (\tau(K) + Y(0)) = Y

From this we can isolate $Y(0)$ as:

Y(0) = \frac{Y - \sum\limits_{i=1}^{K}e(i)\tau(i)}{e(0) + \sum\limits_{i=1}^{K} e(i)} = Y - \sum\limits_{i=1}^{K}e(i)\tau(i)

Where we used the fact that all the propensity scores should sum up to 1.

Finally we can compute all $Y(k)$ as $Y(k) = Y(0) + \tau(k)$.

I extracted this idea for the binary case from this code snippet from this reference: Screenshot 2024-07-22 at 17 30 34

Any thoughts on this @kklein ?