plot data & model

plt.R

plt <- function(df, mod){
  # predict only for range in site
  sites <- levels(df$site)
  value_0p <- c(sapply(sites, function(y)  seq(min(df$value_0[df$site == y]), 
                                               max(df$value_0[df$site == y]), 
                                               length.out = 333)))
  # use the mean for init
  pdat <- data.frame(value_0 = value_0p, 
                     site = rep(sites, each = 333), 
                     init = mean(df$init)
  )
  
  if ('lm' %in% class(mod)) {
    pdat$fit <- predict(mod, newdata = pdat)
    crit <- qnorm(0.975)
    pdat$err <- crit*predict(mod, newdata = pdat, se.fit = TRUE)$se.fit
    pdat$upr <- pdat$fit + pdat$err
    pdat$lwr <- pdat$fit - pdat$err
  }
  if ('glm' %in% class(mod)) {
    # calculate error on link scale and then transform to response scale
    fit <- predict(mod, newdata = pdat, type = 'link')
    crit <- qnorm(0.975)
    err <- crit*predict(mod, newdata = pdat, type = 'link', se.fit = TRUE)$se.fit
    upr <- pdat$fit + pdat$err
    lwr <- pdat$fit - pdat$err
    # inverse link
    pdat$fit <- mod$family$linkinv(fit)
    pdat$upr <- mod$family$linkinv(upr)
    pdat$lwr <- mod$family$linkinv(lwr)
  }

  # extract dep variable
  dep <- all.vars(mod$terms)[1]
  
  # plot data, model & error
  p <- ggplot() +
    geom_ribbon(data = pdat, aes(x = value_0, ymax = upr, ymin = lwr, fill = site), alpha = 0.2) +
    geom_line(data = pdat, aes(x = value_0, y = fit, col = site, group = site)) +
    geom_point(data = df, aes_string(x = 'value_0', y = dep, col = 'site')) +
    facet_wrap(~site) +
    theme_bw()
  p
}