Add switches to fitting scripts

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,
@@ -253,6 +256,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/fit_pyrenew_model.py

@@ -75,6 +75,26 @@ def fit_and_save_model(
         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 @@ def main(
     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,7 +340,9 @@ def main(
         n_warmup=n_warmup,
         n_samples=n_samples,
         n_chains=n_chains,
-        sample_ed_visits=True,
+        sample_ed_visits=sample_ed_visits,
+        sample_hospital_admissions=sample_hospital_admissions,
+        sample_wastewater=sample_wastewater,
     )
     logger.info("Model fitting complete")
 
@@ -348,7 +353,9 @@ def main(
         model_run_dir,
         "pyrenew_e",
         n_days_past_last_training,
-        predict_ed_visits=True,
+        predict_ed_visits=sample_ed_visits,
+        predict_hospital_admissions=sample_hospital_admissions,
+        predict_wastewater=sample_wastewater,
     )
 
     logger.info(
@@ -528,6 +535,26 @@ def main(
         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))