Vectorize log score gradient calculation

R/loo_model_weights.R

@@ -275,18 +275,14 @@ stacking_weights <-
       # gradient of the objective function
       stopifnot(length(w) == K - 1)
       w_full <- c(w, 1 - sum(w))
-      grad <- rep(0, K - 1)
       mixture_lpd <- matrixStats::rowLogSumExps(
         sweep(lpd_point, 2, log(w_full), "+")
       )
-      for (k in 1:(K - 1)) {
-        grad[k] <- sum(exp_diff_over_exp(
-          lpd_point[, k],
-          lpd_point[, K],
-          mixture_lpd
-        ))
-      }
-      return(-grad)
+      -colSums(exp_diff_over_exp(
+        lpd_point[, seq_len(K - 1), drop = FALSE],
+        matrix(lpd_point[, K], N, K - 1),
+        matrix(mixture_lpd, N, K - 1)
+      ))
     }
 
     ui <- rbind(rep(-1, K - 1), diag(K - 1))  # K-1 simplex constraint matrix

tests/testthat/test_model_weighting.R

@@ -146,6 +146,23 @@ test_that("loo_model_weights (stacking and pseudo-BMA) gives expected result", {
   expect_identical(w3, w3_b)
 })
 
+test_that("stacking_weights gives expected result", {
+  set.seed(10)
+  lpd_point <- matrix(rnorm(2000, -1), 400, 5)
+
+  actual <- stacking_weights(lpd_point)
+
+  expect_s3_class(actual, "stacking_weights")
+  expect_named(actual, paste0("model", 1:5))
+  expect_equal(sum(actual), 1)
+  expect_equal(
+    as.numeric(actual),
+    c(0.158382400055713, 0.176065699214342, 0.19918082992972,
+      0.264490615659226, 0.201880455141),
+    tolerance = 1e-6
+  )
+})
+
 test_that("Bayesian bootstrap gives expected result", {
   lpd_point <- rbind(
     c(-0.2, -0.8, -1.1),