Make which quantities are sampled more readily configurable

pipelines/batch/setup_prod_job.py

@@ -21,6 +21,9 @@ def main(
     pool_id: str,
     diseases: str | list[str],
     output_subdir: str | Path = "./",
+    sample_ed_visits: bool = False,
+    sample_hospital_admissions: bool = False,
+    sample_wastewater: bool = False,
     container_image_name: str = "pyrenew-hew",
     container_image_version: str = "latest",
     n_training_days: int = 90,
@@ -155,24 +158,35 @@ def main(
         ],
     )
 
+    sample_ed_visits_flag = "--sample-ed-visits " if sample_ed_visits else ""
+    sample_hospital_admissions_flag = (
+        "--sample-ed-visits " if sample_hospital_admissions else ""
+    )
+    sample_wastewater_flag = (
+        "--sample-ed-visits " if sample_hospital_admissions else ""
+    )
+
     base_call = (
         "/bin/bash -c '"
         "python pipelines/forecast_state.py "
         "--disease {disease} "
         "--state {state} "
-        "--n-training-days {n_training_days} "
-        "--n-warmup {n_warmup} "
-        "--n-samples {n_samples} "
+        f"--n-training-days {n_training_days} "
+        f"--n-warmup {n_warmup} "
+        f"--n-samples {n_samples} "
         "--facility-level-nssp-data-dir nssp-etl/gold "
         "--state-level-nssp-data-dir "
         "nssp-archival-vintages/gold "
         "--param-data-dir params "
         "--output-dir {output_dir} "
         "--priors-path pipelines/priors/prod_priors.py "
         "--report-date {report_date} "
-        "--exclude-last-n-days {exclude_last_n_days} "
+        f"--exclude-last-n-days {exclude_last_n_days} "
         "--no-score "
         "--eval-data-path "
+        f"{sample_ed_visits_flag}"
+        f"{sample_hospital_admissions_flag}"
+        f"{sample_wastewater_flag}"
         "nssp-etl/latest_comprehensive.parquet"
         "'"
     )
@@ -196,10 +210,6 @@ def main(
                 state=state,
                 disease=disease,
                 report_date="latest",
-                n_warmup=n_warmup,
-                n_samples=n_samples,
-                n_training_days=n_training_days,
-                exclude_last_n_days=exclude_last_n_days,
                 output_dir=str(Path("output", output_subdir)),
             ),
             container_settings=container_settings,
@@ -253,6 +263,27 @@ def main(
     default="latest",
 )
 
+
+parser.add_argument(
+    "--sample-ed-visits",
+    type=bool,
+    action=argparse.BooleanOptionalAction,
+    help="If provided, fit to and predict ED visit data.",
+)
+parser.add_argument(
+    "--sample-hospital-admissions",
+    type=bool,
+    action=argparse.BooleanOptionalAction,
+    help=("If provided, fit to and predict hospital admissions data."),
+)
+parser.add_argument(
+    "--sample-wastewater",
+    type=bool,
+    action=argparse.BooleanOptionalAction,
+    help="If provided, fit to and predict wastewater data.",
+)
+
+
 parser.add_argument(
     "--n-training-days",
     type=int,

pipelines/build_pyrenew_model.py

@@ -18,6 +18,9 @@
 
 def build_model_from_dir(
     model_dir: Path,
+    sample_ed_visits: bool = False,
+    sample_hospital_admissions: bool = False,
+    sample_wastewater: bool = False,
 ) -> tuple[PyrenewHEWModel, PyrenewHEWData]:
     data_path = Path(model_dir) / "data" / "data_for_model_fit.json"
     prior_path = Path(model_dir) / "priors.py"
@@ -47,12 +50,20 @@
         jnp.array(model_data["generation_interval_pmf"]),
     )  # check if off by 1 or reversed
 
-    data_observed_disease_ed_visits = jnp.array(
-        model_data["data_observed_disease_ed_visits"]
+    data_observed_disease_ed_visits = (
+        jnp.array(model_data["data_observed_disease_ed_visits"])
+        if sample_ed_visits
+        else None
     )
-    data_observed_disease_hospital_admissions = jnp.array(
-        model_data["data_observed_disease_hospital_admissions"]
+    data_observed_disease_hospital_admissions = (
+        jnp.array(model_data["data_observed_disease_hospital_admissions"])
+        if sample_hospital_admissions
+        else None
     )
+
+    # placeholder
+    data_observed_disease_wastewater = None if sample_wastewater else None
+
     population_size = jnp.array(model_data["state_pop"])
 
     ed_right_truncation_pmf_rv = DeterministicVariable(
@@ -133,7 +144,7 @@
         data_observed_disease_hospital_admissions=(
             data_observed_disease_hospital_admissions
         ),
-        data_observed_disease_wastewater=None,  # placeholder
+        data_observed_disease_wastewater=data_observed_disease_wastewater,
         right_truncation_offset=right_truncation_offset,
         first_ed_visits_date=first_ed_visits_date,
         first_hospital_admissions_date=first_hospital_admissions_date,

pipelines/fit_pyrenew_model.py

@@ -12,6 +12,9 @@
 def fit_and_save_model(
     model_run_dir: str,
     model_name: str,
+    sample_ed_visits: bool = False,
+    sample_hospital_admissions: bool = False,
+    sample_wastewater: bool = False,
     n_warmup: int = 1000,
     n_samples: int = 1000,
     n_chains: int = 4,
@@ -26,12 +29,17 @@
             "rng_key must be an integer with which "
             "to seed :func:`jax.random.key`"
         )
-    (my_model, my_data) = build_model_from_dir(model_run_dir)
+    (my_model, my_data) = build_model_from_dir(
+        model_run_dir,
+        sample_ed_visits=sample_ed_visits,
+        sample_hospital_admissions=sample_hospital_admissions,
+        sample_wastewater=sample_wastewater,
+    )
     my_model.run(
         data=my_data,
-        sample_ed_visits=True,
-        sample_hospital_admissions=True,
-        sample_wastewater=False,
+        sample_ed_visits=sample_ed_visits,
+        sample_hospital_admissions=sample_hospital_admissions,
+        sample_wastewater=sample_wastewater,
         num_warmup=n_warmup,
         num_samples=n_samples,
         rng_key=rng_key,
@@ -67,6 +75,26 @@
         required=True,
         help="Name of the model to use for generating predictions.",
     )
+
+    parser.add_argument(
+        "--sample-ed-visits",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, fit to ED visit data.",
+    )
+    parser.add_argument(
+        "--sample-hospital-admissions",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help=("If provided, fit to hospital admissions data."),
+    )
+    parser.add_argument(
+        "--sample-wastewater",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, fit to wastewater data.",
+    )
+
     parser.add_argument(
         "--n-warmup",
         type=int,

pipelines/forecast_state.py

@@ -201,7 +201,10 @@
     exclude_last_n_days: int = 0,
     score: bool = False,
     eval_data_path: Path = None,
-):
+    sample_ed_visits: bool = False,
+    sample_hospital_admissions: bool = False,
+    sample_wastewater: bool = False,
+) -> None:
     logging.basicConfig(level=logging.INFO)
     logger = logging.getLogger(__name__)
 
@@ -337,14 +340,22 @@
         n_warmup=n_warmup,
         n_samples=n_samples,
         n_chains=n_chains,
+        sample_ed_visits=sample_ed_visits,
+        sample_hospital_admissions=sample_hospital_admissions,
+        sample_wastewater=sample_wastewater,
     )
     logger.info("Model fitting complete")
 
     logger.info("Performing posterior prediction / forecasting...")
 
     n_days_past_last_training = n_forecast_days + exclude_last_n_days
     generate_and_save_predictions(
-        model_run_dir, "pyrenew_e", n_days_past_last_training
+        model_run_dir,
+        "pyrenew_e",
+        n_days_past_last_training,
+        predict_ed_visits=sample_ed_visits,
+        predict_hospital_admissions=sample_hospital_admissions,
+        predict_wastewater=sample_wastewater,
     )
 
     logger.info(
@@ -524,6 +535,26 @@
         type=Path,
         help=("Path to a parquet file containing compehensive truth data."),
     )
+
+    parser.add_argument(
+        "--sample-ed-visits",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, fit to ED visit data.",
+    )
+    parser.add_argument(
+        "--sample-hospital-admissions",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help=("If provided, fit to hospital admissions data."),
+    )
+    parser.add_argument(
+        "--sample-wastewater",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, fit to wastewater data.",
+    )
+
     args = parser.parse_args()
     numpyro.set_host_device_count(args.n_chains)
     main(**vars(args))

pipelines/generate_predictive.py

@@ -9,13 +9,23 @@
 
 
 def generate_and_save_predictions(
-    model_run_dir: str | Path, model_name: str, n_forecast_points: int
+    model_run_dir: str | Path,
+    model_name: str,
+    n_forecast_points: int,
+    predict_ed_visits: bool = False,
+    predict_hospital_admissions: bool = False,
+    predict_wastewater: bool = False,
 ) -> None:
     model_run_dir = Path(model_run_dir)
     model_dir = Path(model_run_dir, model_name)
     if not model_dir.exists():
         raise FileNotFoundError(f"The directory {model_dir} does not exist.")
-    (my_model, my_data) = build_model_from_dir(model_run_dir)
+    (my_model, my_data) = build_model_from_dir(
+        model_run_dir,
+        sample_ed_visits=predict_ed_visits,
+        sample_hospital_admissions=predict_hospital_admissions,
+        sample_wastewater=predict_wastewater,
+    )
 
     my_model._init_model(1, 1)
     fresh_sampler = my_model.mcmc.sampler
@@ -31,9 +41,9 @@
 
     posterior_predictive = my_model.posterior_predictive(
         data=forecast_data,
-        sample_ed_visits=True,
-        sample_hospital_admissions=True,
-        sample_wastewater=False,
+        sample_ed_visits=predict_ed_visits,
+        sample_hospital_admissions=predict_hospital_admissions,
+        sample_wastewater=predict_wastewater,
     )
 
     idata = az.from_numpyro(
@@ -73,6 +83,28 @@
         default=0,
         help="Number of time points to forecast (Default: 0).",
     )
+    parser.add_argument(
+        "--predict-ed-visits",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, generate posterior predictions for ED visits.",
+    )
+    parser.add_argument(
+        "--predict-hospital-admissions",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help=(
+            "If provided, generate posterior predictions "
+            "for hospital admissions."
+        ),
+    )
+    parser.add_argument(
+        "--predict-wastewater",
+        type=bool,
+        action=argparse.BooleanOptionalAction,
+        help="If provided, generate posterior predictions for wastewater.",
+    )
+
     args = parser.parse_args()
 
     generate_and_save_predictions(**vars(args))