reuse filtered population during results processing when possible (#660)

Author: rmudambi

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**3.5.3 - 09/23/25**
+
+  - Reuse filtered population during results processing when possible
+
 **3.5.2 - 09/23/25**
 
   - Fix type hint in get_output_model_name_string() utility method

src/vivarium/framework/results/context.py

@@ -52,7 +52,7 @@ class ResultsContext:
         objects to be produced keyed by the observation name.
     grouped_observations
         Dictionary of observation details. It is of the format
-        {lifecycle_state: {(pop_filter, stratifications): list[Observation]}}.
+        {lifecycle_state: {pop_filter: {stratifications: list[Observation]}}}.
         Allowable lifecycle_states are "time_step__prepare", "time_step",
         "time_step__cleanup", and "collect_metrics".
     logger
@@ -65,8 +65,8 @@ def __init__(self) -> None:
         self.excluded_categories: dict[str, list[str]] = {}
         self.observations: dict[str, Observation] = {}
         self.grouped_observations: defaultdict[
-            str, defaultdict[tuple[str, tuple[str, ...] | None], list[Observation]]
-        ] = defaultdict(lambda: defaultdict(list))
+            str, defaultdict[str, defaultdict[tuple[str, ...] | None, list[Observation]]]
+        ] = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
 
     @property
     def name(self) -> str:
@@ -112,18 +112,18 @@ def set_stratifications(self) -> None:
         """
         used_stratifications: set[str] = set()
         for state_observations in self.grouped_observations.values():
-            for observation_details in state_observations.items():
-                (_, stratification_names), observations = observation_details
-                if stratification_names is None:
-                    continue
-
-                used_stratifications |= set(stratification_names)
-                for observation in observations:
-                    observation.stratifications = tuple(
-                        self.stratifications[name]
-                        for name in stratification_names
-                        if name in self.stratifications
-                    )
+            for pop_filter_observations in state_observations.values():
+                for stratification_names, observations in pop_filter_observations.items():
+                    if stratification_names is None:
+                        continue
+
+                    used_stratifications |= set(stratification_names)
+                    for observation in observations:
+                        observation.stratifications = tuple(
+                            self.stratifications[name]
+                            for name in stratification_names
+                            if name in self.stratifications
+                        )
 
         if unused_stratifications := set(self.stratifications.keys()) - used_stratifications:
             self.logger.info(
@@ -272,8 +272,8 @@ def register_observation(
             **kwargs,
         )
         self.observations[name] = observation
-        self.grouped_observations[observation.when][
-            (observation.pop_filter, stratifications)
+        self.grouped_observations[observation.when][observation.pop_filter][
+            stratifications
         ].append(observation)
         return observation
 
@@ -318,26 +318,36 @@ def gather_results(
 
         # Optimization: We store all the producers by pop_filter and stratifications
         # so that we only have to apply them once each time we compute results.
-        for (pop_filter, stratification_names), observations in self.grouped_observations[
+        for pop_filter, stratification_observations in self.grouped_observations[
             lifecycle_state
         ].items():
-            observations = [obs for obs in observations if obs in event_observations]
-            if not observations:
+            event_pop_filter_observations = [
+                observation
+                for observations in stratification_observations.values()
+                for observation in observations
+                if observation in event_observations
+            ]
+            if not event_pop_filter_observations:
                 continue
 
-            # Results production can be simplified to
-            # filter -> groupby -> aggregate in all situations we've seen.
-            filtered_pop = self._filter_population(
-                population, pop_filter, stratification_names
-            )
-            if filtered_pop.empty:
+            filtered_population = self._filter_population(population, pop_filter)
+            if filtered_population.empty:
                 continue
-            else:
+
+            for stratification_names, observations in stratification_observations.items():
+                observations = [
+                    obs for obs in observations if obs in event_pop_filter_observations
+                ]
+                if not observations:
+                    continue
+
                 pop: pd.DataFrame | DataFrameGroupBy[tuple[str, ...] | str, bool]
-                if stratification_names is None:
-                    pop = filtered_pop
-                else:
-                    pop = self._get_groups(stratification_names, filtered_pop)
+                pop = self._drop_na_stratifications(filtered_population, stratification_names)
+                if pop.empty:
+                    continue
+                if stratification_names is not None:
+                    pop = self._get_groups(stratification_names, pop)
+
                 for observation in observations:
                     results = observation.observe(pop, stratification_names)
                     yield (results, observation.name, observation.results_updater)
@@ -357,7 +367,8 @@ def get_observations(self, event: Event) -> list[Observation]:
         """
         return [
             observation
-            for observations in self.grouped_observations[event.name].values()
+            for stratification_observations in self.grouped_observations[event.name].values()
+            for observations in stratification_observations.values()
             for observation in observations
             if observation.to_observe(event)
         ]
@@ -436,24 +447,24 @@ def get_required_values(
             required_values.update(stratification.requires_values)
         return list(required_values)
 
-    def _filter_population(
-        self,
-        population: pd.DataFrame,
-        pop_filter: str,
-        stratification_names: tuple[str, ...] | None,
+    def _filter_population(self, population: pd.DataFrame, pop_filter: str) -> pd.DataFrame:
+        """Filter out simulants not to observe."""
+        return population.query(pop_filter) if pop_filter else population.copy()
+
+    def _drop_na_stratifications(
+        self, population: pd.DataFrame, stratification_names: tuple[str, ...] | None
     ) -> pd.DataFrame:
         """Filter out simulants not to observe."""
-        pop = population.query(pop_filter) if pop_filter else population.copy()
         if stratification_names:
             # Drop all rows in the mapped_stratification columns that have NaN values
             # (which only exist if the mapper returned an excluded category).
-            pop = pop.dropna(
+            population = population.dropna(
                 subset=[
                     get_mapped_col_name(stratification)
                     for stratification in stratification_names
                 ]
             )
-        return pop
+        return population
 
     @staticmethod
     def _get_groups(

tests/framework/results/test_context.py

@@ -752,45 +752,47 @@ def get_required_resources_kwargs(
 
 
 @pytest.mark.parametrize(
-    "pop_filter, stratifications",
-    [
-        ('familiar=="cat"', tuple()),
-        ('familiar=="spaghetti_yeti"', tuple()),
-        ("", ("new_col1",)),
-        ("", ("new_col1", "new_col2")),
-        ('familiar=="cat"', ("new_col1",)),
-        ("", tuple()),
-    ],
-    ids=[
-        "pop_filter",
-        "pop_filter_empties_dataframe",
-        "single_excluded_stratification",
-        "two_excluded_stratifications",
-        "pop_filter_and_excluded_stratification",
-        "no_pop_filter_or_excluded_stratifications",
-    ],
+    "pop_filter",
+    ['familiar=="cat"', 'familiar=="spaghetti_yeti"', ""],
+    ids=["pop_filter", "pop_filter_empties_dataframe", "no_pop_filter"],
 )
-def test__filter_population(pop_filter: str, stratifications: tuple[str, ...]) -> None:
+def test__filter_population(pop_filter: str) -> None:
     population = BASE_POPULATION.copy()
-    if stratifications:
-        # Make some of the stratifications missing to mimic mapping to excluded categories
-        population["new_col1"] = "new_value1"
-        population.loc[population["tracked"] == True, "new_col1"] = np.nan
-        if len(stratifications) == 2:
-            population["new_col2"] = "new_value2"
-            population.loc[population["new_col1"].notna(), "new_col2"] = np.nan
-        # Add on the post-stratified columns
-        for stratification in stratifications:
-            mapped_col = f"{stratification}_mapped_values"
-            population[mapped_col] = population[stratification]
 
     filtered_pop = ResultsContext()._filter_population(
-        population=population, pop_filter=pop_filter, stratification_names=stratifications
+        population=population, pop_filter=pop_filter
     )
     expected = population.copy()
     if pop_filter:
         familiar = pop_filter.split("==")[1].strip('"')
         expected = expected[expected["familiar"] == familiar]
+    assert filtered_pop.equals(expected)
+
+
+@pytest.mark.parametrize(
+    "stratifications",
+    [tuple(), ("new_col1",), ("new_col1", "new_col2")],
+    ids=[
+        "no_stratifications",
+        "single_excluded_stratification",
+        "two_excluded_stratifications",
+    ],
+)
+def test__drop_na_stratifications(stratifications: tuple[str, ...]) -> None:
+    population = BASE_POPULATION.copy()
+    population["new_col1"] = "new_value1"
+    population.loc[population["tracked"] == True, "new_col1"] = np.nan
+    population["new_col2"] = "new_value2"
+    population.loc[population["new_col1"].notna(), "new_col2"] = np.nan
+    # Add on the post-stratified columns
+    for stratification in stratifications:
+        mapped_col = f"{stratification}_mapped_values"
+        population[mapped_col] = population[stratification]
+
+    filtered_pop = ResultsContext()._drop_na_stratifications(
+        population=population, stratification_names=stratifications
+    )
+    expected = population.copy()
     for stratification in stratifications:
         expected = expected[expected[stratification].notna()]
     assert filtered_pop.equals(expected)

tests/framework/results/test_interface.py

@@ -195,9 +195,9 @@ def test_register_stratified_observation(mocker: MockerFixture) -> None:
 
     grouped_observations = interface._manager._results_context.grouped_observations
     assert len(grouped_observations) == 1
-    ((filter, stratifications), observations) = list(
-        grouped_observations["some-when"].items()
-    )[0]
+    filter = list(grouped_observations["some-when"].keys())[0]
+    stratifications = list(grouped_observations["some-when"][filter])[0]
+    observations = grouped_observations["some-when"][filter][stratifications]
     assert filter == "some-filter"
     assert isinstance(stratifications, tuple)  # for mypy in following set(stratifications)
     assert set(stratifications) == {
@@ -239,13 +239,15 @@ def test_register_unstratified_observation(mocker: MockerFixture) -> None:
         results_gatherer=lambda _: pd.DataFrame(),
         results_updater=lambda _, __: pd.DataFrame(),
     )
-    observations = interface._manager._results_context.grouped_observations
-    assert len(observations) == 1
-    ((filter, stratification), observation) = list(observations["some-when"].items())[0]
+    grouped_observations = interface._manager._results_context.grouped_observations
+    assert len(grouped_observations) == 1
+    filter = list(grouped_observations["some-when"].keys())[0]
+    stratifications = list(grouped_observations["some-when"][filter])[0]
+    observations = grouped_observations["some-when"][filter][stratifications]
     assert filter == "some-filter"
-    assert stratification is None
-    assert len(observation) == 1
-    obs = observation[0]
+    assert stratifications is None
+    assert len(observations) == 1
+    obs = observations[0]
     assert obs.name == "some-name"
     assert obs.pop_filter == "some-filter"
     assert obs.when == "some-when"
@@ -346,26 +348,33 @@ def test_register_multiple_adding_observations(mocker: MockerFixture) -> None:
     )
     # Test observation gets added
     assert len(interface._manager._results_context.grouped_observations) == 1
-    # Test for default pop_filter
-    assert ("tracked==True", ()) in interface._manager._results_context.grouped_observations[
-        lifecycle_states.TIME_STEP_CLEANUP
-    ]
+    assert (
+        interface._manager._results_context.grouped_observations[
+            lifecycle_states.TIME_STEP_CLEANUP
+        ]["tracked==True"][()][0].name
+        == "living_person_time"
+    )
+
     interface.register_adding_observation(
         name="undead_person_time",
-        pop_filter="undead == True",
+        pop_filter="undead==True",
         when=lifecycle_states.TIME_STEP_PREPARE,
         aggregator=_silly_aggregator,
     )
     # Test new observation gets added
     assert len(interface._manager._results_context.grouped_observations) == 2
-    # Preserve other observation and its pop filter
-    assert ("tracked==True", ()) in interface._manager._results_context.grouped_observations[
-        lifecycle_states.TIME_STEP_CLEANUP
-    ]
-    # Test for overridden pop_filter
-    assert ("undead == True", ()) in interface._manager._results_context.grouped_observations[
-        lifecycle_states.TIME_STEP_PREPARE
-    ]
+    assert (
+        interface._manager._results_context.grouped_observations[
+            lifecycle_states.TIME_STEP_CLEANUP
+        ]["tracked==True"][()][0].name
+        == "living_person_time"
+    )
+    assert (
+        interface._manager._results_context.grouped_observations[
+            lifecycle_states.TIME_STEP_PREPARE
+        ]["undead==True"][()][0].name
+        == "undead_person_time"
+    )
 
 
 @pytest.mark.parametrize("resource_type", ["value", "column"])
@@ -484,13 +493,15 @@ def test_register_concatenating_observation(mocker: MockerFixture) -> None:
         requires_values=["some-value", "some-other-value"],
         results_formatter=lambda _, __: pd.DataFrame(),
     )
-    observations = interface._manager._results_context.grouped_observations
-    assert len(observations) == 1
-    ((filter, stratification), observation) = list(observations["some-when"].items())[0]
+    grouped_observations = interface._manager._results_context.grouped_observations
+    assert len(grouped_observations) == 1
+    filter = list(grouped_observations["some-when"].keys())[0]
+    stratifications = list(grouped_observations["some-when"][filter])[0]
+    observations = grouped_observations["some-when"][filter][stratifications]
     assert filter == "some-filter"
-    assert stratification is None
-    assert len(observation) == 1
-    obs = observation[0]
+    assert stratifications is None
+    assert len(observations) == 1
+    obs = observations[0]
     assert obs.name == "some-name"
     assert obs.pop_filter == "some-filter"
     assert obs.when == "some-when"