Conditionally update rast values from another raster using terra

I am using the lapp functin of {terra} in R and I want to update rast_a with values from rast_b or rast_c (and some other math) depending on the value in each cell of rast_a.

sample data

rast_a <- rast(ncol = 2, nrow = 2)
values(rast_a) <- 1:4

rast_b <- rast(ncol = 2, nrow = 2)
values(rast_b) <- c(2,2,2,2)

rast_c <- rast(ncol = 2, nrow = 2)
values(rast_c) <- c(3,3,3,3)

Problem

This is my (wrong) attempt.

my_update_formula <- function(a, b, c) {

  a[a == 1] <- b[a == 1] + 10 + 20 - 30
  a[a == 2] <- c[a == 2] + 10 + 50 - 50
  
  return(a)
}

result <- lapp(c(rast_a, rast_b, rast_c),
              fun = my_update_formula)

values(result)

     lyr1
[1,]    3
[2,]    3
[3,]    3
[4,]    4

The actual result should be 2,3,3,4. But because of the operations inside the formula, the first value gets updated twice. First it is changed from 1 to 2 (correctly) but then it fulfills the condition of the second line of code also, and is changed again (I don’t want that to happen).

How can I solve this please?

>Solution :

You can change your formula to

f1 <- function(a, b, c) {
    d <- a
    d[a == 1] <- b[a == 1]
    d[a == 2] <- c[a == 2] + 10
    d
}

#or
f2 <- function(a, b, c) {
    i <- a == 1
    j <- a == 2
    a[i] <- b[i]
    a[j] <- c[j] + 10
    return(a)
}


lapp(c(rast_a, rast_b, rast_c), fun = f1) |> values()
#     lyr1
#[1,]    2
#[2,]   13
#[3,]    3
#[4,]    4

lapp(c(rast_a, rast_b, rast_c), fun = f2) |> values()
#     lyr1
#[1,]    2
#[2,]   13
#[3,]    3
#[4,]    4

You can get the same result with

x <- ifel(rast_a==1, rast_b, 
     ifel(rast_a == 2, rast_c + 10, rast_a))