Add to_observe

Author: albrja

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**4.3.10 - 09/23/25**
+
+  - Bugfix: Fix bug in PublicHealthObserver to retain stratification columns in results
+
 **4.3.9 - 09/16/25**
 
   - Bugfix: require alive column in register_transition_count_observation()

src/vivarium_public_health/results/observer.py

@@ -12,6 +12,7 @@
 
 import pandas as pd
 from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
 from vivarium.framework.results import Observer
 
 from vivarium_public_health.results.columns import COLUMNS
@@ -38,6 +39,7 @@ def register_adding_observation(
         excluded_stratifications: list[str] = [],
         aggregator_sources: list[str] | None = None,
         aggregator: Callable[[pd.DataFrame], float | pd.Series] = len,
+        to_observe: Callable[[Event], bool] = lambda event: True,
     ) -> None:
         """Registers an adding observation to the results system.
 
@@ -73,6 +75,9 @@ def register_adding_observation(
             List of population view columns to be used in the `aggregator`.
         aggregator
             Function that computes the quantity for this observation.
+        to_observe
+            Function that takes an event and returns a boolean indicating whether
+            the observation should be performed for that event.
         """
         builder.results.register_adding_observation(
             name=name,
@@ -85,6 +90,7 @@ def register_adding_observation(
             excluded_stratifications=excluded_stratifications,
             aggregator_sources=aggregator_sources,
             aggregator=aggregator,
+            to_observe=to_observe,
         )
 
     def format_results(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
@@ -207,7 +213,7 @@ def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
     def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
         """Get the 'sub_entity' column.
 
-        This method should be overwritten in subclasses to provide the 'sub_entity' column.
+        This method can be overwritten in subclasses to provide the 'sub_entity' column.
 
         Parameters
         ----------

tests/results/test_disability_observer.py

@@ -72,6 +72,7 @@ def test_disability_observer_setup(mocker):
         excluded_stratifications=observer.configuration.exclude,
         aggregator_sources=cause_pipelines,
         aggregator=observer.disability_weight_aggregator,
+        to_observe=mocker.ANY,
     )
 
     assert set(observer.disability_classes) == set([DiseaseState, RiskAttributableDisease])