Add parallel processing scripts for dengue forecasting and raking

- Implemented `04_rake_dengue_parallel.py` for parallel execution of dengue raking forecasts using Jobmon.
- Implemented `05_as_dengue_shifts_parallel.py` for parallel execution of age-sex shifts in dengue forecasts.
- Created `rake_dengue.py` script to handle the raking of dengue incidence predictions with detailed input/output path logic and regression calculations.
- Configured Jobmon workflows with appropriate compute resources and task templates for efficient processing.
- Added error handling and logging for workflow execution status.

Author: bcreiner

src/idd_forecast_mbp/02_data_prep/05_malaria_modeling_dataframe.py

@@ -254,7 +254,7 @@
 ]
 # Apply log transformations
 for col in covariates_to_log_transform:
-    as_md_df[f"log_{col}"] = np.log(as_md_df[col] + 1e-6)
+    as_md_df[f"log_{col}"] = np.log(as_md_df[col])
 
 as_md_modeling_df = as_md_df.merge(malaria_stage_3_df, on=["location_id", "year_id"], how="left")
 as_md_modeling_df = as_md_modeling_df[~as_md_modeling_df["A0_af"].isna()]

src/idd_forecast_mbp/02_data_prep/06_dengue_modeling_dataframe.py

@@ -172,7 +172,7 @@
     dengue_df[f"logit_{col}"] = np.log(clipped_values / (1 - clipped_values))
 
 aa_A0_dengue_df = dengue_df[(dengue_df["location_id"] == dengue_df["A0_location_id"]) & (dengue_df["year_id"] == 2022)].copy()
-aa_A0_dengue_df = aa_A0_dengue_df[aa_A0_dengue_df['aa_dengue_mort_count'] > 1].copy()
+aa_A0_dengue_df = aa_A0_dengue_df[aa_A0_dengue_df['aa_dengue_inc_count'] > 100].copy()
 A0_dengue_ids = aa_A0_dengue_df['A0_location_id'].unique()
 
 
@@ -189,7 +189,9 @@
 
 write_parquet(dengue_df, aa_ge3_dengue_stage_1_modeling_df_path)
 
-dengue_df = dengue_df[dengue_df['A0_location_id'].isin(A0_dengue_ids)].copy() 
+
+
+
 
 
 
@@ -198,15 +200,16 @@
 ### Prepares the final dataset for modeling by selecting relevant columns,
 ### merging to the age-sex-location-year level, and saving the final dataset.
 ###----------------------------------------------------------###
-dengue_stage_2_df = dengue_df.copy()
-# Drop any columns that have yn = 0
-dengue_stage_2_df = dengue_stage_2_df[dengue_stage_2_df["yn"] == 1].drop(columns=["yn"])
-# Subset down to level 5 locations
+dengue_stage_2_df = dengue_df[dengue_df['A0_location_id'].isin(A0_dengue_ids)].copy() 
+#
 dengue_stage_2_df = dengue_stage_2_df[dengue_stage_2_df["level"] == 5].drop(columns=["level"])
+
+
+
 # Create the A0_af factor variable
 dengue_stage_2_df["A0_location_id"] = dengue_stage_2_df["A0_location_id"].astype(int)
 dengue_stage_2_df['A0_af'] = 'A0_' + dengue_stage_2_df['A0_location_id'].astype(str)
-dengue_stage_2_df['A0_af'] = dengue_stage_2_df['A0_af'].astype('category')
+
 # dengue_stage_2_df = dengue_stage_2_df.drop(columns=['aa_dengue_inc_count', 'aa_dengue_inc_rate', 'aa_dengue_mort_count', 'aa_dengue_mort_rate', 'aa_dengue_cfr'])
 # Get the as data
 md_location_ids = dengue_stage_2_df["location_id"].unique().tolist()
@@ -216,20 +219,21 @@
                                          columns=as_merge_variables + ["dengue_mort_rate","dengue_inc_rate","dengue_mort_count","dengue_inc_count","population","aa_population"],
                                          filters=[year_filter, md_location_filter, age_filter, sex_filter])
 
+as_md_df = as_md_df[as_md_df['dengue_inc_rate'] > 0].copy()
 
 as_md_df["dengue_cfr"] = as_md_df["dengue_mort_rate"] / as_md_df["dengue_inc_rate"]
-as_md_df.loc[as_md_df["dengue_inc_rate"] == 0, "dengue_cfr"] = 0
 
 covariates_to_log_transform = [
-    "dengue_mort_rate",
     "dengue_inc_rate"
 ]
 for col in covariates_to_log_transform:
-    as_md_df[f"log_{col}"] = np.log(as_md_df[col] + 1e-6)
+    as_md_df[f"log_{col}"] = np.log(as_md_df[col])
+    
 
 covariates_to_logit_transform = ['dengue_cfr']
 for col in covariates_to_logit_transform:
-    clipped_values = as_md_df[col].clip(lower=0.001, upper=0.999)
+    clipped_values = as_md_df[col].clip(upper=0.99)
+    print(f"Range of {col}: {as_md_df[col].min()} to {as_md_df[col].max()}")
     as_md_df[f"logit_{col}"] = np.log(clipped_values / (1 - clipped_values))
 
 dengue_stage_2_df = dengue_stage_2_df.drop(columns=["population"])
@@ -239,7 +243,14 @@
 as_md_modeling_df = as_md_modeling_df[~(as_md_modeling_df["age_group_id"] == 2)]
 as_md_modeling_df["as_id"] = "a" + as_md_modeling_df["age_group_id"].astype(str) + "_s" + as_md_modeling_df["sex_id"].astype(str)
 
-write_parquet(as_md_modeling_df, as_md_dengue_modeling_df_path)
+columns_to_keep = as_merge_variables + ['A0_af', 'as_id', 'dengue_suitability', 'log_dengue_inc_rate', 'logit_urban_1km_threshold_300',
+                                        'log_gdppc_mean', 'logit_dengue_cfr', 'dengue_mort_rate']
+
+as_md_modeling_df = as_md_modeling_df[columns_to_keep]
+# Drop any columns that have yn = 0
+# dengue_stage_2_df = dengue_stage_2_df[dengue_stage_2_df["yn"] == 1].drop(columns=["yn"])
+
+
 
 cause_columns = list([col for col in as_md_modeling_df.columns if cause in col and "suit" not in col])
 base_md_modeling_df = as_md_modeling_df[(as_md_modeling_df['age_group_id'] == reference_age_group_id) & (as_md_modeling_df['sex_id'] == reference_sex_id)].copy()
@@ -253,6 +264,8 @@
                                                 on=aa_merge_variables,
                                                 how='left')
 
+as_md_modeling_df = as_md_modeling_df[as_md_modeling_df['dengue_mort_rate'] > 0]
 
+write_parquet(as_md_modeling_df, as_md_dengue_modeling_df_path)
 write_parquet(base_md_modeling_df, base_md_modeling_df_path)
 write_parquet(rest_md_modeling_df, rest_md_modeling_df_path)

src/idd_forecast_mbp/02_data_prep/09_forecasted_dengue_dataframes_parallel.py

@@ -55,10 +55,10 @@
 
 # Define the task template for processing each year batch
 task_template = tool.get_task_template(
-    template_name="hierarchy_generation",
+    template_name="as_dengue_forecasting_df_creation",
     default_cluster_name="slurm",
     default_compute_resources={
-        "memory": "50G",
+        "memory": "70G",
         "cores": 1,
         "runtime": "5m",
         "queue": "all.q",

src/idd_forecast_mbp/02_data_prep/forecasted_draw_specific_dengue_dataframes.py

@@ -83,13 +83,33 @@
 
 # Get the unique values of A0_location_id
 years = list(range(2022, 2023))
-year_filter = ('year_id', 'in', years)
+year_filter = ('year_id', '==', 2022)
 
+dengue_df = read_parquet_with_integer_ids(aa_full_cause_df_path_template,
+                                           filters=[year_filter, level_filter(hierarchy_df, start_level = 3)])
 
-base_md_modeling_df = read_parquet_with_integer_ids(base_md_modeling_df_path,
-                                                    filters = [year_filter])
+dengue_df = dengue_df[dengue_df['dengue_inc_count'] > 100].copy()
+dengue_df = dengue_df.rename(columns={'location_id':'A0_location_id'})
+A0_location_ids = dengue_df['A0_location_id'].unique()
 
-dengue_modeling_location_ids = base_md_modeling_df['location_id'].unique()
+age_filter = ('age_group_id', '==', reference_age_group_id)
+sex_filter = ('sex_id', '==', reference_sex_id)
+
+md_dengue_df = read_parquet_with_integer_ids(as_full_cause_df_path_template,
+                                           filters=[year_filter, age_filter, sex_filter, level_filter(hierarchy_df, start_level = 5)])
+
+md_dengue_df = md_dengue_df.merge(hierarchy_df[['location_id', 'A0_location_id']],
+                                  on='location_id', how='left')
+
+md_dengue_df = md_dengue_df[md_dengue_df['A0_location_id'].isin(A0_location_ids)].copy()
+
+md_dengue_df['base_log_dengue_inc_rate'] = np.log(md_dengue_df['dengue_inc_rate'])
+
+
+# base_md_modeling_df = read_parquet_with_integer_ids(base_md_modeling_df_path,
+#                                                     filters = [year_filter])
+
+dengue_modeling_location_ids = md_dengue_df['location_id'].unique()
 dengue_modeling_location_filter = ('location_id', 'in', dengue_modeling_location_ids)
 
 
@@ -112,14 +132,14 @@
 
 
 
-cause_columns = list([col for col in base_md_modeling_df.columns if cause in col and "suit" not in col])
-columns_to_keep = aa_merge_variables + ['base_log_dengue_inc_rate']
-base_md_modeling_df = base_md_modeling_df[columns_to_keep].copy()
-base_md_modeling_df = base_md_modeling_df.drop(columns=['year_id'])
+
+
+md_dengue_df = md_dengue_df[['location_id', 'base_log_dengue_inc_rate']].copy()
+
 
 # Merge in the dengue_stage_2_modeling_df
 forecast_by_draw_df = forecast_by_draw_df.merge(
-    base_md_modeling_df,
+    md_dengue_df,
     how="left",
     on=["location_id"]
 )
@@ -133,10 +153,26 @@
     # Create a new column with the log transformed value
     forecast_by_draw_df[f"log_{col}"] = np.log(forecast_by_draw_df[col] + 1e-6)
 
-as_md_dengue_modeling_df = read_parquet_with_integer_ids(as_md_dengue_modeling_df_path,
-                                                              filters = [year_filter, dengue_modeling_location_filter])
+as_sex_filter = ('sex_id',  '==', reference_sex_id)
+as_age_filter = ('age_group_id', '==', reference_age_group_id)
+
+
+as_md_dengue_modeling_df = read_parquet_with_integer_ids(as_full_cause_df_path_template,
+                                                            filters=[year_filter, dengue_modeling_location_filter, as_age_filter, as_sex_filter])
+
+as_md_dengue_modeling_df["as_id"] = "a" + as_md_dengue_modeling_df["age_group_id"].astype(str) + "_s" + as_md_dengue_modeling_df["sex_id"].astype(str)
+
+as_md_dengue_modeling_df["dengue_cfr"] = as_md_dengue_modeling_df["dengue_mort_rate"] / as_md_dengue_modeling_df["dengue_inc_rate"]
+
+covariates_to_logit_transform = ['dengue_cfr']
+for col in covariates_to_logit_transform:
+    clipped_values = as_md_dengue_modeling_df[col].clip(upper=0.99)
+    print(f"Range of {col}: {as_md_dengue_modeling_df[col].min()} to {as_md_dengue_modeling_df[col].max()}")
+    as_md_dengue_modeling_df[f"logit_{col}"] = np.log(clipped_values / (1 - clipped_values))
+
+
+
 
-cause_columns = list([col for col in as_md_dengue_modeling_df.columns if cause in col and "suit" not in col])
 columns_to_keep = as_merge_variables + ['logit_dengue_cfr', 'as_id']
 as_md_dengue_modeling_df = as_md_dengue_modeling_df[columns_to_keep].copy()
 as_md_dengue_modeling_df['year_to_rake'] = 2022
@@ -154,4 +190,9 @@
     draw=draw
 )
 
+columns_to_keep = as_merge_variables + ['logit_dengue_cfr', 'log_gdppc_mean', 'base_log_dengue_inc_rate', 'dengue_suitability', 'logit_urban_1km_threshold_300', 'as_id', 'A0_af']
+forecast_by_draw_df = forecast_by_draw_df[columns_to_keep].copy()
+
+
+
 write_parquet(forecast_by_draw_df, forecast_by_draw_df_path)

src/idd_forecast_mbp/02_data_prep/forecasted_draw_specific_malaria_dataframes.py

@@ -252,7 +252,7 @@ def generate_dah_scenarios(
 covariates_to_log_transform = [col for col in as_base_malaria_df.columns if 'rate' in col]
 for col in covariates_to_log_transform:
     # Create a new column with the log transformed value
-    as_base_malaria_df[f"log_{col}"] = np.log(as_base_malaria_df[col] + 1e-6)
+    as_base_malaria_df[f"log_{col}"] = np.log(as_base_malaria_df[col])
 
 as_base_malaria_df[f"logit_malaria_pfpr"] = np.log(0.999 * as_base_malaria_df["malaria_pfpr"] / (1 - 0.999 * as_base_malaria_df["malaria_pfpr"]))
 

src/idd_forecast_mbp/03_modeling/explore_new_final_models_dengue.r

@@ -0,0 +1,105 @@
+
+rm(list = ls())
+
+require(glue)
+require(mgcv)
+require(scam)
+require(arrow)
+require(data.table)
+
+
+"%ni%" <- Negate("%in%")
+"%nlike%" <- Negate("%like%")
+
+last_year <- 2022
+
+MODELING_DATA_PATH <- "/mnt/team/idd/pub/forecast-mbp/03-modeling_data"
+
+
+
+as_md_dengue_modeling_df_path = glue("{MODELING_DATA_PATH}/as_md_dengue_modeling_df.parquet")
+base_md_dengue_modeling_df_path = glue("{MODELING_DATA_PATH}/base_md_dengue_modeling_df.parquet")
+rest_md_dengue_modeling_df_path = glue("{MODELING_DATA_PATH}/rest_md_dengue_modeling_df.parquet")
+
+# Read in a parquet file
+#
+
+as_dengue_df = as.data.frame(arrow::read_parquet(as_md_dengue_modeling_df_path))
+as_dengue_df <- as_dengue_df[which(as_dengue_df$dengue_mort_rate > 0),]
+range(as_dengue_df$logit_dengue_cfr)
+as_dengue_df$A0_af <- as.factor(as_dengue_df$A0_af)
+as_dengue_df$as_id = as.factor(as_dengue_df$as_id)
+
+base_dengue_df = as.data.frame(arrow::read_parquet(base_md_dengue_modeling_df_path))
+base_dengue_df$A0_af <- as.factor(base_dengue_df$A0_af)
+base_dengue_df$as_id = as.factor(base_dengue_df$as_id)
+
+# load(glue("{MODELING_DATA_PATH}/2025_06_29_dengue_models.RData"))
+
+
+rest_dengue_df = as.data.frame(arrow::read_parquet(rest_md_dengue_modeling_df_path))
+rest_dengue_df$A0_af <- as.factor(rest_dengue_df$A0_af)
+rest_dengue_df$as_id = as.factor(rest_dengue_df$as_id)
+
+base_dengue_df$dengue_suit_fraction <- base_dengue_df$dengue_suitability / 365
+base_dengue_df$dengue_suit_fraction <- pmin(pmax(base_dengue_df$dengue_suit_fraction, 0.001), 0.999)
+base_dengue_df$logit_dengue_suitability <- log(base_dengue_df$dengue_suit_fraction / (1 - base_dengue_df$dengue_suit_fraction))
+
+
+mod_inc_base <- scam(base_log_dengue_inc_rate ~  s(logit_dengue_suitability, k = 6, bs = 'mpi') + 
+                       s(logit_urban_1km_threshold_300, k = 6, bs = 'mpi') + A0_af,
+                     data = base_dengue_df,
+                     optimizer = "efs",
+                     control = list(maxit = 300))  # Limit iterations 
+
+
+
+mod_cfr_all <- lm(logit_dengue_cfr ~ log_gdppc_mean + as_id + A0_af,
+                  data = as_dengue_df)
+
+
+write.csv(levels(as_dengue_df$as_id), file = glue("{MODELING_DATA_PATH}/as_id_levels.csv"), row.names = FALSE)
+write.csv(coef(mod_cfr_all), file = glue("{MODELING_DATA_PATH}/mod_cfr_all_coefficients.csv"), row.names = TRUE)
+
+strip_model <- function(model) {
+  model$model <- NULL
+  model$fitted.values <- NULL
+  model$residuals <- NULL
+  model$effects <- NULL
+  model$qr$qr <- NULL
+  model$linear.predictors <- NULL
+  model$weights <- NULL
+  model$prior.weights <- NULL
+  model$data <- NULL
+  model$family <- NULL
+  model$deviance <- NULL
+  model$aic <- NULL
+  model$null.deviance <- NULL
+  model$iter <- NULL
+  model$df.residual <- NULL
+  model$df.null <- NULL
+  model$y <- NULL
+  model$converged <- NULL
+  model$boundary <- NULL
+  
+  # Keep essential components for prediction
+  attr(model$terms, ".Environment") <- NULL
+  
+  return(model)
+}
+
+# Usage
+stripped_mod_cfr_all <- strip_model(mod_cfr_all)
+
+model_names <- c("mod_inc_base", "mod_cfr_all")
+
+save(list = model_names, file = glue("{MODELING_DATA_PATH}/2025_06_29_dengue_models.RData"))
+stripped_model_names <- c("mod_inc_base", "stripped_mod_cfr_all")
+save(list = model_names, file = glue("{MODELING_DATA_PATH}/2025_06_29_dengue_models_stripped.RData"))
+
+
+
+
+
+
+

src/idd_forecast_mbp/03_modeling/new_final_malaria_models.r

@@ -39,32 +39,38 @@ malaria_pfpr_mod <- scam(logit_malaria_pfpr ~ logit_malaria_suitability +
                            optimizer = "efs",      # Faster optimizer
                            control = list(maxit = 300))  # Limit iterations
 
+mod_df <- past_data[which(past_data$aa_malaria_mort_rate > 0),]
 mortality_scam_mod <- scam(log_aa_malaria_mort_rate ~ s(logit_malaria_pfpr, k = 10, bs = "mpi") + 
                              log_gdppc_mean + 
                              A0_af,
-                           data = past_data,
+                           data = mod_df,
                            optimizer = "efs",      # Faster optimizer
                            control = list(maxit = 300))  # Limit iterations
+
+mod_df <- past_data[which(past_data$aa_malaria_inc_rate > 0),]
 incidence_scam_mod <- scam(log_aa_malaria_inc_rate ~ s(logit_malaria_pfpr, k = 10, bs = "mpi") + 
                              log_gdppc_mean + A0_af,
-                           data = past_data,
+                           data = mod_df,
                            optimizer = "efs",      # Faster optimizer
                            control = list(maxit = 300))  # Limit iterations
 
+mod_df <- past_data[which(past_data$base_malaria_mort_rate > 0),]
 mortality_base_scam_mod <- scam(log_base_malaria_mort_rate ~ s(logit_malaria_pfpr, k = 10, bs = "mpi") + 
                              log_gdppc_mean + 
                              A0_af,
-                           data = past_data,
+                           data = mod_df,
                            optimizer = "efs",      # Faster optimizer
                            control = list(maxit = 300))  # Limit iterations
+
+mod_df <- past_data[which(past_data$base_malaria_inc_rate  > 0),]
 incidence_base_scam_mod <- scam(log_base_malaria_inc_rate ~ s(logit_malaria_pfpr, k = 10, bs = "mpi") + 
                              log_gdppc_mean + A0_af,
-                           data = past_data,
+                           data = mod_df,
                            optimizer = "efs",      # Faster optimizer
                            control = list(maxit = 300))  # Limit iterations
 
 
 model_names <- c("malaria_pfpr_mod", "mortality_scam_mod", "incidence_scam_mod", "mortality_base_scam_mod",
                  "incidence_base_scam_mod")
 
-save(list = model_names, file = glue("{data_path}/2025_06_29_malaria_models.RData"))
+save(list = model_names, file = glue("{data_path}/2025_07_03_malaria_models.RData"))

src/idd_forecast_mbp/04_forecasting/02_forecast_dengue_admin_2s_launcher.r

@@ -18,7 +18,7 @@ write.csv(param_map,  param_map_filepath, row.names = FALSE)
 ## QSUB Command
 job_name <- glue("forecast_dengue")   # name of the job
 thread_flag <- "-c 4" 
-mem_flag <- "--mem=100G" 
+mem_flag <- "--mem=75G" 
 runtime_flag <- "-t 50"
 #jdrive_flag <- "-l archive" # archive nodes can access the J drive. They're a little harder to get though. If you need J drive access, uncomment this and add it to the qsub_command
 queue_flag <- "-p all.q" # long or all