testing/aux_setup2_plot.R

In vsritter/sim.etm: What the Package Does (One Line, Title Case)

rm(list = ls())
set.seed(1047556)
library(tidyverse)

# -----------------------------------------------------------------------------
# ETM
# -----------------------------------------------------------------------------
Q1 <- array(NA, dim = c(3, 3, 2))
Q1[1,2,] <- c(1, 1)
Q1[1,3,] <- c(1, .8)
Q1[2,3,] <- c(1, 1.5)

Q2 <- array(NA, dim = c(3, 3, 2))
Q2[1,2,] <- c(.5, 1)
Q2[1,3,] <- c(.5, .8)
Q2[2,3,] <- c(.5, 1.5)


# Pop strata sizes (30/70)
data1 <- simulate_process(N = 6000, gen = Q1)
data2 <- simulate_process(N = 14000, gen = Q2)


dt_etm1 <- data1 %>%
  arrange(id, time) %>%
  group_by(id) %>%
  mutate(from = paste0("T", state),
         to = paste0("T", c(state[-1], NA)),
         entry = time,
         exit = c(time[-1], NA)) %>%
  filter(!is.na(exit)) %>%
  select(id, from, to, entry, exit) %>%
  ungroup()

dt_etm2 <- data2 %>%
  arrange(id, time) %>%
  group_by(id) %>%
  mutate(from = paste0("T", state),
         to = paste0("T", c(state[-1], NA)),
         entry = time,
         exit = c(time[-1], NA)) %>%
  filter(!is.na(exit)) %>%
  select(id, from, to, entry, exit) %>%
  ungroup()


dt_etm <- dt_etm1 %>%
  bind_rows(dt_etm2 %>% mutate(id = id + max(dt_etm1$id)))

# Sample strata sizes (80/20)
dt_etm_s <- dt_etm1 %>%
  bind_rows(filter(dt_etm2, id <= 1500) %>%
              mutate(id = id + max(dt_etm1$id)))


tra <- apply(Q1+Q2, 1:2, function(x) sum(!is.na(x)) > 0)
n_states <- nrow(tra)
Q_index <- which(tra)
P_index <- sort(c(seq(1, n_states^2, n_states + 1), Q_index))
state_names <- paste0("T", 1:n_states)
trans <- paste0("T", rep(1:n_states, times = n_states), rep(1:n_states, each = n_states))
poss_trans <- trans[P_index]



etm_p1 <- etm::etm(dt_etm1, state_names, tra, "cens", s=0, covariance = F)
etm_p1 <- data.frame(p = as.vector(etm_p1$est),
                     t = rep(etm_p1$time, each = n_states^2),
                     trans = rep(trans, times = length(t)),
                     stringsAsFactors = F) %>%
  filter(trans %in% poss_trans)

etm_p2 <- etm::etm(dt_etm2, state_names, tra, "cens", s=0, covariance = F)
etm_p2 <- data.frame(p = as.vector(etm_p2$est),
                     t = rep(etm_p2$time, each = n_states^2),
                     trans = rep(trans, times = length(t)),
                     stringsAsFactors = F) %>%
  filter(trans %in% poss_trans)

etm_p <- etm::etm(dt_etm, state_names, tra, "cens", s=0, covariance = F)
etm_p <- data.frame(p = as.vector(etm_p$est),
                    t = rep(etm_p$time, each = n_states^2),
                    trans = rep(trans, times = length(t)),
                    stringsAsFactors = F) %>%
  filter(trans %in% poss_trans)

etm_p_s <- etm::etm(dt_etm_s, state_names, tra, "cens", s=0, covariance = F)
etm_p_s <- data.frame(p = as.vector(etm_p_s$est),
                    t = rep(etm_p_s$time, each = n_states^2),
                    trans = rep(trans, times = length(t)),
                    stringsAsFactors = F) %>%
  filter(trans %in% poss_trans)


etm_p %>%
  ggplot(aes(x = t, y = p)) +
  geom_step(color = "black") +
  geom_step(aes(x = t, y = p), color = "gray60", data = etm_p_s) +
  geom_step(aes(x = t, y = p), color = "red", data = etm_p1) +
  geom_step(aes(x = t, y = p), color = "blue", data = etm_p2) +
  facet_wrap(~ trans)