Treatment Effect Heterogeneity Visualization — heterogeneity

Creates publication-ready visualizations for treatment effect heterogeneity (HTE) across subgroups or continuous moderators. Supports forest plots, scatter plots of CATE vs. covariates, and marginal effects plots.

Usage

heterogeneity_plot(
  data,
  cate_var = NULL,
  subgroup_var = NULL,
  moderator_var = NULL,
  outcome = NULL,
  treatment = NULL,
  weights = NULL,
  overall_effect = NULL,
  conf_level = 0.95,
  plot_type = c("forest", "scatter", "violin"),
  title = "Treatment Effect Heterogeneity"
)

Arguments

data: A data frame.
cate_var: Character. Name of the column containing estimated CATEs. Required for CATE-based plots.
subgroup_var: Character. Name of a discrete subgroup variable. If provided, creates a forest-style subgroup HTE plot.
moderator_var: Character. Name of a continuous moderator variable. If provided, creates a scatter/LOESS plot of CATE vs. moderator.
outcome: Character. Outcome variable name (for computing raw subgroup effects when cate_var is not available).
treatment: Character. Treatment variable name.
weights: Numeric vector. Optional analytical weights.
overall_effect: Numeric or NULL. Overall ATE to show as reference line. If NULL, estimated from data.
conf_level: Numeric. Confidence level. Default 0.95.
plot_type: Character. "forest" (default for subgroups), "scatter" (for continuous moderator), or "violin".
title: Character. Plot title.

Value

A ggplot2 object.

Examples

if (FALSE) { # \dontrun{
library(grf)
n <- 500; p <- 5
X <- data.frame(matrix(rnorm(n * p), n, p))
W <- rbinom(n, 1, 0.5)
Y <- pmax(X[, 1], 0) * W + rnorm(n)
cf <- causal_forest(as.matrix(X), Y, W)
df <- cbind(X, cate = predict(cf)$predictions)

# Scatter: CATE vs covariate
heterogeneity_plot(df, cate_var = "cate", moderator_var = "X1",
                   plot_type = "scatter")

# Subgroup forest plot
df$subgroup <- ifelse(X[, 1] > 0, "High X1", "Low X1")
heterogeneity_plot(df, cate_var = "cate", subgroup_var = "subgroup")
} # }