Conditional cumulative sum in R with NA and keeping track of episodic wave

I have the following data:

df <- data.frame(time = 0:9,
y_ex1 = c(1,NA,0,1,NA,NA,NA,1,0,0),
y_ex2 = c(1,NA,0,0,NA,NA,NA,1,0,0),
y_ex3 = c(1,NA,0,1,NA,NA,NA,0,0,0),
y_ex4 = c(1,NA,0,0,NA,NA,NA,0,0,0),
y_ex5 = c(1,1,0,0,1,1,1,0,0,0),
)

And I'd like to score each of these vectors as follows:

df$y_ex1_scored <- c(1, 1.5, 0, 1, 2, 3, 4, 5, 0, 0)
df$y_ex2_scored <- c(1, 1.5, 0, 0, .5, 1, 1.5, 2.5, 0, 0)
df$y_ex3_scored <- c(1, 1.5, 0, 1, 1.5, 2, 2.5, 0, 0, 0)
df$y_ex4_scored <- c(1, 1.5, 0, 0, 0, 0, 0, 0, 0, 0)
df$y_ex5_scored <- c(1, 2, 0, 0, 1, 2, 3, 0, 0, 0)

Basically, what I am trying to do is calculate a cumulative sum until I have a 0 or an NA. If I hit a 0, I want the sum to start all over again at 0. If I hit an NA, I want the score to depend on what the value was before and after the NA. If the values are the same, then I want the cumulative sum to continue. For example, if they are both 1, I want the sum to continue as though the data were not missing. If they are both 0, the NAs should be replaced with 0s. If the values are different, I want to add a .5 where each NA.

I also want to keep track of each wave and a wave occurs when the values shift from 0 to 1 or 1 to 0. So for the y_ex1 variable, it should look like

y_ex1_wave <- c(1, 1, 2, 2, 2, 2, 2, 2, 3, 3)

and for y_ex4 it would be

y_ex4_wave <- c(1, 1, 2, 2, 2, 2, 2, 2, 2, 2)

Happy to clarify this, I have been writing a bunch of terrible loops in R to try to do this and have spent far to0 long and have made little progress. Happy to have a base R or tidyverse solution.

Here's a function that'll help:

fun <- function(z) {
  r <- rle(replace(z, is.na(z), Inf))
  for (na in which(is.infinite(r$values))) {
    if (na %in% c(1L, length(r$values))) {
      stop("oops?")
    } else {
      r$values[na] <- ifelse(r$values[na-1] == r$values[na+1], r$values[na-1], 0.5)
    }
  }
  z2 <- inverse.rle(r)
  ave(z2, cumsum(z2 == 0), FUN = cumsum)
}

Demonstration:

out <- as.data.frame(lapply(df[,2:6], fun))
names(out) <- paste0(names(out), "_scored")
all.equal(out, df[,7:11])
# [1] TRUE
out
#    y_ex1_scored y_ex2_scored y_ex3_scored y_ex4_scored y_ex5_scored
# 1           1.0          1.0          1.0          1.0            1
# 2           1.5          1.5          1.5          1.5            2
# 3           0.0          0.0          0.0          0.0            0
# 4           1.0          0.0          1.0          0.0            0
# 5           2.0          0.5          1.5          0.0            1
# 6           3.0          1.0          2.0          0.0            2
# 7           4.0          1.5          2.5          0.0            3
# 8           5.0          2.5          0.0          0.0            0
# 9           0.0          0.0          0.0          0.0            0
# 10          0.0          0.0          0.0          0.0            0

Brief walk-through of the function:

• rle is for run-length encoding, which converts something like c(1,1,2,2,2,3) into c(1,2,3) (values) and c(2,3,1) (lengths of each value). This is useful here so that we can easily keep track of NAs for revaluing.
• Unfortunately, rle (and most of R) treats NA as "it could be anything", so (from ?rle) "Missing values are regarded as unequal to the previous value". Because of this, NA will not be collapsed (as in rows 5-7); however, as a trick, and assuming you have all finite values, I replace NAs with Inf so that I can do what needs to be done. Any "sentinel" value would work here (such as -99), I though Inf was unambiguous. And since we replace all of these values with the preceding value or 0.5 later, we don't care that the intermediate value was Inf.
• We need to iterate over each of the collapsed (by rle) infinite values (formerly NA), and look at the preceding and following. Note: I explicitly stop (error) if the infinite (NA) values are first or last, but for a more robust function, you will likely need to replace stop("oops") with some better logic.
• By replacing the singular Inf with either the preceding value or 0.5, we can then use inverse.rle which (as might guess) converts the pair of values-and-lengths to a single vector, as before. (You might be comforted to know that inverse.rle(rle(vec)) should be an identical vector.)
• The call to ave(..) does a grouping function: cumsum(z2 == 0) helps us break the vector into groups where the first is a 0 (or the whole original vector unbroken, if 0 is never found). On each of these subvectors, we call cumsum. For instance, if vec <- c(1,2,0,3,4), then ave(vec, cumsum(vec==0), FUN=cumsum) will first call cumsum(c(1,2)) then call cumsum(c(0,3,4)), and the cumulative-sums are placed in the correct places in the returned vector. (The ave function also works well with grouping operations that do not assume contiguous groups, but that's for another question.)

From there, we call this function on each column using lapply.

Edit: I don't want to assume what logic you would want to use for first/last value being NA, but you might do something like this:

fun <- function(z) {
  r <- rle(replace(z, is.na(z), Inf))
  for (na in which(is.infinite(r$values))) {
    if (na == 1L) {
      r$values[na] <- 0
    } else if (na == length(r$values)) {
      r$values[na] <- r$values[na-1]
    } else {
      r$values[na] <- ifelse(r$values[na-1] == r$values[na+1], r$values[na-1], 0.5)
    }
  }
  z2 <- inverse.rle(r)
  ave(z2, cumsum(z2 == 0), FUN = cumsum)
}

You could use the following:

fn <- function(x){
x <- zoo::na.approx(x)
x[x>0&x<1] <- 0.5
grp <- data.table::rleid(!x)
data.frame(wave = grp, scored = ave(x, grp, FUN = cumsum))
}
cbind(df1, lapply(df1[-1], fn))
y_ex1.wave y_ex1.scored y_ex2.wave y_ex2.scored y_ex3.wave y_ex3.scored y_ex4.wave y_ex4.scored y_ex5.wave y_ex5.scored
1           1          1.0          1          1.0          1          1.0          1          1.0          1            1
2           1          1.5          1          1.5          1          1.5          1          1.5          1            2
3           2          0.0          2          0.0          2          0.0          2          0.0          2            0
4           3          1.0          2          0.0          3          1.0          2          0.0          2            0
5           3          2.0          3          0.5          3          1.5          2          0.0          3            1
6           3          3.0          3          1.0          3          2.0          2          0.0          3            2
7           3          4.0          3          1.5          3          2.5          2          0.0          3            3
8           3          5.0          3          2.5          4          0.0          2          0.0          4            0
9           4          0.0          4          0.0          4          0.0          2          0.0          4            0
10          4          0.0          4          0.0          4          0.0          2          0.0          4            0

In base R you could do:

fn <- function(x){
y <- cummax(seq_along(x)*!is.na(x))
idx <-which(is.na(x))
x[idx] <- (x[y[idx]] + x[rank(y,ties.method = 'max')[idx]+1])/2
grp <- cumsum(!x)
data.frame(score = ave(x,grp , FUN = cumsum), wave = grp)
}
cbind(df1, lapply(df1[-1], fn))
time y_ex1 y_ex2 y_ex3 y_ex4 y_ex5 y_ex1.score y_ex1.wave y_ex2.score y_ex2.wave
1     0     1     1     1     1     1         1.0          0         1.0          0
2     1    NA    NA    NA    NA     1         1.5          0         1.5          0
3     2     0     0     0     0     0         0.0          1         0.0          1
4     3     1     0     1     0     0         1.0          1         0.0          2
5     4    NA    NA    NA    NA     1         2.0          1         0.5          2
6     5    NA    NA    NA    NA     1         3.0          1         1.0          2
7     6    NA    NA    NA    NA     1         4.0          1         1.5          2
8     7     1     1     0     0     0         5.0          1         2.5          2
9     8     0     0     0     0     0         0.0          2         0.0          3
10    9     0     0     0     0     0         0.0          3         0.0          4