Fix the inverse affine transformation in loo_moment_match_split()

R/loo_moment_matching.R

@@ -268,9 +268,11 @@ loo_moment_match_i <- function(i,
   lwfi <- rep(-matrixStats::logSumExp(rep(0, S)),S)
 
   # initialize objects that keep track of the total transformation
-  total_shift <- rep(0, npars)
-  total_scaling <- rep(1, npars)
-  total_mapping <- diag(npars)
+  total <- list(
+    shift = rep(0, npars),
+    scaling = rep(1, npars),
+    mapping = diag(npars)
+  )
 
   # try several transformations one by one
   # if one does not work, do not apply it and try another one
@@ -302,7 +304,7 @@ loo_moment_match_i <- function(i,
     }
     if (quantities_i$ki < ki) {
       uparsi <- trans$upars
-      total_shift <- total_shift + trans$shift
+      total <- add_step_to_total(total, trans, cov)
 
       lwi <- quantities_i$lwi
       lwfi <- quantities_i$lwfi
@@ -332,8 +334,7 @@ loo_moment_match_i <- function(i,
     }
     if (quantities_i$ki < ki) {
       uparsi <- trans$upars
-      total_shift <- total_shift + trans$shift
-      total_scaling <- total_scaling * trans$scaling
+      total <- add_step_to_total(total, trans, cov)
 
       lwi <- quantities_i$lwi
       lwfi <- quantities_i$lwfi
@@ -365,8 +366,7 @@ loo_moment_match_i <- function(i,
       }
       if (quantities_i$ki < ki) {
         uparsi <- trans$upars
-        total_shift <- total_shift + trans$shift
-        total_mapping <- trans$mapping %*% total_mapping
+        total <- add_step_to_total(total, trans, cov)
 
         lwi <- quantities_i$lwi
         lwfi <- quantities_i$lwfi
@@ -389,7 +389,7 @@ loo_moment_match_i <- function(i,
   if (split && (iterind > 1)) {
     # compute split transformation
     split_obj <- loo_moment_match_split(
-      x, upars, cov, total_shift, total_scaling, total_mapping, i,
+      x, upars, cov, total$shift, total$scaling, total$mapping, i,
       log_prob_upars = log_prob_upars, log_lik_i_upars = log_lik_i_upars,
       cores = 1, r_eff_i = r_eff_i, is_method = is_method, ...
     )
@@ -486,6 +486,39 @@ update_quantities_i <- function(x, upars, i, orig_log_prob,
 
 
 
+#' Add one accepted moment matching step to the total transformation
+#'
+#' Each step maps the current draws `u` to `L (u - mean(u)) + mean(u) + shift`,
+#' where `L` is the identity, `diag(scaling)` or `mapping`. The shifts add up,
+#' and the linear parts compose in the order in which the steps were accepted.
+#' After any accepted step the loop in `loo_moment_match_i()` starts again with
+#' the shift step, so a scaling step can follow a covariance step. With
+#' `cov = TRUE` the scaling is therefore applied to the rows of `mapping`
+#' (`diag(scaling) %*% mapping`), and `scaling` stays one. With `cov = FALSE`
+#' there is no covariance step, and the scalings are kept in `scaling`.
+#'
+#' @noRd
+#' @param total A list with `shift`, `scaling` and `mapping`.
+#' @param trans The output of `shift()`, `shift_and_scale()` or
+#'   `shift_and_cov()`.
+#' @param cov Logical; the `cov` argument of the moment matching.
+#' @return The updated `total`.
+#'
+add_step_to_total <- function(total, trans, cov) {
+  total$shift <- total$shift + trans$shift
+  if (!is.null(trans$scaling)) {
+    if (cov) {
+      total$mapping <- trans$scaling * total$mapping
+    } else {
+      total$scaling <- total$scaling * trans$scaling
+    }
+  }
+  if (!is.null(trans$mapping)) {
+    total$mapping <- trans$mapping %*% total$mapping
+  }
+  total
+}
+
 #' Shift a matrix of parameters to their weighted mean.
 #' Also calls update_quantities_i which updates the importance weights based on
 #' the supplied model object.

R/split_moment_matching.R

@@ -13,9 +13,15 @@
 #' @param total_shift A vector representing the total shift made by the moment
 #'   matching algorithm.
 #' @param total_scaling A vector representing the total scaling of marginal
-#'   variance made by the moment matching algorithm.
-#' @param total_mapping A vector representing the total covariance
-#'   transformation made by the moment matching algorithm.
+#'   variance made by the moment matching algorithm. With `cov = TRUE`,
+#'   [loo_moment_match()] includes the scaling in `total_mapping` instead, and
+#'   this vector is all ones.
+#' @param total_mapping A square matrix representing the total linear
+#'   transformation made by the moment matching algorithm: the product of the
+#'   linear parts of the accepted steps, in the order in which they were
+#'   accepted. Used only when `cov = TRUE`. The transformation is
+#'   `T(theta) = total_mapping %*% diag(total_scaling) %*% (theta - m) + m +
+#'   total_shift`, where `m` is the mean of `upars`.
 #' @param i Observation index.
 #' @param log_prob_upars A function that takes arguments `x` and `upars` and
 #'   returns a matrix of log-posterior density values of the unconstrained
@@ -56,12 +62,15 @@ loo_moment_match_split <- function(x, upars, cov, total_shift, total_scaling,
   attributes(upars_trans) <- attributes(upars)
 
   # inverse accumulated affine transformation
-  upars_trans_inv <- sweep(upars, 2, mean_original, "-")
+  # T(theta) = M diag(s) (theta - m) + m + shift, so
+  # T^-1(z) = diag(s)^-1 M^-1 (z - m - shift) + m:
+  # remove the shift first, then undo the mapping and the scaling.
+  upars_trans_inv <- sweep(upars, 2, mean_original + total_shift, "-")
   if (cov) {
     upars_trans_inv <- tcrossprod(upars_trans_inv, solve(total_mapping))
   }
   upars_trans_inv <- sweep(upars_trans_inv, 2, total_scaling, "/")
-  upars_trans_inv <- sweep(upars_trans_inv, 2, mean_original - total_shift, "+")
+  upars_trans_inv <- sweep(upars_trans_inv, 2, mean_original, "+")
   attributes(upars_trans_inv) <- attributes(upars)
 
   # first half of upars_trans_half are T(theta)
@@ -82,9 +91,14 @@ loo_moment_match_split <- function(x, upars, cov, total_shift, total_scaling,
   log_liki_half <- log_lik_i_upars(x, upars = upars_trans_half, i = i, ...)
 
   # compute weights
-  log_prob_half_trans_inv <- (log_prob_half_trans_inv -
-                              log(prod(total_scaling)) -
-                              log(det(total_mapping)))
+  # log of the absolute Jacobian determinant of T; the mapping is part of T
+  # only when cov = TRUE
+  log_det_jacobian <- sum(log(total_scaling))
+  if (cov) {
+    log_det_jacobian <- log_det_jacobian +
+      as.numeric(determinant(total_mapping, logarithm = TRUE)$modulus)
+  }
+  log_prob_half_trans_inv <- log_prob_half_trans_inv - log_det_jacobian
   stable_S <- log_prob_half_trans > log_prob_half_trans_inv
 
   lwi_half <- -log_liki_half + log_prob_half_trans
@@ -122,7 +136,7 @@ loo_moment_match_split <- function(x, upars, cov, total_shift, total_scaling,
   # currently ignores chain information
   # since we have two proposal distributions
   # compute S_eff separately from both and take the smaller
-  take <- seq_len(S)[-seq_len(S_half)]
+  take <- seq_len(S_half)
   log_liki_half_1 <- log_liki_half[take, drop = FALSE]
   dim(log_liki_half_1) <- c(length(take), 1, 1)
   take <- seq_len(S)[-seq_len(S_half)]