add weight aesthetic

Author: teunbrand

R/stat-ellipse.R

@@ -76,6 +76,8 @@ stat_ellipse <- function(mapping = NULL, data = NULL,
 #' @export
 StatEllipse <- ggproto("StatEllipse", Stat,
   required_aes = c("x", "y"),
+  optional_aes = "weight",
+  dropped_aes = "weight",
 
   setup_params = function(data, params) {
     params$type <- params$type %||% "t"
@@ -96,6 +98,9 @@ calculate_ellipse <- function(data, vars, type, level, segments){
   dfn <- 2
   dfd <- nrow(data) - 1
 
+  weight <- data$weight %||% rep(1, nrow(data))
+  weight <- weight / sum(weight)
+
   if (!type %in% c("t", "norm", "euclid")) {
     cli::cli_inform("Unrecognized ellipse type")
     ellipse <- matrix(NA_real_, ncol = 2)
@@ -104,11 +109,12 @@ calculate_ellipse <- function(data, vars, type, level, segments){
     ellipse <- matrix(NA_real_, ncol = 2)
   } else {
     if (type == "t") {
-      v <- MASS::cov.trob(data[,vars])
+      # Prone to convergence problems when `sum(weight) != nrow(data)`
+      v <- MASS::cov.trob(data[,vars], wt = weight * nrow(data))
     } else if (type == "norm") {
-      v <- stats::cov.wt(data[,vars])
+      v <- stats::cov.wt(data[,vars], wt = weight)
     } else if (type == "euclid") {
-      v <- stats::cov.wt(data[,vars])
+      v <- stats::cov.wt(data[,vars], wt = weight)
       v$cov <- diag(rep(min(diag(v$cov)), 2))
     }
     shape <- v$cov