From 22d60534e52ab50f7e521d8a77aae0d6aa8114c2 Mon Sep 17 00:00:00 2001
From: Ben Goodrich <goodrich.ben@gmail.com>
Date: Wed, 16 Mar 2022 10:52:13 -0400
Subject: [PATCH] more fixes for sparse matrices

---
 src/stan_files/bernoulli.stan                  | 16 ++++++++--------
 src/stan_files/data/NKX.stan                   |  4 ++--
 src/stan_files/functions/common_functions.stan |  2 +-
 tests/testthat/test_stan_functions.R           |  2 +-
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/stan_files/bernoulli.stan b/src/stan_files/bernoulli.stan
index c2371ffe5..cf4b3561c 100644
--- a/src/stan_files/bernoulli.stan
+++ b/src/stan_files/bernoulli.stan
@@ -22,14 +22,14 @@ data {
   // stuff for the sparse case
   int<lower=0> nnz_X0;                       // number of non-zero elements in the implicit X0 matrix
   vector[nnz_X0] w_X0;                       // non-zero elements in the implicit X0 matrix
-  int<lower=0, upper = K - 1> v_X0[nnz_X0];  // column indices for w_X0
+  int<lower=1, upper = K> v_X0[nnz_X0];      // column indices for w_X0
   // where the non-zeros start in each row of X0
-  int<lower=0, upper = rows(w_X0) + 1> u_X0[dense_X ? 0 : N[1] + 1]; 
+  int<lower=1, upper = rows(w_X0) + 1> u_X0[dense_X ? 0 : N[1] + 1]; 
   int<lower=0> nnz_X1;                       // number of non-zero elements in the implicit X1 matrix
   vector[nnz_X1] w_X1;                       // non-zero elements in the implicit X1 matrix
-  int<lower=0, upper = K - 1> v_X1[nnz_X1];  // column indices for w_X1
+  int<lower=1, upper = K> v_X1[nnz_X1];      // column indices for w_X1
   // where the non-zeros start in each row of X1
-  int<lower=0, upper = rows(w_X1) + 1> u_X1[dense_X ? 0 : N[2] + 1]; 
+  int<lower=1, upper = rows(w_X1) + 1> u_X1[dense_X ? 0 : N[2] + 1]; 
   // declares prior_PD, has_intercept, link, prior_dist, prior_dist_for_intercept
 #include /data/data_glm.stan
 
@@ -59,12 +59,12 @@ data {
   int<lower=0> num_non_zero[2];     // number of non-zero elements in the Z matrices
   vector[num_non_zero[1]] w0;       // non-zero elements in the implicit Z0 matrix
   vector[num_non_zero[2]] w1;       // non-zero elements in the implicit Z1 matrix
-  int<lower=0, upper = q - 1> v0[num_non_zero[1]]; // column indices for w0
-  int<lower=0, upper = q - 1> v1[num_non_zero[2]]; // column indices for w1
+  int<lower=1, upper = q> v0[num_non_zero[1]]; // column indices for w0
+  int<lower=1, upper = q> v1[num_non_zero[2]]; // column indices for w1
   // where the non-zeros start in each row of Z0
-  int<lower=0, upper = rows(w0) + 1> u0[t > 0 ? N[1] + 1 : 0];  
+  int<lower=1, upper = rows(w0) + 1> u0[t > 0 ? N[1] + 1 : 0];  
   // where the non-zeros start in each row of Z1
-  int<lower=0, upper = rows(w1) + 1> u1[t > 0 ? N[2] + 1 : 0];  
+  int<lower=1, upper = rows(w1) + 1> u1[t > 0 ? N[2] + 1 : 0];  
   int<lower=0, upper=1> special_case;     // whether we only have to deal with (1|group)
 }
 transformed data {
diff --git a/src/stan_files/data/NKX.stan b/src/stan_files/data/NKX.stan
index 6437c7383..88dde9bff 100644
--- a/src/stan_files/data/NKX.stan
+++ b/src/stan_files/data/NKX.stan
@@ -10,9 +10,9 @@
   // stuff for the sparse case
   int<lower=0> nnz_X;                      // number of non-zero elements in the implicit X matrix
   vector[nnz_X] w_X;                       // non-zero elements in the implicit X matrix
-  int<lower=0, upper = K - 1> v_X[nnz_X];  // column indices for w_X
+  int<lower=1, upper = K> v_X[nnz_X];  // column indices for w_X
   // where the non-zeros start in each row of X
-  int<lower=0, upper = rows(w_X) + 1> u_X[dense_X ? 0 : N + 1]; 
+  int<lower=1, upper = rows(w_X) + 1> u_X[dense_X ? 0 : N + 1]; 
   
   // smooths
   int<lower=0> K_smooth;
diff --git a/src/stan_files/functions/common_functions.stan b/src/stan_files/functions/common_functions.stan
index 25ada737b..63c9741e1 100644
--- a/src/stan_files/functions/common_functions.stan
+++ b/src/stan_files/functions/common_functions.stan
@@ -241,7 +241,7 @@
     int V[t,N];
     int pos = 1;
     if (t > 0) for (j in 1:N) for (i in 1:t) {
-      V[i,j] = v[pos] + 1;
+      V[i,j] = v[pos]; // + 1
       pos += 1;
     }
     return V;
diff --git a/tests/testthat/test_stan_functions.R b/tests/testthat/test_stan_functions.R
index 1389eba69..c41331d97 100644
--- a/tests/testthat/test_stan_functions.R
+++ b/tests/testthat/test_stan_functions.R
@@ -402,7 +402,7 @@ if (require(lme4) && require(HSAUR3)) test_that("the Stan equivalent of lme4's Z
       })) ) {
       V <- matrix(parts$v, nrow = sum(p), ncol = nrow(Z))
       expect_true(all(V == 
-                      t(as.matrix(as.data.frame(make_V(nrow(Z), nrow(V), parts$v - 1L))))))
+                      t(as.matrix(as.data.frame(make_V(nrow(Z), nrow(V), parts$v))))))
       expect_equal(Zb@x, apply(V, 2, FUN = function(v) sum(b[v])))
     }
   }