brownian.R

library(tidyverse)
theme_set(theme_bw())
STEPS = 251
ITER = 5000 

#Function to get the change to X and Y corresponding
#to each die roll
get_offset = function(roll)
{
  nx = ny = NA
  if (roll == 1) { nx = .5 ;ny = sqrt(3 / 4) }
  else if (roll == 2) { nx = 1; ny = 0 }
  else if (roll == 3) { nx = .5; ny = -sqrt(3 / 4) }
  else if (roll == 4) { nx = -.5; ny = -sqrt(3 / 4) }
  else if (roll == 5) { nx = -1; ny = 0 }
  else if (roll == 6) { nx = -.5; ny = sqrt(3 / 4) }
  else { stop("case fail") }
  return(list(nx, ny))
}

#Create a new row in the dataframe of step coordinates
#Each row picks up from where previous left off
#Row 1 is special (begins and ends at origin)
#Parameters:
#  i: row to update
#  noback: if TRUE, re-rolls when a step would undo last step
update_curr = function(i, noback) {
  if (i < 2 | i > STEPS) {
    stop("i must be between 2 and STEPS")
  }
  roll = sample(1:6, 1)
  if (noback) {
    while (roll == (df[i - 1, 'roll'] + 3) %% 6) {
      roll = sample(1:6, 1)
    }
  }
  df[i, 'roll'] <<- roll
  tmp = get_offset(roll)
  #Start point of this row is where last row ended
  df[i, c("x", "y")] <<- c(df[i - 1, "xend"], df[i - 1, "yend"])
  #End point of this row is start point plus changes
  df[i, c("xend", "yend")] <<-
    c(df[i, "x"] + tmp[[1]], df[i, "y"] + tmp[[2]])
}

#List to collect up the iterations
retList = vector('list', ITER)

system.time({
  for (k in 1:ITER) {
    df = data.frame(
      x = c(0, rep(NA, STEPS - 1)),
      y = c(0, rep(NA, STEPS - 1)),
      xend = c(0, rep(NA, STEPS - 1)),
      yend = c(0, rep(NA, STEPS - 1)),
      draw = rep(0, STEPS),
      step = 0:(STEPS - 1),
      group = k
    )
    for (i in 2:STEPS) {
      update_curr(i, FALSE)
    }
    retList[[k]] = df
  }
})

bdf = do.call('rbind', retList) #collect results into big df

#get root mean square distance and average X Y position by step
sdf = bdf %>% mutate(dist = sqrt((xend) ^ 2 + (yend) ^ 2)) %>%
  group_by(step) %>%
  summarise(
    root_mean_sq_av = sqrt(mean(dist ^ 2)),
    mean_x = mean(xend),
    mean_y = mean(yend)
  ) %>%
  ungroup()

#Plot the last walk
mdf = df[c(1, nrow(df)),]
limit = max(c(-min(df$xend), max(df$xend), -min(df$yend), max(df$yend)))
p = ggplot(df, aes(x = x, y = y, xend = xend, yend = yend,color = step))
p = p + coord_cartesian(xlim = c(-limit, limit),
                        ylim = c(-limit, limit)) +
  geom_segment() + scale_color_gradient(low = "blue", high = "red") +
  labs(x = "X position", y = "Y position", title = "One random walk") +
  geom_point(data = mdf, size = 2, aes(x = xend, y = yend)) +
  theme(legend.position = "bottom")
p
ggsave("lastwalk.png",
  p,height = 4, width = 4, dpi = 600)

#Ave position
p = ggplot(sdf %>% gather(id, position, mean_x, mean_y),
           aes(x = step, y = position, color = id))
p = p + geom_line() +
  coord_cartesian(xlim = c(0, STEPS),
                  ylim = c(-limit * .25, limit * .25)) +
  scale_colour_discrete(
    name  = "",
    breaks = c("mean_x", "mean_y"),
    labels = c("Average X position", "Average Y position")) +
  labs(x = "Step", y = "Position", title = "Average X,Y positions at each step") +
  theme(legend.position = "bottom")
p
ggsave("ave_position.png",
  p, height = 4, width = 4, dpi = 600)

#Root mean square distance from origin
p = ggplot(sdf, aes(x = step, y = root_mean_sq_av))
p = p + geom_point(size = .1) +
  geom_line(aes(x = step, y = (step) ^ (1 / 2)), linetype = 'dashed') +
  labs(x = "Steps", y = "Root mean square distance",
       title = "Average distance from start")
p
ggsave("dist_from_origin.png",
  p, height = 4, width = 4, dpi = 600)

#Zoom in on the last plot and save it
p = p + coord_cartesian(xlim = c(0, 10), ylim = c(0, 4)) +
  geom_point(size = 2) +
  labs(title = "Zoomed in average distance from start")
p
ggsave("zoomed_dist_from_origin.png",
  p, height = 4, width = 4, dpi = 600)