Help merging the stancode_opt branch

DESCRIPTION

@@ -16,7 +16,7 @@ URL: https://github.com/ScottClaessens/coevolve, https://scottclaessens.github.i
 BugReports: https://github.com/ScottClaessens/coevolve/issues
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE, roclets = c("namespace", "rd", "srr::srr_stats_roclet"))
-RoxygenNote: 7.3.2
+RoxygenNote: 7.3.3
 Suggests: 
     knitr,
     rmarkdown,
@@ -48,7 +48,7 @@ Imports:
     tidyr
 VignetteBuilder: knitr
 Depends: 
-    R (>= 3.5.0)
+    R (>= 4.1.0)
 LazyData: true
 Additional_repositories: https://stan-dev.r-universe.dev/
 Remotes:

R/coev_make_stancode.R

@@ -659,7 +659,9 @@ write_transformed_pars_block <- function(data, distributions, id, dist_mat,
       "  matrix[J,J] Q = diag_matrix(Q_sigma^2); // drift matrix\n"
     ),
     "  matrix[J,J] Q_inf; // asymptotic covariance matrix\n",
-    "  array[N_tree, N_seg] matrix[J,J] VCV_tips; // vcov matrix for drift\n"
+    "  array[N_tree, N_seg] matrix[J,J] VCV_tips; // vcov matrix for drift\n",
+    "  array[N_tree, N_seg] matrix[J,J] L_VCV_tips; ",
+    "// Cholesky factor of VCV_tips\n"
   )
   # add distance random effects if distance matrix specified by user
   if (!is.null(dist_mat)) {
@@ -708,7 +710,6 @@ write_transformed_pars_block <- function(data, distributions, id, dist_mat,
     "  }\n",
     "  // calculate asymptotic covariance\n",
     "  Q_inf = ksolve(A, Q);\n",
-    "  array[N_unique_lengths] matrix[J,J] L_VCV_tips_cache;\n",
     "  {\n",
     "    array[N_unique_lengths] matrix[J,J] A_delta_cache;\n",
     "    array[N_unique_lengths] matrix[J,J] VCV_cache;\n",
@@ -728,11 +729,11 @@ write_transformed_pars_block <- function(data, distributions, id, dist_mat,
     "        }\n",
     "      }\n",
     "    }\n",
-    "    L_VCV_tips_cache = L_VCV_cache;\n",
     "    for (t in 1:N_tree) {\n",
     "      // setting ancestral states and placeholders\n",
     "      eta[t, node_seq[t, 1]] = eta_anc[t];\n",
     "      VCV_tips[t, node_seq[t, 1]] = diag_matrix(rep_vector(-99, J));\n",
+    "      L_VCV_tips[t, node_seq[t, 1]] = diag_matrix(rep_vector(1.0, J));\n",
     "      for (i in 2:N_seg) {\n",
     "        matrix[J,J] A_delta;\n",
     "        matrix[J,J] VCV;\n",
@@ -756,15 +757,18 @@ write_transformed_pars_block <- function(data, distributions, id, dist_mat,
     "          eta[t, node_seq[t, i]] = to_vector(\n",
     "            A_delta * eta[t, parent[t, i]] + (A_solve * b) + drift_seg\n",
     "          );\n",
-    "          VCV_tips[t, node_seq[t, i]] = diag_matrix(rep_vector(-99, J));",
-    "\n",
+    "          VCV_tips[t, node_seq[t, i]] = ",
+    "diag_matrix(rep_vector(-99, J));\n",
+    "          L_VCV_tips[t, node_seq[t, i]] = ",
+    "diag_matrix(rep_vector(1.0, J));\n",
     "        }\n",
     "        // if is a tip, omit, we'll deal with it in the model block;\n",
     "        else {\n",
     "          eta[t, node_seq[t, i]] = to_vector(\n",
     "            A_delta * eta[t, parent[t, i]] + (A_solve * b)\n",
     "          );\n",
     "          VCV_tips[t, node_seq[t, i]] = VCV;\n",
+    "          L_VCV_tips[t, node_seq[t, i]] = L_VCV;\n",
     "        }\n",
     "      }\n",
     "    }\n",
@@ -777,19 +781,8 @@ write_transformed_pars_block <- function(data, distributions, id, dist_mat,
       sc_transformed_parameters,
       "  for (t in 1:N_tree) {\n",
       "    for (i in 1:N_tips) {\n",
-      paste0(
-        "      if (tip_to_seg[t, i] > 0 && ",
-        "length_index[t, tip_to_seg[t, i]] > 0) {\n"
-      ),
-      paste0(
-        "        tdrift[t,i] = L_VCV_tips_cache[",
-        "length_index[t, tip_to_seg[t, i]]] * ",
-        "to_vector(terminal_drift[t][i,]);\n"
-      ),
-      "      } else {\n",
-      "        tdrift[t,i] = cholesky_decompose(VCV_tips[t,i]) * ",
+      "      tdrift[t,i] = L_VCV_tips[t, i] * ",
       "to_vector(terminal_drift[t][i,]);\n",
-      "      }\n",
       "    }\n",
       "  }\n"
     )
@@ -1044,14 +1037,18 @@ write_model_block <- function(data, distributions, id, dist_mat, priors,
       sc_model <- paste0(
         sc_model,
         "        lp[t] = multi_normal_cholesky_lpdf(tdrift | rep_vector(0.0, ",
-        "J), cholesky_decompose(",
+        "J), ",
         ifelse(
           !is.null(measurement_error),
           # add squared standard errors to diagonal of VCV_tips
-          "add_diag(VCV_tips[t, tip_id[i]], se[i,])",
-          "VCV_tips[t, tip_id[i]]"
+          paste0(
+            "cholesky_decompose(",
+            "add_diag(VCV_tips[t, tip_id[i]], se[i,])",
+            ")"
+          ),
+          "L_VCV_tips[t, tip_id[i]]"
         ),
-        "));\n"
+        ");\n"
       )
     }
   }

tests/testthat/test-coev_make_stancode_indexing.R

@@ -0,0 +1,176 @@
+# Test suite for verifying correct array indexing in generated Stan code
+#
+# This test suite catches bugs where segment indices are confused with node
+# indices. Key insight: VCV_tips and L_VCV_tips are indexed by NODE NUMBER
+# (1 to N_seg), not by segment traversal order.
+#
+# The tip_to_seg mapping converts tip node numbers to segment indices,
+# so using tip_to_seg to index VCV_tips would access the WRONG elements.
+
+test_that("VCV_tips and L_VCV_tips do not use tip_to_seg for access", {
+  # Generate Stan code for a simple model
+
+  sc <- coev_make_stancode(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # VCV_tips should NOT be accessed via tip_to_seg
+
+  # (tip_to_seg returns segment indices, but VCV_tips is indexed by node number)
+  expect_false(
+    grepl("VCV_tips\\[t,\\s*tip_to_seg", sc),
+    info = "VCV_tips should not use tip_to_seg indexing"
+  )
+  expect_false(
+    grepl("L_VCV_tips\\[t,\\s*tip_to_seg", sc),
+    info = "L_VCV_tips should not use tip_to_seg indexing"
+  )
+})
+
+test_that("VCV_tips storage uses node_seq during tree traversal", {
+  sc <- coev_make_stancode(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # When storing VCV_tips during tree traversal, should use node_seq
+  expect_true(
+    grepl("VCV_tips\\[t, node_seq\\[t,", sc),
+    info = "VCV_tips storage should use node_seq indexing"
+  )
+  expect_true(
+    grepl("L_VCV_tips\\[t, node_seq\\[t,", sc),
+    info = "L_VCV_tips storage should use node_seq indexing"
+  )
+})
+
+test_that("segment-indexed arrays use correct traversal indexing", {
+  sc <- coev_make_stancode(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # z_drift should be accessed by segment index (i-1 in the loop)
+  expect_true(
+    grepl("z_drift\\[t, i-1\\]", sc),
+    info = "z_drift should use segment indexing"
+  )
+
+  # length_index should be accessed by segment index
+  expect_true(
+    grepl("length_index\\[t, i\\]", sc),
+    info = "length_index should use segment indexing"
+  )
+})
+
+test_that("tip_id maps observations to tips correctly in model block", {
+  sc <- coev_make_stancode(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # eta should be accessed via tip_id for observations
+  expect_true(
+    grepl("eta\\[t,\\s*tip_id\\[i\\]\\]", sc),
+    info = "eta should use tip_id to access tip states for observations"
+  )
+})
+
+test_that("measurement error model uses correct VCV_tips indexing", {
+  skip_on_cran()
+
+  # Create test data with measurement error
+  d <- authority$data
+  d$x <- rnorm(nrow(d))
+  d$y <- rnorm(nrow(d))
+  d$x_se <- rexp(nrow(d), 5)
+  d$y_se <- rexp(nrow(d), 5)
+
+  sc <- coev_make_stancode(
+    data = d,
+    variables = list(x = "normal", y = "normal"),
+    id = "language",
+    tree = authority$phylogeny,
+    measurement_error = list(x = "x_se", y = "y_se")
+  )
+
+  # With measurement error, VCV_tips access should still not use tip_to_seg
+  expect_false(
+    grepl("VCV_tips\\[t,\\s*tip_to_seg", sc),
+    info = "Measurement error model should not use tip_to_seg for VCV_tips"
+  )
+
+  # Should use tip_id for observation-level access in model block
+  expect_true(
+    grepl("VCV_tips\\[t,\\s*tip_id\\[i\\]\\]", sc),
+    info = "Measurement error model should use tip_id for VCV_tips access"
+  )
+})
+
+test_that("normal distribution model uses correct L_VCV_tips indexing", {
+  skip_on_cran()
+
+  # Create test data with normal variables (no measurement error)
+  d <- authority$data
+  d$x <- rnorm(nrow(d))
+  d$y <- rnorm(nrow(d))
+
+  sc <- coev_make_stancode(
+    data = d,
+    variables = list(x = "normal", y = "normal"),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # L_VCV_tips access should not use tip_to_seg
+
+  expect_false(
+    grepl("L_VCV_tips\\[t,\\s*tip_to_seg", sc),
+    info = "Normal model should not use tip_to_seg for L_VCV_tips"
+  )
+
+  # Should use tip_id for observation-level access
+  expect_true(
+    grepl("L_VCV_tips\\[t,\\s*tip_id\\[i\\]\\]", sc),
+    info = "Normal model should use tip_id for L_VCV_tips access"
+  )
+})
+
+test_that("generated quantities uses correct indexing for yrep", {
+  sc <- coev_make_stancode(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny
+  )
+
+  # In generated quantities, tip_id should be used to link observations to tips
+  expect_true(
+    grepl("terminal_drift_rep\\[t,\\s*tip_id\\[i\\]\\]", sc),
+    info = "Generated quantities should use tip_id for terminal_drift_rep"
+  )
+})