jdyen

# 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)]

jdyen

tf_expand <- function(x, repeats) {
  
  x <- tf$constant(x, dtype = tf$float64)
  repeats <- tf$constant(repeats, dtype = tf$int32)
  
  shape <- tf$reduce_sum(repeats)
  
  idx <- tf$concat(c(tf$constant(0L, dtype = tf$int32)[tf$newaxis],
                     tf$cumsum(repeats[0L:(tf$size(repeats) - 1L)])),
                   axis = 0L)

jdyen

# extract a node from a greta array
get_node <- greta:::get_node

# convert node to greta array
as.greta_array <- greta:::as.greta_array

# extract unique values without keeping an extra data node around
extract_unique <- function(params) {
  
  node <- get_node(params)