From 03d2dc34c49afe9876fc5702c46a00a3a8b47600 Mon Sep 17 00:00:00 2001
From: Joe Roe <joeroe@hey.com>
Date: Sun, 3 Nov 2024 08:17:42 +0100
Subject: [PATCH] Add weighted KDE

---
 R/cal_aggregate.R  | 20 +++++++++++++++++---
 man/cal_density.Rd |  3 +++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/R/cal_aggregate.R b/R/cal_aggregate.R
index 0ceb700..65a02b8 100644
--- a/R/cal_aggregate.R
+++ b/R/cal_aggregate.R
@@ -85,10 +85,13 @@ sum.c14_cal <- function(x, range = cal_age_common(x), normalise = FALSE, ...) {
 #' data(shub1_c14)
 #' shub1_cal <- c14_calibrate(shub1_c14$c14_age, shub1_c14$c14_error)
 #' cal_density(shub1_cal)
+#'
+#' # Stratify and weight bootstrap estimation by phase
+#' cal_density(shub1_cal, strata = shub1_c14$phase)
 cal_density <- function(x, bw = 30, ..., times = 25, bootstrap = TRUE, strata = NULL) {
   # TODO: guard against character argument to bw?
-  age_grid <- cal_age_common(x)
 
+  # Bootstrapping
   if (isTRUE(bootstrap)) {
     bootstraps <- cal_bootstraps(x, times = times, strata = strata)
     age_sample <- purrr::map(bootstraps, \(x) do.call(c, cal_sample(x, 1)))
@@ -97,9 +100,20 @@ cal_density <- function(x, bw = 30, ..., times = 25, bootstrap = TRUE, strata =
     age_sample <- cal_sample(x, times)
   }
 
-  kdes <- purrr::map(age_sample, stats::density, bw = bw, from = min(age_grid),
-              to = max(age_grid), ...)
+  # Weights
+  if (is.null(strata)) weights <- NULL
+  else {
+    weights <- purrr::list_c(purrr::map(split(x, strata), function(y, n) {
+     rep(length(y) / n, length(y))
+    }, n = length(x)))
+  }
+
+  # Density estimation
+  age_grid <- cal_age_common(x)
+  kdes <- purrr::map(age_sample, stats::density, bw = bw, weights = weights,
+                     from = min(age_grid), to = max(age_grid), ...)
 
+  # Combine results
   x <- kdes[[1]]$x # TODO: is it always safe to assume x are all equal?
   y <- do.call(rbind, purrr::map(kdes, "y")) # to matrix for faster calculations
   tibble::tibble(
diff --git a/man/cal_density.Rd b/man/cal_density.Rd
index 8a91085..669a2bf 100644
--- a/man/cal_density.Rd
+++ b/man/cal_density.Rd
@@ -55,6 +55,9 @@ See \insertCite{McLaughlin2019;textual}{c14} and
 data(shub1_c14)
 shub1_cal <- c14_calibrate(shub1_c14$c14_age, shub1_c14$c14_error)
 cal_density(shub1_cal)
+
+# Stratify and weight bootstrap estimation by phase
+cal_density(shub1_cal, strata = shub1_c14$phase)
 }
 \references{
 \insertAllCited{}