add missing args to to_forecast_data, code review suggestions

pyrenew_hew/pyrenew_hew_data.py

@@ -34,7 +34,7 @@ def __init__(
         ww_log_lod: ArrayLike = None,
         date_observed_disease_wastewater: ArrayLike = None,
         data_observed_disease_wastewater_conc: ArrayLike = None,
-        pop_fraction: ArrayLike = jnp.array([1]),
+        pop_fraction: ArrayLike = None,
     ) -> None:
         self.n_ed_visits_data_days_ = n_ed_visits_data_days
         self.n_hospital_admissions_data_days_ = n_hospital_admissions_data_days
@@ -220,5 +220,7 @@ def to_forecast_data(self, n_forecast_points: int) -> Self:
             ww_observed_subpops=self.ww_observed_subpops,
             ww_observed_times=self.ww_observed_times,
             lab_site_to_subpop_map=self.lab_site_to_subpop_map,
+            ww_log_lod=self.ww_log_lod,
+            pop_fraction=self.pop_fraction,
             data_observed_disease_wastewater_conc=None,
         )

pyrenew_hew/pyrenew_wastewater_data.py

@@ -27,31 +27,31 @@ def __init__(
     def site_subpop_spine(self):
         ww_data_present = self.data_observed_disease_wastewater is not None
         if ww_data_present:
-            add_auxiliary_subpop = (
-                self.population_size
-                > self.data_observed_disease_wastewater.select(
-                    pl.col("site_pop", "site", "lab", "lab_site_index")
-                )
-                .unique()
-                .get_column("site_pop")
-                .sum()
-            )
             site_indices = (
                 self.data_observed_disease_wastewater.select(
                     ["site_index", "site", "site_pop"]
                 )
                 .unique()
                 .sort("site_index")
             )
+
+            total_pop_ww = (
+                self.data_observed_disease_wastewater.unique(
+                    ["site_pop", "site"]
+                )
+                .get_column("site_pop")
+                .sum()
+            )
+
+            total_pop_no_ww = self.population_size - total_pop_ww
+            add_auxiliary_subpop = total_pop_no_ww > 0
+
             if add_auxiliary_subpop:
                 aux_subpop = pl.DataFrame(
                     {
                         "site_index": [None],
                         "site": [None],
-                        "site_pop": (
-                            self.population_size
-                            - site_indices.get_column("site_pop").sum()
-                        ),
+                        "site_pop": [total_pop_no_ww],
                     }
                 )
             else:
@@ -192,11 +192,16 @@ def lab_site_to_subpop_map(self):
 
     def to_pyrenew_hew_data_args(self):
         return {
-            attr: value
-            for attr, value in (
-                (attr, getattr(self, attr))
-                for attr, prop in self.__class__.__dict__.items()
-                if isinstance(prop, property)
-            )
-            if isinstance(value, ArrayLike)
+            attr: getattr(self, attr)
+            for attr in [
+                "n_ww_lab_sites",
+                "ww_censored",
+                "ww_uncensored",
+                "ww_log_lod",
+                "ww_observed_lab_sites",
+                "ww_observed_subpops",
+                "ww_observed_times",
+                "data_observed_disease_wastewater_conc",
+                "lab_site_to_subpop_map",
+            ]
         }

tests/test_pyrenew_wastewater_data.py

@@ -60,6 +60,21 @@ def test_pyrenew_wastewater_data():
     assert forecast_data.data_observed_disease_wastewater_conc is None
     assert data.data_observed_disease_wastewater_conc is not None
 
+    assert np.array_equal(data.ww_censored, forecast_data.ww_censored)
+    assert np.array_equal(data.ww_uncensored, forecast_data.ww_uncensored)
+    assert np.array_equal(data.ww_log_lod, forecast_data.ww_log_lod)
+    assert np.array_equal(
+        data.ww_observed_lab_sites, forecast_data.ww_observed_lab_sites
+    )
+    assert np.array_equal(
+        data.ww_observed_subpops, forecast_data.ww_observed_subpops
+    )
+    assert np.array_equal(
+        data.ww_observed_times, forecast_data.ww_observed_times
+    )
+    assert np.array_equal(data.n_ww_lab_sites, forecast_data.n_ww_lab_sites)
+    assert np.array_equal(data.pop_fraction, forecast_data.pop_fraction)
+
     assert np.array_equal(
         data.data_observed_disease_wastewater_conc,
         ww_data["log_genome_copies_per_ml"],