R/nimble_model.R

In vlandau/tempo: Temporally-Explicit Models for Phenological Observations

# TODO add argument names throughout, make sure variable names correspond to the math.
nimble_model <- function(beta_random_indices, correlated) {
   # Everything returned as strings to avoid undefined variable
   # notes in package check
   if (is.null(beta_random_indices)) {
      "nimbleCode({
         # Likelihood
         for (n in 1:N) {
            prob[n, 1:time_steps] <- ilogit(X[n, , ] %*% Beta[])
            censored[n] ~ dinterval(z[n], y[n, 1:2])
            z[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])
            
            y_rep[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])
         }
         
         z_mean <- mean(z[1:N])
         yrep_mean <- mean(y_rep[1:N])
         p_mean <- step(yrep_mean - z_mean) # pp check for mean

         z_var <- var(z[1:N])
         yrep_var <- var(y_rep[1:N])
         p_var <- step(yrep_var - z_var) # pp check for variances
         
         # Priors
         for (b in 1:n_beta) {
            Beta[b] ~ dnorm(mu_beta[b], sd = sigma_beta[b])
         }
      })"
   } else if (length(beta_random_indices) == 1 | !correlated) {
      "nimbleCode({
         # Likelihood
         for (n in 1:N) {
         # group_ids is vector of group indices
            prob[n, 1:time_steps] <-
              ilogit(X[n, , ] %*% Beta[1:n_beta, group_ids[n]])

            censored[n] ~ dinterval(z[n], y[n, 1:2])
            z[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])
            
            y_rep[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])

         }

         ### group-level betas ###
         for (group_j in 1:n_group) {
            # Random betas
            if (n_random_betas == 1) {
               eps[1, group_j] ~ dnorm(0, sd = sd_eps)
               Beta[beta_random_indices, group_j] <-
                 mean_beta[beta_random_indices] + eps[1, group_j]
            } else {
               for (rb in 1:n_random_betas) {
                  eps[rb, group_j] ~ dnorm(0, sd = sd_eps[rb])
                  Beta[beta_random_indices[rb], group_j] <-
                    mean_beta[beta_random_indices[rb]] + eps[rb, group_j]
               }
            }

            # Fixed betas
            if (n_fixed_betas != 0) {
               if (n_fixed_betas == 1) {
                  Beta[beta_fixed_indices, group_j] <-
                    mean_beta[beta_fixed_indices]
               } else {
                  for (fb in 1:n_fixed_betas) {
                     Beta[beta_fixed_indices[fb], group_j] <-
                       mean_beta[beta_fixed_indices[fb]]
                  }
               }
            }
            
         }

         z_mean <- mean(z[1:N])
         yrep_mean <- mean(y_rep[1:N])
         p_mean <- step(yrep_mean - z_mean)

         z_var <- var(z[1:N])
         yrep_var <- var(y_rep[1:N])
         p_var <- step(yrep_var - z_var)
         
         ### Priors ###
         # means
         for (p in 1:n_beta) {
            mean_beta[p] ~ dnorm(mu_beta[p], sd = sigma_beta[p])
         }

         # Random effect variance
         if (n_random_betas == 1) {
            sd_eps ~ dinvgamma(0.001, 0.001)
         } else {
            for (s in 1:n_random_betas) {
               sd_eps[s] ~ dinvgamma(0.001, 0.001)
            }
         }
      })"
   } else {
      "nimbleCode({
         ### Likelihood ###
         for (n in 1:N) {
            # group_ids is vector of group indices
            prob[n, 1:time_steps] <-
              ilogit(X[n, , ] %*% Beta[1:n_beta, group_ids[n]])

            censored[n] ~ dinterval(z[n], y[n, 1:2])
            z[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])
            
            y_rep[n] ~ dtvgeom_nimble(prob[n, 1:time_steps])
         }

         ### correlations and covariance for random effects ###
         # calculate Rho Matrix
         for (r in 1:n_rho) {
            Rho[rho_indices[r, 1], rho_indices[r, 2]] <- rho[r]
            Rho[rho_indices[r, 2], rho_indices[r, 1]] <- rho[r]
         }

         # calculate Sigma
         Sigma[1:n_random_betas, 1:n_random_betas] <-
           diag(sd_eps[1:n_random_betas]) %*% 
           (Rho[1:n_random_betas, 1:n_random_betas] + diag(n_random_betas)) %*%
           diag(sd_eps[1:n_random_betas])

         ### group-level betas ###
         for (group_j in 1:n_group) {
            # Random betas
            eps[1:n_random_betas, group_j] ~ dmnorm(mean_eps[1:n_random_betas],
                                                    cov = Sigma[1:n_random_betas,
                                                                1:n_random_betas])

            for (rb in 1:n_random_betas) {
               Beta[beta_random_indices[rb], group_j] <-
                 mean_beta[beta_random_indices[rb]] + eps[rb, group_j]
            }


            # Fixed betas
            if (n_fixed_betas != 0) {
               if (n_fixed_betas == 1) {
                  Beta[beta_fixed_indices, group_j] <-
                    mean_beta[beta_fixed_indices]
               } else {
                  for (fb in 1:n_fixed_betas) {
                     Beta[beta_fixed_indices[fb], group_j] <-
                       mean_beta[beta_fixed_indices[fb]]
                  }
               }
            }
         }
         
         z_mean <- mean(z[1:N])
         yrep_mean <- mean(y_rep[1:N])
         p_mean <- step(yrep_mean - z_mean)

         z_var <- var(z[1:N])
         yrep_var <- var(y_rep[1:N])
         p_var <- step(yrep_var - z_var)


         ### Priors ###
         # Means
         for (p in 1:n_beta) {
            mean_beta[p] ~ dnorm(mu_beta[p], sd = sigma_beta[p])
         }

         # Random effects
         for (s in 1:n_random_betas) {
            sd_eps[s] ~ dinvgamma(0.001, 0.001)
         }

         # Correlations
         for (w in 1:n_rho) {
            rho[w] ~ dunif(-1, 1)
         }
      })"
   }
}