Improve identifiability of non-stationary GP

NEWS.md

@@ -17,6 +17,10 @@
 - Added a model overview vignette with an architecture diagram showing how the package's models connect.
 - Added a model features vignette providing a quick reference to all modelling options with links to detailed documentation.
 
+## Model changes
+
+- The cumulated non-stationary Gaussian process used to model Rt over time is now mean-centred (`gp -= mean(gp)` in `inst/stan/functions/rt.stan`), so `R0` represents the mean Rt over the trajectory rather than the initial value. This eliminates the `(R0, drift)` ridge in the joint posterior that was responsible for stuck chains and catastrophic R-hat values on some seeds. No API change and no change to the `alpha` prior — verified across previously stuck seeds (R-hat goes from up to 6.10 down to <1.01, treedepth hits from hundreds down to zero).
+
 ## Bug fixes
 
 - Fixed a bug in `forecast_infections()` where the summary call to extract dates was using modified args instead of the original fit dimensions, causing a date-dimension mismatch when extending the R trajectory beyond the original observation period.

inst/stan/functions/rt.stan

@@ -48,6 +48,10 @@ vector update_Rt(int t, real R0, vector noise, array[] int bps,
     } else {
       gp[2:(gp_n + 1)] = noise;
       gp = cumulative_sum(gp);
+      // Identifiability: subtract the trajectory mean so log R0 is the mean
+      // log Rt over the window rather than the initial value. Eliminates
+      // the (R0, drift) ridge in the joint posterior.
+      gp -= mean(gp);
     }
     logR = logR + gp;
   }

tests/testthat/test-stan-rt.R

@@ -9,9 +9,14 @@ test_that("update_Rt works to produce multiple Rt estimates with a static gaussi
   )
 })
 test_that("update_Rt works to produce multiple Rt estimates with a non-static gaussian process", {
+  # Non-stationary GP: cumulated trajectory is centred so log R0 = mean log Rt
+  # over the window rather than the initial value (eliminates the (R0, drift)
+  # ridge in the joint posterior). For noise = rep(0.1, 9), gp_n = 9:
+  #   gp = cumsum(noise) = c(0, 0.1, 0.2, ..., 0.9), mean = 0.45,
+  #   centred = c(-0.45, -0.35, ..., 0.45). log Rt = log(1.2) + centred.
   expect_equal(
-    round(update_Rt(10, 1.2, rep(0.1, 9), rep(10, 0), numeric(0), 0), 2),
-    c(1.20, 1.33, 1.47, 1.62, 1.79, 1.98, 2.19, 2.42, 2.67, 2.95)
+    round(update_Rt(10, 1.2, rep(0.1, 9), rep(10, 0), numeric(0), 0), 3),
+    c(0.765, 0.846, 0.935, 1.033, 1.141, 1.262, 1.394, 1.541, 1.703, 1.882)
   )
 })
 test_that("update_Rt works to produce multiple Rt estimates with a non-static stationary gaussian process", {
@@ -53,9 +58,12 @@ test_that("update_Rt works when Rt is variable and a breakpoint is present", {
     round(update_Rt(5, 1.2, rep(0, 5), c(1, 1, 2, 2, 2), 0.1, 1), 2),
     c(1.2, 1.2, rep(1.33, 3))
   )
+  # Non-stationary GP: see explanation in the earlier non-static GP test.
+  # Here gp_n = 4, gp_centred = c(-0.2, -0.1, 0, 0.1, 0.2), breakpoint adds
+  # 0.1 from t = 3 onward.
   expect_equal(
-    round(update_Rt(5, 1.2, rep(0.1, 4), c(1, 1, 2, 2, 2), 0.1, 0), 2),
-    c(1.20, 1.33, 1.62, 1.79, 1.98)
+    round(update_Rt(5, 1.2, rep(0.1, 4), c(1, 1, 2, 2, 2), 0.1, 0), 3),
+    c(0.982, 1.086, 1.326, 1.466, 1.620)
   )
 })