Minor change task 3

R/fit_task3.R

@@ -122,7 +122,7 @@ fit_two_pop_model <- function(
   times_plot <- times %>%
     dplyr::group_by(.data$par) %>%
     dplyr::mutate(q_low = stats::quantile(x, .05), q_high = stats::quantile(x, .95)) %>%
-    dplyr::filter(x >= .data$q_low, x <= .data$q_high) %>%
+    #dplyr::filter(x >= .data$q_low, x <= .data$q_high) %>%
     ggplot2::ggplot(mapping = ggplot2::aes(x=.data$x, col=.data$par)) +
     ggplot2::geom_density() +
     ggplot2::labs(x = "Time", y="Density") +
@@ -133,8 +133,10 @@ fit_two_pop_model <- function(
   nr_first_obs <- best_fit$draws(format = "df", variables = "nr")[,1] %>% unlist() %>% as.numeric()
   ns_first_obs <- best_fit$draws(format = "df", variables = "ns")[,1] %>% unlist() %>% as.numeric()
 
-  nr_first_obs <- nr_first_obs[t0_r_draws <= 0]
-  ns_first_obs <- ns_first_obs[t0_r_draws <= 0]
+  nr_first_obs[nr_first_obs == 1e-9] = 0
+
+  #nr_first_obs <- nr_first_obs[t0_r_draws <= 0]
+  #ns_first_obs <- ns_first_obs[t0_r_draws <= 0]
 
   f_r = nr_first_obs / (nr_first_obs + ns_first_obs)
   fr_plot <- dplyr::tibble(x = f_r) %>%
@@ -157,7 +159,7 @@ fit_two_pop_model <- function(
   rates_plot <- rates %>%
     dplyr::group_by(.data$par) %>%
     dplyr::mutate(q_low = stats::quantile(x, .05), q_high = stats::quantile(x, .95)) %>%
-    dplyr::filter(x >= .data$q_low, x <= .data$q_high) %>%
+    #dplyr::filter(x >= .data$q_low, x <= .data$q_high) %>%
     ggplot2::ggplot(mapping = ggplot2::aes(x=.data$x, col=.data$par)) +
     ggplot2::geom_density() +
     ggplot2::geom_vline(xintercept = 0, linetype = "dashed") +
@@ -194,17 +196,20 @@ fit_two_pop_data <- function(x, factor_size, variational, chains, iter, cores) {
   #model <- get_model(model_name = "two_pop")
 
   m1 <- get_model("two_pop_single")
-  m2 <- get_model("two_pop_both")
+  m2 <- get_model("two_pop_pre")
+  m3 <- get_model("two_pop_post")
 
   tmp <- utils::capture.output(f1 <- m1$sample(input_data, parallel_chains = chains, iter_warmup = iter, iter_sampling = iter, chains = chains))
   tmp <- utils::capture.output(f2 <- m2$sample(input_data, parallel_chains = chains, iter_warmup = iter, iter_sampling = iter, chains = chains))
+  tmp <- utils::capture.output(f3 <- m3$sample(input_data, parallel_chains = chains, iter_warmup = iter, iter_sampling = iter, chains = chains))
 
-  fits <- list(f1, f2)
+  fits <- list(f1, f2, f3)
 
   loo1 <- f1$loo()
   loo2 <- f2$loo()
+  loo3 <- f3$loo()
 
-  loos <- loo::loo_compare(list(loo1, loo2))
+  loos <- loo::loo_compare(list(loo1, loo2, loo3))
   fit_model <- fits[[as.numeric(stringr::str_replace(rownames(loos)[1], "model", ""))]]
 
   sampling <- "mcmc"

R/getter.R

@@ -10,10 +10,10 @@ get_model <- function(model_name) {
     "piecewise_changepoints" = "piecewise_linear_regression.stan",
     #"two_pop" = "two_population.stan",
     "two_pop_both" = 'two_pop_both_v2.stan',
-    "two_pop_single" = 'two_pop_single.stan'
+    "two_pop_single" = 'two_pop_single.stan',
+    "two_pop_pre" = 'two_pop_pre.stan',
+    "two_pop_post" = 'two_pop_post.stan'
     #"pw_lin_fixed_b" = "pw_linear_b_fixed.stan",
-
-
   )
 
   if (!(model_name) %in% names(all_paths)) stop("model_name not recognized")

inst/cmdstan/two_pop_post.stan

@@ -0,0 +1,68 @@
+data {
+  int<lower=1> S; // Number of steps
+  array[S] int<lower=0> N; // observations
+  array[S] real T;         // observations
+}
+
+parameters {
+  real<lower=0> rho_r;        // Parameter rho_r (rate for recoverN)
+  real<lower=0> rho_s;        // Parameter rho_s (rate for signal decaN)
+  real<lower=T[1]> t0_r;                   // Parameter t_r (time shift)
+  real<lower=T[1]> t_end;
+  // real<lower=0, upper=1.5> f_s;
+}
+
+model {
+  vector[S] mu;               // Expected values for N given x
+  vector[S] ns;
+  vector[S] nr;
+
+  // Priors
+  // n0 ~ normal(N[1], 0.1 * N[1]);
+  rho_r ~ normal(0, 1);       // Prior for rho_r
+  rho_s ~ normal(0, 1);       // Prior for rho_s
+  t0_r ~ normal(T[1], 1);         // Prior for t_r
+  t_end ~ normal(T[S], 1);
+
+  for (i in 1:S) {
+    if (T[i] >= t_end) {
+      ns[i] = 1e-9;
+    } else {
+      ns[i] = exp(-rho_s * (T[i] - t_end));
+    }
+
+    if (T[i] >= t0_r) {
+      nr[i] = exp(rho_r * (T[i] - t0_r));
+    } else {
+      nr[i] = 1e-9;
+    }
+
+    mu[i] = nr[i] + ns[i];
+  }
+
+  // Likelihood (assuming normallN distributed noise)
+  N ~ poisson(mu);
+}
+
+generated quantities {
+  vector[S] log_lik;             // Log-likelihood for each observation
+  vector[S] ns;
+  vector[S] nr;
+  vector[S] yrep;               // Expected values for N given x
+
+  for (i in 1:S) {
+    if (T[i] >= t_end) {
+      ns[i] = 1e-9;
+    } else {
+      ns[i] = exp(-rho_s * (T[i] - t_end));
+    }
+
+    if (T[i] >= t0_r) {
+      nr[i] = exp(rho_r * (T[i] - t0_r));
+    } else {
+      nr[i] = 1e-9;
+    }
+    yrep[i] = nr[i] + ns[i];
+    log_lik[i] = poisson_lpmf(N[i] | yrep[i]); // Log-likelihood calculation
+  }
+}

inst/cmdstan/two_pop_pre.stan

@@ -0,0 +1,69 @@
+data {
+  int<lower=1> S; // Number of steps
+  array[S] int<lower=0> N; // observations
+  array[S] real T;         // observations
+}
+
+parameters {
+  real<lower=0> rho_r;        // Parameter rho_r (rate for recoverN)
+  real<lower=0> rho_s;        // Parameter rho_s (rate for signal decaN)
+  real<upper=T[1]> t0_r;                   // Parameter t_r (time shift)
+  real<lower=T[1]> t_end;
+  // real<lower=0, upper=1.5> f_s;
+}
+
+
+model {
+  vector[S] mu;               // Expected values for N given x
+  vector[S] ns;
+  vector[S] nr;
+
+  // Priors
+  // n0 ~ normal(N[1], 0.1 * N[1]);
+  rho_r ~ normal(0, 1);       // Prior for rho_r
+  rho_s ~ normal(0, 1);       // Prior for rho_s
+  t0_r ~ normal(T[1], 1);         // Prior for t_r
+  t_end ~ normal(T[S], 1);
+
+  for (i in 1:S) {
+    if (T[i] >= t_end) {
+      ns[i] = 1e-9;
+    } else {
+      ns[i] = exp(-rho_s * (T[i] - t_end));
+    }
+
+    if (T[i] >= t0_r) {
+      nr[i] = exp(rho_r * (T[i] - t0_r));
+    } else {
+      nr[i] = 1e-9;
+    }
+
+    mu[i] = nr[i] + ns[i];
+  }
+
+  // Likelihood (assuming normallN distributed noise)
+  N ~ poisson(mu);
+}
+
+generated quantities {
+  vector[S] log_lik;             // Log-likelihood for each observation
+  vector[S] ns;
+  vector[S] nr;
+  vector[S] yrep;               // Expected values for N given x
+
+  for (i in 1:S) {
+    if (T[i] >= t_end) {
+      ns[i] = 1e-9;
+    } else {
+      ns[i] = exp(-rho_s * (T[i] - t_end));
+    }
+
+    if (T[i] >= t0_r) {
+      nr[i] = exp(rho_r * (T[i] - t0_r));
+    } else {
+      nr[i] = 1e-9;
+    }
+    yrep[i] = nr[i] + ns[i];
+    log_lik[i] = poisson_lpmf(N[i] | yrep[i]); // Log-likelihood calculation
+  }
+}