Issue 105: Generate test forecasts and scoring output

.gitignore

@@ -10,6 +10,8 @@
 *.xls
 *.xlsx
 *.rds
+*.pickle
+*.nc
 
 # Documents
 *.doc
@@ -395,3 +397,7 @@ notebooks/*.md
 private_data/*
 *_files/
 .vscode/settings.json
+
+# Test data exceptions to the general data exclusion
+!pipelines/tests/covid-19_r_run_test_inference/TestDir/data.csv
+!pipelines/tests/covid-19_r_run_test_inference/TestDir/eval_data.tsv

pipelines/convert_inferencedata_to_parquet.R

@@ -1,11 +1,24 @@
-library(forecasttools)
-library(readr)
-library(arrow)
-library(fs)
-library(argparser)
-library(dplyr)
-library(stringr)
-library(tidyr)
+script_packages <- c(
+  "argparser",
+  "arrow",
+  "dplyr",
+  "forecasttools",
+  "fs",
+  "ggplot2",
+  "lubridate",
+  "readr",
+  "scoringutils",
+  "stringr",
+  "tidyr"
+)
+
+## load in packages without messages
+purrr::walk(script_packages, \(pkg) {
+  suppressPackageStartupMessages(
+    library(pkg, character.only = TRUE)
+  )
+})
+
 
 tidy_and_save_mcmc <- function(model_run_dir,
                                file_name_prefix = "",

pipelines/iteration_helpers/loop_postprocess.sh

@@ -13,5 +13,6 @@ BASE_DIR="$1"
 for SUBDIR in "$BASE_DIR"/*/; do
     # Run the R script with the current subdirectory as the model_dir argument
     echo "$SUBDIR"
+    Rscript convert_inferencedata_to_parquet.R "$SUBDIR"
     Rscript postprocess_state_forecast.R "$SUBDIR"
 done

pipelines/tests/README.md

@@ -0,0 +1,10 @@
+# Test data folder
+
+This folder is aimed at running test-mode scripts for validating the inference
+pipeline on the test data. The test data is stored in subdirectories.
+
+To run the test scripts, execute the following command from the `pipelines` directory:
+
+```bash
+% bash ./tests/test_run.sh ./tests/covid-19_r_2024-01-29_f_2023-11-01_t_2024-01-29/model_runs 1000 28
+```

pipelines/tests/covid-19_r_2024-01-29_f_2023-11-01_t_2024-01-29/model_runs/TD/priors.py

@@ -0,0 +1,68 @@
+import jax.numpy as jnp
+import numpyro.distributions as dist
+import pyrenew.transformation as transformation
+from numpyro.infer.reparam import LocScaleReparam
+from pyrenew.randomvariable import DistributionalVariable, TransformedVariable
+
+i0_first_obs_n_rv = DistributionalVariable(
+    "i0_first_obs_n_rv",
+    dist.Beta(1, 10),
+)
+
+initialization_rate_rv = DistributionalVariable(
+    "rate", dist.Normal(0, 0.01), reparam=LocScaleReparam(0)
+)
+
+r_logmean = jnp.log(1)
+r_logsd = jnp.log(jnp.sqrt(2))
+
+log_r_mu_intercept_rv = DistributionalVariable(
+    "log_r_mu_intercept_rv", dist.Normal(r_logmean, r_logsd)
+)
+
+eta_sd_rv = DistributionalVariable(
+    "eta_sd", dist.TruncatedNormal(0.04, 0.02, low=0)
+)
+
+autoreg_rt_rv = DistributionalVariable("autoreg_rt", dist.Beta(2, 40))
+
+
+inf_feedback_strength_rv = TransformedVariable(
+    "inf_feedback",
+    DistributionalVariable(
+        "inf_feedback_raw",
+        dist.LogNormal(jnp.log(50), jnp.log(2)),
+    ),
+    transforms=transformation.AffineTransform(loc=0, scale=-1),
+)
+# Could be reparameterized?
+
+p_ed_visit_mean_rv = DistributionalVariable(
+    "p_ed_visit_mean",
+    dist.Normal(
+        transformation.SigmoidTransform().inv(0.005),
+        0.3,
+    ),
+)  # logit scale
+
+
+p_ed_visit_w_sd_rv = DistributionalVariable(
+    "p_ed_visit_w_sd_sd", dist.TruncatedNormal(0, 0.01, low=0)
+)
+
+
+autoreg_p_ed_visit_rv = DistributionalVariable(
+    "autoreg_p_ed_visit_rv", dist.Beta(1, 100)
+)
+
+ed_visit_wday_effect_rv = TransformedVariable(
+    "ed_visit_wday_effect",
+    DistributionalVariable(
+        "ed_visit_wday_effect_raw",
+        dist.Dirichlet(jnp.array([5, 5, 5, 5, 5, 5, 5])),
+    ),
+    transformation.AffineTransform(loc=0, scale=7),
+)
+
+# Based on looking at some historical posteriors.
+phi_rv = DistributionalVariable("phi", dist.LogNormal(6, 1))

pipelines/tests/test_run.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+
+# Check if the base directory is provided as an argument
+if [ -z "$1" ]; then
+    echo "Usage: $0 <base_dir>"
+    exit 1
+fi
+
+# Base directory containing subdirectories
+BASE_DIR="$1"
+N_SAMPLES=$2
+N_AHEAD=$3
+
+# Iterate over each subdirectory in the base directory
+echo "TEST-MODE: Running loop over subdirectories in $BASE_DIR"
+echo "For $N_SAMPLES samples on 1 chain, and $N_AHEAD forecast points"
+for SUBDIR in "$BASE_DIR"/*/; do
+    echo "TEST-MODE: Inference for $SUBDIR"
+    python fit_model.py "$SUBDIR" --n-chains 1 --n-samples $N_SAMPLES
+    echo "TEST-MODE: Finished inference"
+    echo "TEST-MODE: Generating posterior predictions for $SUBDIR"
+    python generate_predictive.py "$SUBDIR" --n-forecast-points $N_AHEAD
+    echo "TEST-MODE: Finished generating posterior predictions"
+    echo "TEST-MODE: Converting inferencedata to parquet for $SUBDIR"
+    Rscript convert_inferencedata_to_parquet.R "$SUBDIR"
+    echo "TEST-MODE: Finished converting inferencedata to parquet"
+    echo "TEST-MODE: Forecasting baseline models for $SUBDIR"
+    Rscript timeseries_forecasts.R "$SUBDIR" --n-forecast-days $N_AHEAD --n-samples $N_SAMPLES
+    echo "TEST-MODE: Finished forecasting baseline models"
+    echo "TEST-MODE: Postprocessing state forecast for $SUBDIR"
+    Rscript postprocess_state_forecast.R "$SUBDIR"
+    echo "TEST-MODE: Finished postprocessing state forecast"
+    echo "TEST-MODE: Scoring forecast for $SUBDIR"
+    Rscript score_forecast.R "$SUBDIR"
+    echo "TEST-MODE: Finished scoring forecast"
+done

pipelines/timeseries_forecasts.R

@@ -230,7 +230,7 @@ p <- arg_parser(
   "Forecast other (non-target-disease) ED visits for a given location."
 ) |>
   add_argument(
-    "model-run-dir",
+    "model_run_dir",
     help = "Directory containing the model data and output.",
   ) |>
   add_argument(