# ============================================================================
# Random walks: recurrence, transience, and limits on an infinite state space
# ============================================================================
#
# A simple random walk on the integer lattice Z^d starts at the origin and, at
# every step, moves to one of its 2d nearest neighbors with equal probability.
# The state space is infinite. Does the walker keep coming back?
#
# Part I works through everything for d = 1. Part II compares d = 1, 2, 3.

## ===========================================================================
## PART I. A WALK ON THE INTEGERS
## ===========================================================================
##
## A gambler with no barriers wins or loses $1 on each bet with equal
## probability. X_t is their net winnings after t bets, starting from X_0 = 0.
##
## After a first look at some paths, we ask of this chain: Is it irreducible?
## Is it periodic? Is it recurrent? Do its n-step transition probabilities have
## limits? Does the distribution of X_t settle down?

## ---------------------------------------------------------------------------
## 1. PATHS
## ---------------------------------------------------------------------------

## --- One step ----------

x <- 0L                                        # start at 0
x <- x + sample(c(-1L, 1L), 1); x              # one step; run this line several times

## --- A path, step by step ----------

x <- integer(21)                               # x[t + 1] will hold X_t
for (t in 1:20) {
  x[t + 1] <- x[t] + sample(c(-1L, 1L), 1)
}
x

## --- Wrap it in a function ----------
##
## `moves` lists the possible steps, each equally likely. We change it later.

walk_1d <- function(m, moves = c(-1L, 1L)) {
  x <- integer(m + 1)
  for (t in 1:m) {
    x[t + 1] <- x[t] + sample(moves, 1)
  }
  x
}

set.seed(5710)
x <- walk_1d(2000)

plot(0:2000, x, type = "l", xlab = "t (steps)", ylab = expression(X[t]))
abline(h = 0, lty = 3, col = "red")

## Check: what was the gambler's biggest loss, and after how many bets?

## ---------------------------------------------------------------------------
## 2. IS IT IRREDUCIBLE?
## ---------------------------------------------------------------------------
##
## Can the walker get from 0 to any other state? Walk until it first reaches a
## target state, and count the steps.

target   <- 10
position <- 0L
n_steps  <- 0
while (position != target && n_steps < 1e6) {   # give up after a million steps
  position <- position + sample(c(-1L, 1L), 1)
  n_steps  <- n_steps + 1
}
n_steps

## Check: rerun this block a few times, then with target <- -5 and target <- 25.
##        How many steps did each run take?
## Q. Does the walker always get there? Can a simulation show that the chain is
##    irreducible, or does that take an argument?

## ---------------------------------------------------------------------------
## 3. IS IT PERIODIC?
## ---------------------------------------------------------------------------
##
## Back to the 2000-step path x from section 1. When is the walker at 0?

at0   <- x == 0L                               # TRUE at each step where the walker is at 0
steps <- 0:2000                                # the step number of each entry of x

steps_at0 <- steps[at0]                        # the steps t at which X_t = 0
steps_at0

length(steps_at0)                              # visits to 0, counting the start
steps_at0[2]                                   # first return time T_0: first visit after step 0

## Check: what do the steps in steps_at0 have in common? What was the longest
##        wait between two visits to 0?
## Q. What is the implication for the period of state 0? Why might we want to
##    fix this?

## --- A lazy walker ----------
##
## Adds 0 to the possible moves, so -1, 0, and 1 are equally likely.

set.seed(5710)
x_lazy <- walk_1d(2000, moves = c(-1L, 0L, 1L))

## Check: at which steps is the lazy walker at 0? Are they all even?

## ---------------------------------------------------------------------------
## 4. IS IT RECURRENT?
## ---------------------------------------------------------------------------
##
## One path can't tell us whether the walker is sure to come back. We need
## thousands, all moving at once: row r of X is walker r's path.

## --- Five walkers, step by step ----------

m     <- 20                                    # steps
n_sim <- 5                                     # walkers

X <- matrix(0L, nrow = n_sim, ncol = m + 1,
            dimnames = list(sim = NULL, step = 0:m))
for (t in 1:m) {
  X[, t + 1] <- X[, t] + sample(c(-1L, 1L), n_sim, replace = TRUE)   # every walker steps
}
X

## --- Wrap it in a function ----------
##
## `start` is where the walkers begin: one position for everyone, or one for each
## walker. We use it again in section 6.

sim_walks_1d <- function(m, n_sim, moves = c(-1L, 1L), start = 0L) {
  X <- matrix(0L, nrow = n_sim, ncol = m + 1,
              dimnames = list(sim = NULL, step = 0:m))
  X[, 1] <- start                                # positions at step 0
  for (t in 1:m) {
    X[, t + 1] <- X[, t] + sample(moves, n_sim, replace = TRUE)
  }
  X
}

# same seed, one walker: sim_walks_1d reproduces walk_1d exactly
set.seed(1); one_walker <- walk_1d(500)
set.seed(1); one_row    <- sim_walks_1d(500, n_sim = 1)
all(one_walker == one_row[1, ])

## --- Thousands of walkers ----------

m     <- 2000
n_sim <- 2000

set.seed(83126)
X <- sim_walks_1d(m, n_sim)

I0 <- X == 0L                                  # at0 from section 3, one row per walker
I0[1:5, 1:13]                                  # first five walkers, steps 0 to 12

## Check: how many walkers are at 0 after exactly 2 steps?

## --- Lens 1: does the walker come back? ----------

steps <- 0:m

# walker 1, exactly as in section 3
steps_at0 <- steps[I0[1, ]]
steps_at0[2]

# every walker: the same calculation, in a loop
T0 <- rep(Inf, n_sim)                          # Inf means "not back within m steps"
for (r in 1:n_sim) {
  steps_at0 <- steps[I0[r, ]]                  # steps at which walker r is at 0
  if (length(steps_at0) > 1) {
    T0[r] <- steps_at0[2]                      # its first visit after step 0
  }
}
T0[1:20]

# for each n, the share of walkers whose first return came by step n
back_by <- rep(NA, m)
for (n in 1:m) {
  back_by[n] <- mean(T0 <= n)
}

plot(1:m, back_by, type = "l", log = "x", ylim = c(0, 1),
     xlab = "n (steps, log scale)", ylab = "share of walkers back at 0 by step n")

## Check: how many walkers never came back? What was the longest return time
##        among those who did?
## Q. Is the curve heading toward f_0 = P(T_0 < infinity | X_0 = 0) = 1?

## --- Lens 2: how many times does it come back? ----------

# a running count of visits to 0 for each walker: column t + 1 counts visits by step t
visits_so_far <- matrix(0L, nrow = n_sim, ncol = m + 1, dimnames = dimnames(I0))
visits_so_far[, 1] <- I0[, 1]                  # the start counts as a visit
for (t in 1:m) {
  visits_so_far[, t + 1] <- visits_so_far[, t] + I0[, t + 1]   # one more if at 0 at step t
}
visits_so_far[1:5, 1:13]

N0 <- visits_so_far[, m + 1]                   # each walker's total visits over all m steps
summary(N0)

avg_visits <- colMeans(visits_so_far)          # average over walkers, at each step

plot(0:m, avg_visits, type = "l",
     xlab = "n (steps)", ylab = "average visits to 0 by step n")

## Check: which walker visited 0 most often, and how many times?
## Q. Is the average number of visits leveling off?

## --- Lens 3: n-step transition probabilities ----------

mean(I0[, "2"])                  # share of walkers at 0 after 2 steps: estimates p_00^(2)

p00 <- colMeans(I0)              # the same share for every step: p00["n"] estimates p_00^(n)
round(p00[1:11], 3)              # steps 0 to 10

plot(0:60, p00[1:61], type = "h", lwd = 2, ylim = c(0, 1),
     xlab = "n (steps)", ylab = expression(p["00"]^(n)))

# a running total of p00 gives back the average visits from lens 2
all.equal(cumsum(p00), avg_visits)

## Check: compare p00 at n = 2 and n = 4 with the exact values 1/2 and 3/8.
## Q. Why does the running total of p_00^(n) equal the average number of visits?
##    What does recurrence, sum_n p_00^(n) = infinity, look like in lens 2's plot?

## ---------------------------------------------------------------------------
## 5. LIMITS
## ---------------------------------------------------------------------------
##
## Where are p_00^(n) and p_0,10^(n) headed as n grows? The 2000 steps in X are
## not enough to tell. To go further with more walkers, keep only each walker's
## current position, and record the share at 0 and at 10 after every step.

m     <- 16384                                 # steps
n_sim <- 10000                                 # walkers

set.seed(83126)
position   <- rep(0L, n_sim)                   # every walker starts at 0
p00_long   <- rep(NA, m)                       # p00_long[n]: share of walkers at 0 after n steps
p0_10_long <- rep(NA, m)                       # p0_10_long[n]: share of walkers at 10 after n steps
for (n in 1:m) {
  position      <- position + sample(c(-1L, 1L), n_sim, replace = TRUE)   # every walker steps
  p00_long[n]   <- mean(position == 0L)
  p0_10_long[n] <- mean(position == 10L)
}

n_by_4 <- 4^(1:7)                              # 4, 16, 64, ..., 16384: each 4 times the last
round(p00_long[n_by_4], 4)

## Check: what happens to p_00^(n) each time n is multiplied by 4? Where is that
##        headed?

even_n <- seq(2, m, by = 2)                    # at odd n no walker is at 0 or 10 (section 3)

plot(even_n, p00_long[even_n], type = "l", log = "x", ylim = c(0, 0.5),
     xlab = "n (steps, log scale)", ylab = "share of walkers at 0 or 10 after n steps")
lines(even_n, p0_10_long[even_n], col = "red")
points(n_by_4, p00_long[n_by_4], pch = 19)     # the values printed above
abline(h = 0, lty = 3)
legend("topright", legend = expression(p["00"]^(n), p["0,10"]^(n)),
       col = c("black", "red"), lty = 1, bty = "n")

## Q. What happens to p_0j^(n) as n grows, for any fixed position j?
## Q. Does p_00^(n) converge, even though the walk is periodic? Compare with what
##    periodicity did on a finite chain.
## Q. On a finite chain the rows of P^n sum to 1, and so does their limit. Is
##    there a limiting distribution here?

## ---------------------------------------------------------------------------
## 6. DOES THE DISTRIBUTION OF X_t SETTLE DOWN?
## ---------------------------------------------------------------------------
##
## Column "t" of X holds every walker's position after t steps: a sample from
## the distribution of X_t.

X[1:10, "100"]                                 # the first ten walkers, after 100 steps

counts <- table(X[, "100"])                    # number of walkers at each position
counts

op <- par(mfrow = c(1, 3), mar = c(4, 4.2, 3, 1))
for (t in c("10", "100", "1000")) {
  counts <- table(X[, t])                      # number of walkers at each position after t steps
  plot(counts, xlim = c(-100, 100), ylim = c(0, 550), axes = FALSE,
       main = paste("after", t, "steps"), xlab = "position", ylab = "number of walkers")
  axis(1); axis(2)                             # plain axes, not a label for every position
}
par(op)

## Check: what share of walkers are within 10 of the origin after 10 steps? After 1000?
## Q. Is the distribution of X_t settling down to a stationary distribution?

## --- Starting spread out ----------
##
## Walkers that all start at 0 form a bell that keeps spreading. Instead, spread
## 5000 walkers evenly over the positions -200 to 200, then follow a simple walk
## and a lazy walk from the same starting positions.

set.seed(83126)
start <- sample(-200:200, 5000, replace = TRUE)       # a starting position for each walker

X_spread      <- sim_walks_1d(m = 1000, n_sim = 5000, start = start)
X_spread_lazy <- sim_walks_1d(m = 1000, n_sim = 5000, start = start,
                              moves = c(-1L, 0L, 1L))

bins <- seq(-1200, 1200, by = 10)                     # count walkers in bins 10 positions wide

op <- par(mfrow = c(2, 3), mar = c(4, 4.2, 3, 1))
for (t in c("0", "100", "1000")) {
  hist(X_spread[, t], breaks = bins, xlim = c(-400, 400), ylim = c(0, 200),
       main = paste("simple, after", t, "steps"), xlab = "position", ylab = "number of walkers")
}
for (t in c("0", "100", "1000")) {
  hist(X_spread_lazy[, t], breaks = bins, xlim = c(-400, 400), ylim = c(0, 200),
       main = paste("lazy, after", t, "steps"), xlab = "position", ylab = "number of walkers")
}
par(op)

## Check: how many simple walkers are between -50 and 50 at step 0? At step 1000?
##        How many are more than 200 from the origin at step 1000?
## Q. Which part of the picture stays put, and which part changes? Does it matter
##    that the simple walk is periodic?
## Q. Show that mu(j) = 1 for every position j satisfies mu P = mu, for both walks.
## Q. Why can't mu be scaled into a stationary distribution?

## ===========================================================================
## PART II. MORE DIMENSIONS
## ===========================================================================
##
##   d = 2: a pedestrian on a city grid, picking a direction at each corner
##   d = 3: a molecule hopping between sites of a crystal lattice
##
## A step now picks one of the d coordinates, then a direction, so each of the
## 2d neighbors has probability 1/(2d).

step_lattice <- function(pos) {                 # a row per walker, a column per coordinate
  n_sim <- nrow(pos)
  d     <- ncol(pos)
  coord <- sample.int(d, n_sim, replace = TRUE)       # which coordinate each walker changes
  dir   <- sample(c(-1L, 1L), n_sim, replace = TRUE)  # and in which direction
  for (j in 1:d) {
    movers <- coord == j                              # the walkers changing coordinate j
    pos[movers, j] <- pos[movers, j] + dir[movers]
  }
  pos
}

pos <- matrix(0L, nrow = 5, ncol = 2)            # five pedestrians at the corner (0, 0)
pos <- step_lattice(pos); pos                    # run this line several times

## --- Simulate ----------
##
## Keeping every coordinate at every step gets expensive. The three lenses below
## only need each walker's distance from the origin, |x_1| + ... + |x_d| (on a
## city grid, blocks from the start). It is 0 exactly when the walker is back.

sim_distance <- function(m, d, n_sim) {
  pos  <- matrix(0L, nrow = n_sim, ncol = d)
  dist <- matrix(0L, nrow = n_sim, ncol = m + 1,
                 dimnames = list(sim = NULL, step = 0:m))
  for (t in 1:m) {
    pos <- step_lattice(pos)
    dist[, t + 1] <- as.integer(rowSums(abs(pos)))    # |x_1| + ... + |x_d|, per walker
  }
  dist
}

m     <- 5000
n_sim <- 1000

set.seed(83126)
dist1 <- sim_distance(m, d = 1, n_sim)
dist2 <- sim_distance(m, d = 2, n_sim)
dist3 <- sim_distance(m, d = 3, n_sim)

dist2[1:5, 1:13]                                 # first five pedestrians, steps 0 to 12

cols <- palette.colors(palette = "Okabe-Ito")[c(1, 6, 7)]   # black, blue, vermillion
dims <- c("d = 1", "d = 2", "d = 3")

## --- Lens 1: distance from the origin ----------

walker_1 <- cbind(dist1[1, ], dist2[1, ], dist3[1, ])   # the first walker in each dimension

matplot(0:m, walker_1, type = "l", lty = 1, col = cols,
        xlab = "t (steps)", ylab = "distance from origin")
legend("topleft", dims, col = cols, lty = 1, bty = "n")

## Q. Which of these walkers came back to the origin, and how often? Look at a
##    few other rows (walkers) before deciding.

## --- Section 4's calculations, as functions ----------
##
## The same loops as in Part I, wrapped so we can repeat them in each dimension.

first_return <- function(I0) {
  steps <- 0:(ncol(I0) - 1)
  T0 <- rep(Inf, nrow(I0))                       # Inf means "not back within the simulation"
  for (r in 1:nrow(I0)) {
    steps_at0 <- steps[I0[r, ]]
    if (length(steps_at0) > 1) {
      T0[r] <- steps_at0[2]
    }
  }
  T0
}

share_back_by <- function(T0, m) {
  back_by <- rep(NA, m)
  for (n in 1:m) {
    back_by[n] <- mean(T0 <= n)
  }
  back_by
}

## --- Lens 2: return probability ----------

T0_1 <- first_return(dist1 == 0L)                # each walker's first return time, d = 1
T0_2 <- first_return(dist2 == 0L)                #                                  d = 2
T0_3 <- first_return(dist3 == 0L)                #                                  d = 3

back_by_d <- cbind(share_back_by(T0_1, m), share_back_by(T0_2, m), share_back_by(T0_3, m))

matplot(1:m, back_by_d, type = "l", lty = 1, col = cols, log = "x", ylim = c(0, 1),
        xlab = "n (steps, log scale)", ylab = "share of walkers back at 0 by step n")
legend("bottomright", dims, col = cols, lty = 1, bty = "n")

## Q. Which dimensions look recurrent? Transient? Can you tell for d = 2?

## --- Lens 3: average visits ----------

p00_1 <- colMeans(dist1 == 0L)                   # share of walkers at the origin after each step
p00_2 <- colMeans(dist2 == 0L)
p00_3 <- colMeans(dist3 == 0L)

avg_visits_d <- cbind(cumsum(p00_1), cumsum(p00_2), cumsum(p00_3))   # running totals, as in lens 3

matplot(1:m, avg_visits_d[-1, ], type = "l", lty = 1, col = cols, log = "x",
        ylim = c(1, 6), xlab = "n (steps, log scale)",
        ylab = "average visits to 0 by step n")    # [-1, ] drops step 0 for the log scale
legend("topleft", dims, col = cols, lty = 1, bty = "n")

## Q. Which curves are headed to infinity? Does that settle d = 2?