Home > OS >  Assign and save values based on function output in R
Assign and save values based on function output in R

Time:07-22

I'm stuck with something that should be very easy.

I have been working with a function that I want to take a vector or single column of data and output a "color" depending on which range the value falls into.

band_code <- function(x){
  x <- x[1]
  if(x < (-0.25)){color <- "cold"}
  else if(x > (-0.25) && x < (0.25)){color <- "off"}
  else if(x > (0.25) && x < (0.75)){color <- "intermediate"}
  else if(x > (0.75) && x < (1.25)){color <- "on"}
  else if(x > (1.25)){color <- "extra"}
  return(color)
}

This works well for simple, single tests:

> band_code(0.76)
[1] "on"

And this works well, using a test vector:

test <- -2:2

for (i in test){
  print(i)
  band = band_code(i)
  print(band)
}

[1] -2
[1] "cold"
[1] -1
[1] "cold"
[1] 0
[1] "off"
[1] 1
[1] "on"
[1] 2
[1] "extra"

But when I try to save the output as a column in an object, it fails:

test <- as.data.frame(test)

for (i in test){
  band = band_code(i)
  test$band <- band
}

> test
  test band
1   -2 cold
2   -1 cold
3    0 cold
4    1 cold
5    2 cold

The "band" values should be as they were before ("cold", "cold", "off", "on", "extra"), but I guess it is repeating the first "cold" all 5 times.

What about saving the output to an object causes the malfunction?

I think it has to do with the first line in the band_code function (x <- x[1]), but after many various attempts I can't get it to work.

Help, please!

CodePudding user response:

You should use cut() for situations like this:

test <- -2:2
test <- as.data.frame(test)
test$vals <- cut(test$test, breaks=c(-Inf, -.25, .25, .75, 1.25, Inf), labels=c("cold", "off", "intermediate", "on", "extra"))
test
#   test  vals
# 1   -2  cold
# 2   -1  cold
# 3    0   off
# 4    1    on
# 5    2 extra

CodePudding user response:

A tidy solution is to use dplyr::case_when, which can be vectorized unlike if and else.

band_code <- function(x) {
  dplyr::case_when(
    x < -0.25 ~ "cold",
    x > (-0.25) & x < (0.25) ~ "off",
    x > (0.25) & x < (0.75) ~ "intermediate",
    x > (0.75) & x < (1.25) ~ "on",
    x > (1.25) ~ "extra",
    TRUE ~ NA_character_
  )
}
data.frame(colors = band_code(-10:10))
  • Related