R/random_walk.R

In matthewkling/windscape: Landscape connectivity by wind

setClass("wind_walk",
         contains = "SpatRaster",
         slots = c(mode = "character",
                   n_iter = "numeric",
                   iter_length = "numeric"))


as_wind_walk <- function(x, mode, n_iter, iter_length){
      if(!inherits(x, "SpatRaster")) stop("x must be a SpatRaster")
      x <- as(as(x, "SpatRaster"), "wind_walk")
      x@mode <- mode
      x@n_iter <- n_iter
      x@iter_length <- iter_length
      x
}

#' Iteration step length of a random walk
#'
#' @param x A \code{wind_walk} generated by \code{random_walk()}.
#'
#' @export
iter_length <- function(x) x@iter_length


#' Maximum iteration duration for a random walk
#'
#' Calculate the maximum possible iteration length for a random walk simulation, which is the
#' residence time of the grid cell with the greatest total conductance. Walks can proceed no
#' faster than this.
#'
#' @param rose A \code{wind_rose}.
#'
#' @export
rw_max_step <- function(rose){
      1 / minmax(sum(rose))[2, ]
}


rw_prob <- function(x, t){
      p <- sum(x)
      p <- 1 / t - p
      p <- c(p, x)
      p / sum(p)
}



#' Simulate diffusion by wind advection
#'
#' This function estimates diffusion across a windscape by Markov random walk. On each iteration,
#' the "particle mass" in each grid cell is dispersed within the local 9-cell neighborhood in
#' proportion with wind conductance. Depending on usage, this "mass" could represent probability,
#' number of individuals, etc. As the simulation proceeds, the mass diffuses across the landscape.
#'
#' The input wind rose raster is converted to a simplex of nine probabilities for each grid
#' cell, giving the rates at which particles are retained in a cell or moved to each of its
#' eight neighbors. Probabilities of moving to a neighboring cell are proportional to conductance
#' in the wind rose data set, with the probabilities of remaining in a cell scaled so that the
#' cell with the highest conductance has a zero retention probability; this allows conductance
#' to be normalized locally to a simplex while remaining proportional across cells and
#' maximizing the dispersal occurring at each iteration.
#'
#' The `mode` argument determines how particle mass propagates over time. The default method,
#' `"pulse"`, tracks the fleeting diffusion of the initial starting particle mass, which will drift
#' and spread like a cloud, eventually leaving the modeling domain (unless it's a closed spatial
#' domain). The `mode = "ratchet"` option runs a propagating simulation in which local particle
#' mass never declines in any cell, with every location continuing to transmit mass at the
#' cumulative maximum rate; this is more akin to a biological process in which dispersing particles
#' reproduce locally after establishment, or in which a continuous stream of particles is released
#' from the original source at every time step.
#'
#' @param rose A \code{wind_rose}.
#' @param init Initial conditions from which to begin diffusion, either a two-column
#'    matrix of coordinates, or a SpatRaster layer with non-negative mass values and the
#'    same spatial properties as \code{rose}. If coordinates, the simulation starts with a
#'    mass of 1 at each coordinate location. If a SpatRaster, diffusion is done
#'    directly on the raster; this could be the ouput from a prior random_walk, or any
#'    other data representing quantities to be spatially dispersed.
#' @param iter Number of simulation iterations (positive integer).
#' @param record Integer vector specifying which iterations to record. Default is to
#'    record only the final iteration, i.e. \code{iter}. One raster layer is returned for
#'    each value in \code{record}.
#' @param mode Propagation mode, either "pulse" (default) or "ratchet". See details.
#' @param timescale A value between 0 and 1, giving the factor by which to
#'    scale the default time step length (which is calculated from the data; see \link{rw_max_step}).
#'    At each iteration, particle mass either exits cells at the rates given in the wind rose object,
#'    or remains in the cell. The default timescale value of 1 sets the timestep length to the
#'    maximum possible value, allowing the simulation to advance as far as possible in space given the
#'    number of iterations, which is computationally optimal. Reducing this value may be useful for
#'    smoothing the simulation dynamics, and/or setting the timestep to a desired duration.
#'
#' @return A \code{wind_walk} raster object.
#'
#' @export
random_walk <- function(rose, init, iter = 100, record = iter, mode = "pulse", timescale = 1){

      disperse <- function(n, p){
            a <- p
            for(k in 1:9) a[,,k] <- n * a[,,k]
            da <- dim(a)
            a[,,2] <- rbind(rep(0, da[2]), cbind(a[1:(da[1]-1), 2:(da[2]), 2], rep(0, da[1]-1))) # SW
            a[,,3] <- cbind(a[1:(da[1]), 2:(da[2]), 3], rep(0, da[1])) # W
            a[,,4] <- rbind(cbind(a[2:(da[1]), 2:(da[2]), 4], rep(0, da[1]-1)), rep(0, da[2])) # NW
            a[,,5] <- rbind(a[2:(da[1]), 1:(da[2]), 5], rep(0, da[2])) # N
            a[,,6] <- rbind(cbind(rep(0, da[1]-1), a[2:(da[1]), 1:(da[2]-1), 6]), rep(0, da[2])) # NE
            a[,,7] <- cbind(rep(0, da[1]), a[1:(da[1]), 1:(da[2]-1), 7]) # E
            a[,,8] <- rbind(rep(0, da[2]), cbind(rep(0, da[1]-1), a[1:(da[1]-1), 1:(da[2]-1), 8])) # SE
            a[,,9] <- rbind(rep(0, da[2]), a[1:(da[1]-1), 1:(da[2]), 9]) # S
            apply(a, c(1, 2), sum)
      }

      diffuse <- function(n, p, i, rec = i, method = "pulse"){
            rec <- sort(rec)
            r <- terra::as.array(rast(n, nlyrs = length(rec), vals = 0))
            n <- matrix(n, nrow(n), byrow = T)
            n0 <- n
            if(0 %in% rec) r[,,1] <- n
            p <- terra::as.array(p)
            pb <- txtProgressBar(min = 0, max = i, initial = 0, style = 3)
            for(j in 1:i){
                  if(method == "pulse") n <- disperse(n, p)
                  if(method == "ratchet") n <- pmax(n, disperse(n, p))
                  # if(method == "feed") n <- n0 + disperse(n, p)
                  # if(method == "stream") n <- pmax(n0, disperse(n, p))

                  if(j %in% rec) r[,,match(j, rec)] <- n
                  setTxtProgressBar(pb, j+1)
            }
            close(pb)
            r
      }

      if(!between(timescale, 0, 1)) stop("'timescale' must be greater than 0 and less than or equal to 1.")
      t <- rw_max_step(rose) * timescale
      message("\titeration timestep: ~", signif(t, 3),
              "\n\tsimulation duration: ~", signif(t * iter, 3),
              "\n(these values are in hours, IF trans == 1 and wind_field units are m/s)")
      p <- rw_prob(rose, t)

      # starting distribution
      if(inherits(init, "matrix")){
            init <- raster::cellFromXY(rose, init)
            n <- rose[[1]]
            values(n) <- 0
            n[unique(init)] <- 1
      }else{
            n <- init
      }

      w <- diffuse(n, p, iter, record, method = mode)
      n <- rast(n, nlyrs = length(record), vals = w)
      names(n) <- paste0("iter", record)
      as_wind_walk(n, mode = mode, n_iter = iter, iter_length = t)
}