CODEWITHSUNDEEP
rggplot2
Luiz Rodrigo
Luiz Rodrigo

Reputation: 936

Adding colors and legends to stat_function() programatically

I'm trying to code a function that plots on the same figure approximations to the solution of an ODE using different step values. I got the ODE approximations right, I just can't figure out how to add colors and legends identifying each function.

I tried to follow this answer, but I can't translate it quite well to a context where the number of functions is not constant.

Here's my code and the output it generates. 

library(purrr)
library(ggplot2)
library(glue)

eulerMethod = function(f, t0, y0, h, memo = 1000) {
  vec = double(memo + 1)

  vec[1] = y0
  for (i in 1:memo) {
    vec[i+1] = vec[i] + h*f(t0 + i*h, vec[i])
  }

  solution = function(t) {
    if (t < t0) return(NaN)
    n = (t-t0)/h

    intN = floor(n)
    if (n == intN)
      return(vec[n+1])
    else # linear interpolation
      return(vec[intN + 1] + (n - intN) * (vec[intN + 2] - vec[intN + 1]))
  }
}

compare = function(f, t0, y0, interval, hs = c(1, 0.5, 0.2, 0.1, 0.05)) {
  fs = map(hs, ~ eulerMethod(f, t0, y0, .)) %>% 
    map(Vectorize)

  # generates "h = 1" "h = 0.5" ... strings
  legends = map_chr(hs, ~ glue("h = {hs[[.]]}"))
  map(1:length(hs), ~ stat_function(fun = fs[[.]],
                                    geom = "line",
                                    aes_(colour = legends[.]))) %>%
    reduce(`+`, .init = ggplot(data.frame(x = interval), aes(x)))
}

# y' = y => solution = exp(x)
compare(function(t, y) y, 0, 1, c(0, 5))

All functions are plotted with the same color and legend.

Upvotes: 2

Views: 2084

Answers (1)

eipi10
eipi10

Reputation: 93761

I haven't used glue before, but it wasn't working the way I think you expected it to. It was just returning five copies of h = 1. I've modified your code to just use paste0 to create the legend values:

compare = function(f, t0, y0, interval, hs = c(1, 0.5, 0.2, 0.1, 0.05)) {
  fs = map(hs, ~ eulerMethod(f, t0, y0, .)) %>% 
    map(Vectorize)

  # generates "h = 1" "h = 0.5" ... strings
  legends = paste0("h = ", hs)

  map(1:length(hs), ~ stat_function(fun = fs[[.]],
                                    geom = "line",
                                    aes_(colour = legends[.]))) %>%
    reduce(`+`, .init = ggplot(data.frame(x = interval), aes(x, colour=.)))
}

compare(function(t, y) y, 0, 1, c(0, 5))

enter image description here

Also, it seems to me the code could be made a bit more straightforward. For example:

compare = function(f, t0, y0, interval, hs = c(1, 0.5, 0.2, 0.1, 0.05)) {

  fs = map(hs, ~ eulerMethod(f, t0, y0, .)) %>% 
    map(Vectorize)

  ggplot(data.frame(x = interval), aes(x)) +
    map(1:length(fs), function(nn) {
      stat_function(fun = fs[[nn]], geom = "line", aes_(colour = factor(hs[nn])))
    }) +
    labs(colour="h")
}

compare(function(t, y) y, 0, 1, c(0, 5))

Upvotes: 2

Related Questions