test-iterate_matrix.R

# simulate from a probabilistic 4-stage transition matrix model

library(greta.dynamics)

k <- 4
n <- 10
survival_fixed <- rep(0.5, k)
stasis_fixed <- c(rep(0.25, k - 1), 1)
stay_fixed <- survival_fixed * stasis_fixed
progress_fixed <- (survival_fixed * (1 - stay_fixed))[1:(k - 1)]
recruit_fixed <- c(1, 5)

# tf params
survival <- uniform(0, 1, dim = k)
stasis <- c(uniform(0, 1, dim = k - 1), 1)
stay <- survival * stasis
progress <- (survival * (1 - stay))[1:(k - 1)]
recruit <- exponential(c(3, 5))

# combine into a matrix
tmat <- zeros(k, k)
diag(tmat) <- stay
progress_idx <- row(tmat) - col(tmat) == 1
tmat[progress_idx] <- progress
tmat[1, k - (1:0)] <- recruit

# expand along multiple dims
tmats <- zeros(n, k, k)
for (i in seq_len(n))
  tmats[i, , ] <- tmat

# make a non-tf matrix
tmat_real <- matrix(0, k, k)
diag(tmat_real) <-stay_fixed
tmat_real[progress_idx] <- progress_fixed
tmat_real[1, k - (1:0)] <- recruit_fixed

# expand non-tf matrix
tmats_real <- array(0, dim = c(n, k, k))
for (i in seq_len(n))
  tmats_real[i, , ] <- tmat_real

# function to project non-tf matrices
iterate_matrix_real <- function(matrix,
                                initial_state = rep(1, ncol(matrix)),
                                niter = 100) {
  matrix_dim <- dim(matrix)
  state_dim <- dim(initial_state)
  state <- initial_state
  matrix_n_dim <- length(matrix_dim)
  state_n_dim <- length(state_dim)
  matrix_multisite <- matrix_n_dim == 3
  state_multisite <- state_n_dim == 3
  multisite <- matrix_multisite | state_multisite
  if (multisite) {
    if (!state_multisite) {
      n <- matrix_dim[1]
      state_list <- replicate(n, state, simplify = FALSE)
      state <- t(do.call(cbind, state_list))
      dim(state) <- c(dim(state), 1)
      state_dim <- dim(state)
      state_n_dim <- 3
    }
    if (!matrix_multisite) {
      n <- state_dim[1]
      dim(matrix) <- c(dim(matrix), 1)
      matrix_list <- replicate(n, matrix, simplify = FALSE)
      matrix <- do.call(abind, c(matrix_list, list(along = 1)))
      matrix_dim <- dim(matrix)
      matrix_n_dim <- 3
    }
    if (matrix_multisite & state_multisite) {
      n <- matrix_dim[1]
    }
    states <- array(0, dim = c(n, matrix_dim[2], niter + 1))
    states[, , 1] <- state
    for (i in seq_len(n)) {
      for (j in seq_len(niter)) {
        states[i, , j + 1] <- matrix[i, , ] %*% states[i, , j]
      }
    }
    states <- states[, , -1]
  } else {
    n <- 1
    states <- matrix(0, nrow = matrix_dim[2], ncol = (niter + 1))
    states[, 1] <- state
    for (i in seq_len(niter)) {
      states[, i + 1] <- matrix %*% states[, i]
    }
    states <- states[, -1]
  }
  states
}

# iterate matrix 20 times
iterations_tf <- iterate_matrix(tmat, niter = 20)
iterations_tf_multisite <- iterate_matrix(tmats, niter = 20)
iterations_real <- iterate_matrix_real(tmat_real, niter = 20)
iterations_real_multisite <- iterate_matrix_real(tmats_real, niter = 20)

# calculate real values of tf states
tf_allstates <- calculate(iterations_tf$all_states,
                          list(survival = survival_fixed,
                               stasis = stasis_fixed,
                               recruit = recruit_fixed))
tf_multisite_allstates <- calculate(iterations_tf_multisite$all_states,
                                    list(survival = survival_fixed,
                                         stasis = stasis_fixed,
                                         recruit = recruit_fixed)) 

# check equal
all.equal(tf_allstates, iterations_real)
all.equal(tf_multisite_allstates, iterations_real_multisite)

# multisite version works if dims 2 and 3 of tmats are transposed
tmats_tranposed <- zeros(n, k, k)
for (i in seq_len(n))
  tmats_tranposed[i, , ] <- t(tmat)

iterations_tf_multisite_transposed <- iterate_matrix(tmats_tranposed, niter = 20)
tf_multisite_allstates_transposed <- calculate(iterations_tf_multisite_transposed$all_states,
                                               list(survival = survival_fixed,
                                                    stasis = stasis_fixed,
                                                    recruit = recruit_fixed)) 

# check equal
all.equal(tf_multisite_allstates_transposed, iterations_real_multisite)