Testing for heterogeneity

I read about the following "simple test" for heterogeneity in Athey & Wager 2019 (Estimating Treatment Effects with Causal Forests: An Application). : "A first, simple approach to testing for heterogeneity involves grouping observations according to whether their out-of-bag CATE estimates are above or below the median CATE estimate, and then estimating average treatment effects in these two subgroups separately using the doubly robust approach. This procedure is somewhat heuristic, as the "high" and "low" subgroups are not independent of the scores used to estimate the within-group effects; however, the subgroup definition does not directly depend on the outcomes or treatments (Yi;Wi) themselves, and it appears that this approach can provide at least qualitative insights about the strength of heterogeneity."

In econml, do you think this code snippet correctly implements this heuristic test?

ate_overall = econml_model.ate(X, T0=T0, T1=T)
ate_interval = econml_model.ate_interval(X, T0=T0, T1=T)
effects = econml_model.effect(X=X, T0=T0, T1=T)
median_effect = np.median(effects)
above_median = effect > median_effect
below_median = ~above_median
ate_above = econml_model.ate(
    X.loc[above_median, :],
    T0=T0.loc[above_median, :],
    T1=T.loc[above_median, :],
)
ate_below = econml_model.ate(
    X.loc[below_median, :],
    T0=T0.loc[below_median, :],
    T1=T.loc[below_median, :],
)

Thanks for your help!

Start with the econml_model.effect and ate calls shown in the issue, then consult the cited Athey & Wager 2019 heuristic. Determine whether the subgrouping and treatment-effect estimates match that procedure; done means a documented yes/no answer with any needed clarification or correction.