[Feature Request]: Support for multiple discrete treatments in `doubleml.DoubleMLIRM`

[apoorvalal]

Describe the feature you want to propose or implement

Would be nice to extend the DoubleMLIRM class to accommodate multiple discrete-valued treatments $w \in \mathcal{W}$, where the estimation primatives are the marginal counterfactual mean $Y^{(w)}$ under each treatment, and one can construct treatment effects as contrasts between them. The marginal mean construction is the same as in the binary treatment case

$$ \phi_i^{(w)} = \mu^{(w)}(x) + \frac{1(W_i = w)}{\pi^{w}(x)} \left(Y_i - \mu^{(w)}(x)\right) $$

which can be constructed at the unit level, and the subsequent analysis can be performed as usual (the binary influence function would be a special case where one constructs $\psi_i = \phi_i^{(1)} - \phi_i^{(0)} - \theta$ and takes the mean and variance for point estimate and standard error.

---

Nuances:

• ATT isn't a well defined concept for multiple-treatments (or at the very least one has to stipulate the subpopulation that is allowed to contribute to the estimator of $E[Y^({0}) \mid W = 1]$, typically the 'pure control' (0-valued by convention) subpopulation might be reasonable)

Propose a possible solution or implementation

This approach is implemented in Ed Kennedy's package, and I have a minimal implementation in this library. Wondering if there's interest in accommodating this use case (since this will involve changing IIRM from check_data onwards.

Did you consider alternatives to the proposed solution. If yes, please describe

No response

Comments, context or references

No response