Fix the inverse affine transformation in loo_moment_match_split()

[avehtari]

This PR was assisted by Claude

Summary

loo_moment_match_split() transforms the first half of the draws with T(θ) = M diag(s) (θ − m) + m + b (upars_trans, with s = total_scaling, M = total_mapping, b = total_shift, m = mean_original). For the second half, the untransformed draws θ, the mixture proposal needs the density at T⁻¹(θ). The code computed upars_trans_inv as diag(s)⁻¹ M⁻¹ (θ − m) + m − b, but T⁻¹(θ) = diag(s)⁻¹ M⁻¹ (θ − m − b) + m. The two differ by (diag(s)⁻¹ M⁻¹ − I) b, so they agree only when the moment matching accepted only shift steps. This PR removes the shift first, then undoes the mapping and the scaling, and adds back only the mean.

One-dimensional check of the old formula:

m <- 0; s <- 2; b <- 1
T <- function(th) s * (th - m) + m + b
Tinv_old <- function(z) (z - m) / s + m - b
Tinv_old(T(1))   # 0.5, not 1

Effect

When a scaling or covariance step was accepted, the weights of the untransformed half used the proposal density at the wrong point. In my tests with exact LOO values (normal models, S = 4000) the effect was small when the LOO posterior is wider than the full posterior (scaling 1.0 to 2.1 accepted): over 60 runs the old and the corrected inverse differed by less than 0.002 on average and by at most 0.08, and both had the same error against the exact value. When the LOO posterior is narrower (an outlier in a model with unknown sd, scaling below 1 accepted), the two differed by 0.15 to 0.34, but moment matching was 4 to 9 off the exact value in those runs for another reason. The iwmm package has the same pattern; there, in brute-force use, the effect is larger (+0.60 on average in a two-dimensional example, topipa/iwmm PR).

Tests

• New test "loo_moment_match_split evaluates the inverse of its transformation": it records the draws that loo_moment_match_split() passes to log_prob_upars and checks that the forward transformation maps the second half back to the original draws, with cov = FALSE and cov = TRUE. It fails on master and passes with this PR.
• The split2 snapshot in "loo_moment_match_split works" changes, because that call gives a scaling of 0.7 and a shift directly; elpd of observation 1 changes from −18.48 to −17.71. That call uses a fixed transformation, not one found by moment matching, so neither value is close to the exact LOO value of that observation (−20.15). I updated the snapshot. No other snapshot changes.
• All 182 tests pass locally with NOT_CRAN=true (R 4.6.1).

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.50%. Comparing base (b22ea49) to head (33961a3).

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #405   +/-   ##
=======================================
  Coverage   92.50%   92.50%           
=======================================
  Files          31       31           
  Lines        3107     3107           
=======================================
  Hits         2874     2874           
  Misses        233      233           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 33961a3 is merged into master:

• ✔️loo_function: 1.86s -> 1.86s [-0.48%, +0.34%]
• ✔️loo_matrix: 1.62s -> 1.62s [-0.37%, +0.58%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[jgabry]

Looks good to me, thanks. I reviewed myself and also asked Codex for a review (to intentionally use a different model, since you said Claude assisted). It found several pre-existing issues that aren't directly related to this PR but are related to moment matching in case you want to fix them in the same PR:

Related pre-existing issues:

1. The rebuilt T may not be the transformation moment matching actually found. After any accepted step the loop restarts at the shift step (R/loo_moment_matching.R:280-382), so scaling and covariance steps can interleave (scale, cov, scale, …). The actual composed linear part is then D₂·M₁·D₁. But the accumulators store total_mapping = M₁ and total_scaling = D₂·D₁, and the split code applies M₁·D₂·D₁. Diagonal and full matrices don't commute, so these generally differ. With this PR, T, T⁻¹ and the Jacobian agree with each other, so the mixture estimator is still valid; T is just not exactly the proposal moment matching found. Fixing it means tracking one combined matrix instead of separate scaling and mapping.
2. The Jacobian ignores cov. log(det(total_mapping)) is subtracted even when cov = FALSE, while the mapping itself is only applied when cov is TRUE. In real use the mapping stays the identity when cov = FALSE, so this is harmless. But the existing split2 test passes cov = FALSE with a non-identity mapping, so that snapshot covers a configuration that can't happen in real use.
3. A copy-paste bug in the r_eff calculation. At R/split_moment_matching.R:125 and :128 both halves use take <- seq_len(S)[-seq_len(S_half)], so r_eff_i1 and r_eff_i2 are computed from the same (second) half. The first should be seq_len(S_half).
4. Docs typo: @PARAM total_mapping says "A vector", but it's a matrix.

[jgabry]

This depends on seen[[2]] being the inverse call, which I think is fine.

Optionally, maybe also worth adding these extra checks? But I'm ok to merge without this.

expect_length(seen, 2)
first <- seq_len(as.integer(0.5 * S))
# first call: second half untransformed; second call: first half untransformed
expect_equal(seen[[1]][second, , drop = FALSE], upars[second, , drop = FALSE])
expect_equal(seen[[2]][first, , drop = FALSE], upars[first, , drop = FALSE])