bump plot with jitter

bump.R

library(ggplot2)
library(magrittr)
library(tidyr)
library(dplyr)

sigmoid2 <- function(x_from, x_to, y_from, y_to, smooth = 5, n = 100, direction = "x", location = 0, scale = 1) {
  if(!direction %in% c("x", "y")) {stop("Only the directions x or y is allowed.")}
  
  if(direction == "x") {
    x <- seq(-smooth, smooth, length = n)
    y <- stats::plogis(x, location = location, scale = scale)
    out <- data.frame(x = (x + smooth) / (smooth * 2) * (x_to - x_from) + x_from,
                      y = y * (y_to - y_from) + y_from)
  }
  
  if(direction == "y") {
    y <- seq(-smooth, smooth, length = n)
    x <- stats::plogis(y, location = location, scale = scale)
    out <- data.frame(y = (y + smooth) / (smooth * 2) * (y_to - y_from) + y_from,
                      x = x * (x_to - x_from) + x_from)
  }
  out
}


rank_sigmoid2<- function(x, y, smooth = 8, direction = "x") {
  .df <- dplyr::tibble(x = x,
                       y = y) %>%
    dplyr::mutate(x_lag = dplyr::lag(x),
                  y_lag = dplyr::lag(y)) %>%
    tidyr::drop_na("x_lag") %>% 
    dplyr::mutate(loc = stats::runif(1, -2, 2))
  purrr::pmap_dfr(.df, ~sigmoid2(x_from = ..3, x_to = ..1, y_from = ..4, y_to = ..2, smooth  = smooth, direction = direction, location  = ..5))
}

StatBumpJitter <- ggplot2::ggproto("StatBump", ggplot2::Stat,
                             setup_data = function(data, params) {
                               # Create x_lag, and y_lag to be passed to `compute_group`
                               # Factors need this to be able to compute a sigmoid function
                               data <- data %>%
                                 dplyr::mutate(r = dplyr::row_number()) %>%
                                 dplyr::arrange(x) %>%
                                 dplyr::group_by_at(vars(-PANEL, -group, -x, -y, -r)) %>%
                                 dplyr::mutate(x_lag = dplyr::lag(x),
                                               y_lag = dplyr::lag(y)) %>%
                                 dplyr::ungroup() %>%
                                 dplyr::arrange(r) %>%
                                 dplyr::select(-r) %>%
                                 as.data.frame()
                               data
                             },
                             compute_group = function(data, scales, smooth = 8, direction = "x", location = location, scale = scale) {
                               data <- data %>%
                                 dplyr::arrange(x)
                               
                               # Handling of the special case of factors
                               # Factors come as a df with one row
                               if(nrow(data) == 1) {
                                 if(is.na(data$x_lag) | is.na(data$y_lag)) {
                                   return(data %>% dplyr::slice(0))
                                 } else {
                                   out <- sigmoid2(data$x_lag, data$x, data$y_lag, data$y,
                                                  smooth = smooth, direction = direction)
                                   return(as.data.frame(out))
                                 }
                               }
                               
                               # Normal case
                               out <-rank_sigmoid2(data$x, data$y, smooth = smooth, direction = direction) %>%
                                 dplyr::mutate(key = 1) %>%
                                 dplyr::left_join(data %>%
                                                    dplyr::select(-x, -y) %>%
                                                    dplyr::mutate(key = 1) %>%
                                                    dplyr::distinct(),
                                                  by = "key", relationship = "many-to-many") %>%
                                 dplyr::select(-key) %>%
                                 as.data.frame()
                               out
                             },
                             
                             required_aes = c("x", "y")
)

geom_bumpjitter <- function(mapping = NULL, data = NULL, geom = "line",
                      position = "identity", na.rm = FALSE, show.legend = NA,
                      smooth = 8, direction = "x", inherit.aes = TRUE, ...) {
  ggplot2::layer(
    stat = StatBumpJitter, data = data, mapping = mapping, geom = geom,
    position = position, show.legend = show.legend, inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, smooth = smooth, direction = direction, ...)
  )
}

dfiris <- mutate(iris, id=row_number()) %>% 
  pivot_longer(1:4, names_to = "variable")

ggplot(dfiris, aes(x = variable, y = value, group = id, colour=Species)) +
  geom_bumpjitter(linewidth = 0.1)+
  geom_point(alpha=0.1)+
  guides(colour = guide_legend(override.aes = list(alpha=0.7)))+
  theme_bw(16)

ggsave("iris_bump.png")