if/else when new_expr_vec = TRUE

[sebdalgarno]

I'm trying to work out how to use if/else or ifelse() when new_expr vectorized.

for example, a reprex:

set.seed(123)

nObs <- 500  
ncategory <- 4  

true_b0 <- 0.5  
true_bCategory <- c(-0.8, 1.2, 0.6)  
true_sY <- 0.1

category <- sample(1:ncategory, nObs, replace = TRUE)

y <- numeric(nObs)
for (i in 1:nObs) {
  if (category[i] == 1) {
    eY <- true_b0  # Category 1 is reference
  } else {
    eY <- true_b0 + true_bCategory[category[i] - 1]
  }
  
  y[i] <- rnorm(1, mean = eY, sd = true_sY)
}

data <- data.frame(
  category = factor(category),
  y = y
)

model <- model(code = "data {
  int<lower=1> nObs;
  int<lower=1> ncategory;
  int<lower=1, upper=ncategory> category[nObs];
  real y[nObs];
}
parameters {
  real b0;
  vector[ncategory-1] bCategory;
  real<lower=0> sY;
}
model {
  b0 ~ normal(0, 1);
  bCategory ~ normal(0, 1);
  sY ~ exponential(1);
  
  for (i in 1:nObs) {
    real eY = b0 + (category[i] > 1 ? bCategory[category[i]-1] : 0);
    y[i] ~ normal(eY, sY);
  }
}",
      new_expr = {
        for(i in 1:length(y)) {
          b1[i] <- ifelse(category[i] > 1, bCategory[category[i]-1], 0)
          eY[i] <- b0 + b1[i]
        }
      },
      new_expr_vec =TRUE
)

analysis <- analyse(model, data = data, nthin = 1)

pred <- 
  xnew_data(data, category) %>% 
  predict(analysis, new_data = ., term = "eY")

this does not give the expected answer. not sure if a problem - probably need to dig into the vectorized new_expr code and predict() more

[joethorley]

Ohh - its going to be really hard to vectorize. I would suggest creating a new indexing syntax perhaps something like [[i]] that when in a new expr returns 0 if i is 1 otherwise it subtracts - 1 for i and returns the value of the vector.

[aylapear]

I think we should add in a test in the new_expr_vec() to error if using the base R ifelse() function until we can find a more elegant solution for this. We noted the tidyverse if_else() throws an error due to the recycling that is occurring so no need to build in a test for that as it is already error.

[sebdalgarno]

@nehill197 's interim solution if we don't end up implementing joe's more general solution. define new vector eCat

analysis$model <- update_model(
  analysis$model,
  new_expr = {
  for(i in 1:length(y)) {
    eCat[2:ncategory] <- bCategory
    eCat[1] <- 0
    eY[i] <- b0 + eCat[i]
  }
},
  new_expr_vec = TRUE
)

[joethorley]

very cool - and yes the more general solution could really just have this solution under the hood