Different vector for hosp_report_prob_dist in tests

CDCgov · Apr 3, 2024 · 90627e6 · 90627e6
1 parent 1c2ba0a
commit 90627e6
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/model/src/test/test_model_hospitalizations.py b/model/src/test/test_model_hospitalizations.py
@@ -356,13 +356,22 @@ def test_model_hosp_with_obs_model_weekday_phosp():
 
     # Other random components
     weekday = jnp.array([1, 1, 1, 1, 2, 2])
-    weekday = weekday / weekday.sum()
     weekday = jnp.tile(weekday, 10)
+    weekday = weekday / weekday.sum()
     weekday = weekday[:31]
 
-    hosp_report_prob_dist = DeterministicVariable((weekday,))
     weekday = DeterministicVariable((weekday,))
 
+    hosp_report_prob_dist = jnp.array([0.9, 0.8, 0.7, 0.7, 0.6, 0.4])
+    hosp_report_prob_dist = jnp.tile(hosp_report_prob_dist, 10)
+    hosp_report_prob_dist = hosp_report_prob_dist / hosp_report_prob_dist.sum()
+
+    hosp_report_prob_dist = hosp_report_prob_dist[:31]
+
+    hosp_report_prob_dist = DeterministicVariable(
+        vars=(hosp_report_prob_dist,)
+    )
+
     latent_hospitalizations = HospitalAdmissions(
         infection_to_admission_interval=inf_hosp,
         infections_varname="infections",