Merge branch '198-test-for-to_epiweekly_quantile_table' of https://gi…

…thub.com/CDCgov/pyrenew-hew into 198-test-for-to_epiweekly_quantile_table

.github/workflows/containers.yaml

@@ -20,6 +20,7 @@ jobs:
     outputs:
       tag: ${{ steps.image-tag.outputs.tag }}
       commit-msg: ${{ steps.commit-message.outputs.message }}
+      branch: ${{ steps.branch-name.outputs.branch }}
 
     steps:
 
@@ -115,7 +116,9 @@ jobs:
         with:
           push: true # This can be toggled manually for tweaking.
           tags: |
-            ${{ env.REGISTRY}}/${{ env.IMAGE_NAME }}:${{ needs.build-dependencies-image.outputs.tag }}
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ needs.build-dependencies-image.outputs.tag }}
           file: ./Containerfile
           build-args: |
             TAG=${{ needs.build-dependencies-image.outputs.tag }}
+            GIT_COMMIT_SHA=${{ github.event.pull_request.head.sha || github.sha }}
+            GIT_BRANCH_NAME=${{ needs.build-dependencies-image.outputs.branch }}

.github/workflows/pipeline-run-check.yaml

@@ -0,0 +1,39 @@
+name: Pipeline Run Check
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  run-pipeline:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: "Set up R"
+        uses: r-lib/actions/setup-r@v2
+        with:
+          r-version: "release"
+          use-public-rspm: true
+      - name: "Install poetry"
+        run: pip install poetry
+      - name: "Install pyrenew-hew"
+        run: poetry install
+      - name: "Set up Quarto"
+        uses: quarto-dev/quarto-actions/setup@v2
+      - name: "Set up dependencies for hewr"
+        uses: r-lib/actions/setup-r-dependencies@v2
+        with:
+          working-directory: hewr
+      - name: "Install extra pkgs"
+        run: |
+            pak::local_install("hewr", ask = FALSE)
+            pak::pkg_install("cmu-delphi/epipredict@main", ask = FALSE)
+            pak::pkg_install("cmu-delphi/epiprocess@main", ask = FALSE)
+        shell: Rscript {0}
+      - name: "Run pipeline"
+        run: poetry run bash pipelines/tests/test_forecast_state.sh pipelines/tests

.gitignore

@@ -403,3 +403,6 @@ private_data/*
 # Test data exceptions to the general data exclusion
 !pipelines/tests/covid-19_r_2024-01-29_f_2023-11-01_t_2024-01-29/model_runs/TD/data/data.tsv
 !pipelines/tests/covid-19_r_2024-01-29_f_2023-11-01_t_2024-01-29/model_runs/TD/data/eval_data.tsv
+
+# Ignore test pipe output
+pipelines/tests/private_data/*

Containerfile

@@ -1,6 +1,12 @@
+#syntax=docker/dockerfile:1.7-labs
 ARG TAG=latest
 
 FROM cfaprdbatchcr.azurecr.io/pyrenew-hew-dependencies:${TAG}
+ARG GIT_COMMIT_SHA
+ENV GIT_COMMIT_SHA=$GIT_COMMIT_SHA
+
+ARG GIT_BRANCH_NAME
+ENV GIT_BRANCH_NAME=$GIT_BRANCH_NAME
 
 COPY ./hewr /pyrenew-hew/hewr
 
@@ -13,5 +19,8 @@ RUN Rscript -e "pak::pkg_install('cmu-delphi/epiprocess@main')"
 RUN Rscript -e "pak::pkg_install('cmu-delphi/epipredict@main')"
 RUN Rscript -e "pak::local_install('hewr')"
 
-COPY . .
+COPY --exclude=pipelines/priors . .
 RUN pip install --root-user-action=ignore .
+
+# Copy priors folder last
+COPY pipelines/priors pipelines/priors

hewr/NAMESPACE

@@ -1,6 +1,7 @@
 # Generated by roxygen2: do not edit by hand
 
 export(combine_training_and_eval_data)
+export(generate_exp_growth_pois)
 export(get_all_model_batch_dirs)
 export(make_forecast_figure)
 export(parse_model_batch_dir_path)

hewr/R/generate_exp_growth_process.R

@@ -0,0 +1,26 @@
+#' Generate Exponential Growth Process with Poisson Noise
+#'
+#' This function generates a sequence of samples from an exponential growth
+#' process through Poisson sampling:
+#' ```math
+#' \begin{aligned}
+#' \( \lambda_t &= I_0 \exp(\sum_{t=1}^{t} r_t) \) \\
+#' I_t &\sim \text{Poisson}(\lambda_t).
+#' ```
+#' @param rt A numeric vector of exponential growth rates.
+#' @param initial A numeric value representing the initial value of the process.
+#'
+#' @return An integer vector of Poisson samples generated from the exponential
+#' growth process.
+#'
+#' @examples
+#' rt <- c(0.1, 0.2, 0.15)
+#' initial <- 10
+#' generate_exp_growth_pois(rt, initial)
+#'
+#' @export
+generate_exp_growth_pois <- function(rt, initial) {
+  means <- initial * exp(cumsum(rt))
+  samples <- stats::rpois(length(means), lambda = means)
+  return(samples)
+}

hewr/tests/testthat/test_generate_exp_growth_process.R

@@ -0,0 +1,30 @@
+test_that("generate_exp_growth_pois generates correct number of samples", {
+  rt <- c(0.1, 0.2, 0.15)
+  initial <- 10
+  samples <- generate_exp_growth_pois(rt, initial)
+  expect_length(samples, length(rt))
+})
+
+test_that("generate_exp_growth_pois returns a vector of integers", {
+  rt <- c(0.1, 0.2, 0.15)
+  initial <- 10
+  samples <- generate_exp_growth_pois(rt, initial)
+  expect_type(samples, "integer")
+})
+
+test_that("generate_exp_growth_pois does not return implausible values", {
+  rt <- c(0.1, 0.2, 0.15)
+  initial <- 10
+  analytic_av <- initial * exp(cumsum(rt))
+  analytic_std <- sqrt(analytic_av)
+  samples <- generate_exp_growth_pois(rt, initial)
+  expect_true(all(samples <= initial + 10 * analytic_std))
+  expect_true(all(samples >= initial - 10 * analytic_std))
+})
+
+test_that("generate_exp_growth_pois handles empty growth rates", {
+  rt <- numeric(0)
+  initial <- 10
+  samples <- generate_exp_growth_pois(rt, initial)
+  expect_equal(samples, numeric(0))
+})

pipelines/batch/setup_prod_job.py

@@ -167,7 +167,7 @@ def main(
         "nssp-archival-vintages/gold "
         "--param-data-dir params "
         "--output-dir {output_dir} "
-        "--priors-path config/prod_priors.py "
+        "--priors-path pipelines/priors/prod_priors.py "
         "--report-date {report_date} "
         "--exclude-last-n-days {exclude_last_n_days} "
         "--no-score "

pipelines/batch/setup_test_prod_job.py

@@ -21,7 +21,7 @@
         "--tag",
         type=str,
         help="The tag name to use for the container image version",
-        default=Path(Repository(os.getcwd()).head.name).stem,
+        default=Repository(os.getcwd()).head.shorthand,
     )
 
     args = parser.parse_args()
@@ -95,6 +95,6 @@
         diseases=["COVID-19", "Influenza"],
         container_image_name="pyrenew-hew",
         container_image_version=tag,
-        excluded_locations=locs_to_exclude,
+        locations_exclude=locs_to_exclude,
         test=True,
     )

pipelines/collate_plots.py

@@ -5,8 +5,9 @@
 import os
 from pathlib import Path
 
+from forecasttools import ensure_listlike
 from pypdf import PdfWriter
-from utils import ensure_listlike, get_all_forecast_dirs
+from utils import get_all_forecast_dirs
 
 
 def merge_pdfs_and_save(

pipelines/forecast_state.py

@@ -8,7 +8,10 @@
 
 import numpyro
 import polars as pl
+import tomli_w
+import tomllib
 from prep_data import process_and_save_state
+from pygit2 import Repository
 from save_eval_data import save_eval_data
 
 numpyro.set_host_device_count(4)
@@ -17,6 +20,55 @@
 from generate_predictive import generate_and_save_predictions  # noqa
 
 
+def record_git_info(model_run_dir: Path):
+    metadata_file = Path(model_run_dir, "metadata.toml")
+
+    if metadata_file.exists():
+        with open(metadata_file, "rb") as file:
+            metadata = tomllib.load(file)
+    else:
+        metadata = {}
+
+    try:
+        repo = Repository(os.getcwd())
+        branch_name = repo.head.shorthand
+        commit_sha = str(repo.head.target)
+    except Exception as e:
+        branch_name = os.environ.get("GIT_BRANCH_NAME", "unknown")
+        commit_sha = os.environ.get("GIT_COMMIT_SHA", "unknown")
+
+    new_metadata = {
+        "branch_name": branch_name,
+        "commit_sha": commit_sha,
+    }
+
+    metadata.update(new_metadata)
+
+    metadata_file.parent.mkdir(parents=True, exist_ok=True)
+    with open(metadata_file, "wb") as file:
+        tomli_w.dump(metadata, file)
+
+
+def copy_and_record_priors(priors_path: Path, model_run_dir: Path):
+    metadata_file = Path(model_run_dir, "metadata.toml")
+    shutil.copyfile(priors_path, Path(model_run_dir, "priors.py"))
+
+    if metadata_file.exists():
+        with open(metadata_file, "rb") as file:
+            metadata = tomllib.load(file)
+    else:
+        metadata = {}
+
+    new_metadata = {
+        "priors_path": str(priors_path),
+    }
+
+    metadata.update(new_metadata)
+
+    with open(metadata_file, "wb") as file:
+        tomli_w.dump(metadata, file)
+
+
 def generate_epiweekly(model_run_dir: Path) -> None:
     result = subprocess.run(
         [
@@ -238,8 +290,11 @@ def main(
 
     os.makedirs(model_run_dir, exist_ok=True)
 
-    logger.info(f"Using priors from {priors_path}...")
-    shutil.copyfile(priors_path, Path(model_run_dir, "priors.py"))
+    logger.info("Recording git info...")
+    record_git_info(model_run_dir)
+
+    logger.info(f"Copying and recording priors from {priors_path}...")
+    copy_and_record_priors(priors_path, model_run_dir)
 
     logger.info(f"Processing {state}")
     process_and_save_state(