Testing for heterogeneity

[carl-offerfit]

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!