Feat/zinb simulation and toy model

1_run_simulation.R

@@ -57,9 +57,18 @@ sim_convergence <- tibble()
 # loop through all sample sizes included
 for (sample_size in c(200, 1000, 10000)) {
   message(paste("Running with sample size", sample_size))
-  # loop through all gps model specifications
-  for (gps_mod in 1:4) {
+  # loop through all gps model specifications (now includes 5 for ZINB)
+  for (gps_mod in 1:5) {
     message(paste("Running with gps model", gps_mod))
+
+    # Specific parameters for ZINB (gps_mod == 5).
+    # These are currently set to the defaults in sim_data_generate.
+    # They can be customized here if needed for specific simulation runs.
+    # e.g., zinb_params <- list(zinb_mu = 10, zinb_theta = 2, zinb_pi = 0.25)
+    # And then pass these to sim_data_generate:
+    # sim_data_generate(..., zinb_mu = zinb_params$zinb_mu, ...)
+    # For now, we rely on the defaults in sim_data_generate.
+
     # Loop through two different outcome relationships
     for (outcome_interaction in c("T", "F")) {
 

README.md

@@ -12,6 +12,41 @@ The main code scripts are as follows:
 * `/figures/` saves figures used in publication
 * `/markdown_files/` contains scripts used for exploratory data analysis and making additional plots
 
+## Toy Model for Understanding Simulations
+A new script `toy_model_simulation.R` has been added to help users understand the simulation process.
+* **Purpose**: This script runs a single, simplified simulation scenario. It is heavily commented to explain each step, from data generation to model fitting and results evaluation. It's an excellent starting point for new users to grasp the core mechanics before diving into the more complex, large-scale simulations.
+* **How to Run**:
+    1. Open `toy_model_simulation.R` in RStudio or your preferred R environment.
+    2. Modify the `REPO_DIR` variable at the top of the script if your working directory is not the root of the repository.
+    3. Adjust simulation parameters (e.g., `SAMPLE_SIZE`, `GPS_MODEL_SPEC`, `EXPOSURE_RESPONSE`) as desired to explore different scenarios.
+    4. Run the script. It will print output to the console and can be configured to save a plot of the estimated exposure-response curve.
+* **Key Features**:
+    * Uses the same core functions (`sim_data_generate`, `metrics_from_data`) as the main simulation.
+    * Focuses on clarity and step-by-step execution.
+    * Allows easy experimentation with different simulation settings.
+
+## Simulation Details
+
+### Treatment Generation (`gps_mod`)
+The `gps_mod` parameter in `sim_data_generate` (within `functions/simulation_functions.R`) controls how the exposure (treatment) is generated based on covariates. The following options are available:
+
+*   `gps_mod = 1`: Linear relationship with normal error. Exposure ~ β0 + βX + N(0, σ^2).
+*   `gps_mod = 2`: Linear relationship with t-distributed error. Exposure ~ β0 + βX + t(df).
+*   `gps_mod = 3`: Non-linear relationship (quadratic term) with normal error. Exposure ~ β0 + βX + β_k*X_k^2 + N(0, σ^2).
+*   `gps_mod = 4`: Non-linear relationship (quadratic and interaction terms) with normal error. Exposure ~ β0 + βX + β_k*X_k^2 + β_ij*X_i*X_j + N(0, σ^2).
+*   `gps_mod = 5`: Zero-Inflated Negative Binomial (ZINB) distribution.
+    *   This option generates count data with excess zeros, which can be useful for modeling treatments that are often zero but positive otherwise.
+    *   The mean of the Negative Binomial component (μ_NB) is modeled as a function of covariates: `μ_NB = exp(log(zinb_mu_base) + common_linear_term / 2)`.
+    *   The ZINB generation uses the following parameters (which can be passed to `sim_data_generate`):
+        *   `zinb_mu` (Default: 5): The base mean for the Negative Binomial component. The actual mean for each observation will vary based on its covariates.
+        *   `zinb_theta` (Default: 1): The dispersion parameter (size) for the Negative Binomial component. Higher values mean less dispersion.
+        *   `zinb_pi` (Default: 0.3): The zero-inflation probability. This is the probability that an observation is an "excess" zero, regardless of the NB component.
+
+### Outcome Generation
+The outcome `Y` is generated based on the exposure, confounders, and specified error distribution. Key parameters include:
+* `exposure_response_relationship`: Defines the true functional form between exposure and outcome (e.g., "linear", "sublinear", "threshold").
+* `outcome_interaction`: A boolean (`TRUE`/`FALSE`) indicating whether there's an interaction term between the (transformed) exposure and some confounders in the true outcome model.
+
 ## Data application
 Code related to fitting model to Medicare database is found in `/data_application/`
 

functions/simulation_functions.R

@@ -2,6 +2,24 @@
 library(dplyr)
 library(purrr)
 library(MASS)
+# library(pscl) # For ZINB, if using pscl::rzinb. Or use custom function.
+# library(countreg) # For ZINB using gamlss::rzinb
+
+# Function to generate Zero-Inflated Negative Binomial (ZINB) random variables
+# n: number of observations
+# mu: mean of the negative binomial part
+# theta: dispersion parameter of the negative binomial part (size)
+# pi: zero-inflation probability
+rzinb_custom <- function(n, mu, theta, pi) {
+  # Generate from Negative Binomial
+  nb_counts <- rnbinom(n, size = theta, mu = mu)
+  # Generate zero-inflation component
+  is_zero <- rbinom(n, size = 1, prob = pi)
+  # Combine: if is_zero is 1, then count is 0, otherwise it's from NB
+  counts <- ifelse(is_zero == 1, 0, nb_counts)
+  return(counts)
+}
+
 
 # Function to generate synthetic data for simulations
 # Sample size is sample for simulation 
@@ -16,11 +34,14 @@ sim_data_generate <- function(sample_size = 1000,
                               outcome_interaction = FALSE,
                               outcome_sd = 10, 
                               data_application = FALSE,
-                              data_app_sample = NULL) {
+                              data_app_sample = NULL,
+                              zinb_mu = 5,      # Default mean for ZINB
+                              zinb_theta = 1,   # Default dispersion for ZINB
+                              zinb_pi = 0.3) {  # Default zero-inflation probability for ZINB
 
   # Make sure input parameters are correct
-  if (!gps_mod %in% 1:4) {
-    stop("Invalid value for gps_mod. It should be in 1:4.")
+  if (!gps_mod %in% 1:5) { # Updated to include 5 for ZINB
+    stop("Invalid value for gps_mod. It should be in 1:5.")
   }
 
   if (!exposure_response_relationship %in% c("linear", "sublinear", "threshold")) {
@@ -45,21 +66,50 @@ sim_data_generate <- function(sample_size = 1000,
       rename_with(~ paste0("cf", 1:6))
 
     # Generate appropriate exposure based on gps_mod
-    exposure_df <- 
-      tibble(confounders_large) %>%
-      mutate(
-        exposure = case_when(
-          gps_mod == 1 ~ 9 * cov_function(confounders_large) + 18 + rnorm(sample_size, 0, sqrt(10)),
-          gps_mod == 2 ~ 9 * cov_function(confounders_large) + 18 + sqrt(5) * rt(sample_size, df = 3),
-          gps_mod == 3 ~ 9 * cov_function(confounders_large) + 15 + 2 * (confounders_large[, "cf3"])^2 + rnorm(sample_size, 0, sqrt(10)),
-          gps_mod == 4 ~ 9 * cov_function(confounders_large) + 2 * confounders_large[, "cf3"]^2 + 2 * confounders_large[, "cf1"] * confounders_large[, "cf4"] + 15 + rnorm(sample_size, 0, sqrt(10))
-        )
-      ) %>%
-      filter(exposure > 0) %>%
-      slice_sample(n = sample_size, replace = FALSE) %>%
-      dplyr::select(exposure, starts_with("cf"))
+    # Note: Using 2*sample_size for initial generation, then filtering and sampling.
 
-    # Now seperate out exposure and confounders, name cf for brevity
+    common_linear_term <- cov_function(confounders_large) # This is applied to 2*sample_size rows
+
+    if (gps_mod %in% 1:4) {
+      exposure_values_raw <-
+        case_when(
+          # Ensure rhs of ~ matches length of confounders_large (i.e. 2*sample_size)
+          gps_mod == 1 ~ 9 * common_linear_term + 18 + rnorm(2 * sample_size, 0, sqrt(10)),
+          gps_mod == 2 ~ 9 * common_linear_term + 18 + sqrt(5) * rt(2 * sample_size, df = 3),
+          gps_mod == 3 ~ 9 * common_linear_term + 15 + 2 * (confounders_large[, "cf3"])^2 + rnorm(2 * sample_size, 0, sqrt(10)),
+          gps_mod == 4 ~ 9 * common_linear_term + 2 * confounders_large[, "cf3"]^2 + 2 * confounders_large[, "cf1"] * confounders_large[, "cf4"] + 15 + rnorm(2 * sample_size, 0, sqrt(10))
+        )
+
+      exposure_df <-
+        tibble(confounders_large, exposure_raw = exposure_values_raw) %>%
+        filter(exposure_raw > 0) %>%
+        slice_sample(n = sample_size, replace = ifelse(nrow(.) < sample_size, TRUE, FALSE)) %>%
+        rename(exposure = exposure_raw) %>%
+        dplyr::select(exposure, starts_with("cf"))
+
+    } else if (gps_mod == 5) { # ZINB model
+      log_mu_base <- log(zinb_mu) # Use the zinb_mu parameter from function arguments
+
+      # common_linear_term is already calculated based on confounders_large (2*sample_size)
+      log_mu_nb_values <- log_mu_base + (common_linear_term / 2)
+      mu_nb_values <- exp(log_mu_nb_values)
+
+      mu_nb_values <- pmin(mu_nb_values, 50)
+      mu_nb_values[mu_nb_values <= 0.01] <- 0.01
+
+      exposure_values_raw <- rzinb_custom(n = 2 * sample_size, # Matches length of mu_nb_values
+                                          mu = mu_nb_values,
+                                          theta = zinb_theta, # zinb_theta from function arguments
+                                          pi = zinb_pi)       # zinb_pi from function arguments
+
+      exposure_df <-
+        tibble(confounders_large, exposure_raw = exposure_values_raw) %>%
+        slice_sample(n = sample_size, replace = FALSE) %>%
+        rename(exposure = exposure_raw) %>%
+        dplyr::select(exposure, starts_with("cf"))
+    }
+
+    # Now separate out exposure and confounders, name cf for brevity
     exposure <- exposure_df$exposure
     cf <- as.matrix(exposure_df %>% dplyr::select(-exposure))
   } else {

nb_model_simulation.R

@@ -0,0 +1,137 @@
+# Negative Binomial (NB) Model Simulation Script
+#
+# This script provides an example of how to run a single simulation scenario
+# with a Negative Binomial (NB) distributed treatment.
+# It utilizes the existing ZINB functionality (gps_mod = 5) by setting
+# the zero-inflation probability (zinb_pi) to 0.
+#
+# Its purpose is to illustrate how to achieve NB treatment generation
+# using the functions in functions/simulation_functions.R.
+
+# Load necessary libraries
+# Ensure they are installed in your R environment.
+# install.packages(c("purrr", "dplyr", "tidyr", "ggplot2", "MASS", "parallel", "xgboost", "SuperLearner", "WeightIt", "chngpt", "cobalt", "CausalGPS"))
+.libPaths(new = c("~/R/x86_64-pc-linux-gnu-library/4.2", .libPaths())) # Specific to FASSE cluster
+
+library(purrr)
+library(dplyr)
+library(tidyr)
+library(ggplot2)
+library(MASS)
+# library(parallel) # Not strictly needed for a single run
+library(xgboost)
+# library(SuperLearner)
+library(WeightIt)
+library(chngpt)
+library(cobalt)
+library(CausalGPS)
+
+# --- Configuration ---
+# Define the directory where the repository is located.
+# Adjust this path to match your local setup.
+# REPO_DIR <- "~/Desktop/Francesca_research/Simulation_studies/" # Example local path
+REPO_DIR <- "." # Assumes running from the root of the repository
+
+# Source the simulation functions
+source(paste0(REPO_DIR, "/functions/simulation_functions.R"))
+
+# --- Simulation Parameters for NB Treatment ---
+set.seed(456) # Use a different seed for variety
+
+SAMPLE_SIZE <- 300 # A different sample size for this example
+GPS_MODEL_SPEC <- 5  # Use gps_mod = 5, which is the ZINB mechanism
+EXPOSURE_RESPONSE <- "sublinear" # Example: sublinear relationship
+OUTCOME_INTERACTION_PRESENT <- FALSE
+OUTCOME_SD <- 10
+
+# Parameters for Negative Binomial (achieved via ZINB settings)
+NB_MEAN_PARAM <- 8       # This will be used for `zinb_mu` (base mean for NB component)
+NB_DISPERSION_PARAM <- 1.5 # This will be used for `zinb_theta` (dispersion for NB component)
+NB_ZERO_INFLATION_PROB <- 0 # CRITICAL: Set to 0 for pure Negative Binomial
+
+message(paste("--- Negative Binomial (NB) Treatment Simulation ---"))
+message(paste("Simulating NB by using ZINB mechanism (gps_mod=5) with zero-inflation probability = 0."))
+
+# --- 1. Data Generation for NB Treatment ---
+message(paste("\nGenerating synthetic data with NB treatment (Sample Size:", SAMPLE_SIZE,
+              ", Base NB Mean:", NB_MEAN_PARAM, ", NB Dispersion:", NB_DISPERSION_PARAM, ")"))
+
+sim_data_nb <- sim_data_generate(
+  sample_size = SAMPLE_SIZE,
+  gps_mod = GPS_MODEL_SPEC, # Should be 5
+  exposure_response_relationship = EXPOSURE_RESPONSE,
+  outcome_interaction = OUTCOME_INTERACTION_PRESENT,
+  outcome_sd = OUTCOME_SD,
+  data_application = FALSE,
+  zinb_mu = NB_MEAN_PARAM,
+  zinb_theta = NB_DISPERSION_PARAM,
+  zinb_pi = NB_ZERO_INFLATION_PROB # Key for NB
+)
+
+# Explore the generated data (optional)
+message("\nGenerated NB data summary:")
+print(head(sim_data_nb))
+message(paste("Number of rows in generated data:", nrow(sim_data_nb)))
+summary(sim_data_nb$exposure)
+message(paste("Proportion of actual zeros in exposure:", mean(sim_data_nb$exposure == 0)))
+# For NB (pi=0), any zeros are from the NB distribution itself, not from inflation.
+# If mu is high and theta not too small, actual zeros from NB might be rare.
+
+# --- 2. Model Fitting and Metric Calculation ---
+message("\nFitting models and calculating metrics for NB data...")
+
+metrics_and_predictions_nb <- metrics_from_data(
+  sim_data = sim_data_nb,
+  exposure_response_relationship = EXPOSURE_RESPONSE,
+  outcome_interaction = OUTCOME_INTERACTION_PRESENT
+)
+
+model_metrics_nb <- metrics_and_predictions_nb$metrics
+model_predictions_nb <- metrics_and_predictions_nb$predictions
+correlation_table_nb <- metrics_and_predictions_nb$cor_table
+convergence_status_nb <- metrics_and_predictions_nb$convergence_info
+
+# --- 3. Reviewing Results ---
+message("\n--- NB Model Simulation Results ---")
+
+message("\nConvergence Status of Weighting Methods (NB):")
+print(convergence_status_nb)
+
+message("\nCovariate Balance (Absolute Correlation with Exposure - NB):")
+print(head(correlation_table_nb %>% filter(method == "ent")))
+
+message("\nModel Performance Metrics (Bias and MSE - NB):")
+print(head(model_metrics_nb %>% filter(model == "gam_model", exposure < 5))) # Example for GAM
+
+message("\nPredicted Exposure-Response Curves (NB - First few points):")
+print(head(model_predictions_nb %>% dplyr::select(exposure, true_fit, gam_model, linear_model)))
+
+# --- 4. Visualization (Optional) ---
+plot_data_nb <- model_predictions_nb %>%
+  select(exposure, true_fit, gam_model, linear_model, ent_gam) %>%
+  pivot_longer(cols = -c(exposure, true_fit), names_to = "model_type", values_to = "predicted_outcome")
+
+erc_plot_nb <- ggplot(plot_data_nb, aes(x = exposure)) +
+  geom_line(aes(y = predicted_outcome, color = model_type, linetype = "Estimated ERC")) +
+  geom_line(aes(y = true_fit, linetype = "True ERC"), color = "black") +
+  labs(
+    title = "NB Model: Estimated vs. True Exposure-Response Curve",
+    subtitle = paste("Scenario: ", EXPOSURE_RESPONSE, " ER, N=", SAMPLE_SIZE, ", Base NB Mean=", NB_MEAN_PARAM, sep=""),
+    x = "Exposure Level",
+    y = "Outcome",
+    color = "Model Type",
+    linetype = "Curve Type"
+  ) +
+  scale_linetype_manual(values = c("Estimated ERC" = "solid", "True ERC" = "dashed")) +
+  theme_minimal() +
+  theme(legend.position = "top")
+
+# print(erc_plot_nb) # Display the plot
+# To save:
+# dir.create(paste0(REPO_DIR, "/figures"), showWarnings = FALSE) # Ensure directory exists
+# ggsave(paste0(REPO_DIR, "/figures/nb_model_erc_plot.png"), erc_plot_nb, width = 8, height = 6)
+# message(paste0("\nPlot saved to ", REPO_DIR, "/figures/nb_model_erc_plot.png (if figures directory exists)"))
+
+message("\nNB model simulation finished.")
+message("This script uses gps_mod=5 with zinb_pi=0 to simulate Negative Binomial treatment.")
+message("Modify parameters as needed to explore other NB scenarios.")