Support callable and tuple(str, callable) for population filter (#139)

Albrja/mic-7064/Support callable and tuple(str, callable) for population filter


Support callable and tuple(str, callable) for population filter
- *Category*: Feature
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-7064

Changes and notes
-Support callable and tuple(str, callable) for population filter

Author: albrja

libs/engine/CHANGELOG.rst

@@ -1,3 +1,7 @@
+**5.2.0 - 06/29/26**
+
+- Support a callable population filter (or ``(query, callable)`` tuple) for all observations
+
 **5.1.7 - 06/22/26**
 
 - Pin vivarium-build-utils to v4.x and update Makefile to use ``vivarium.build_utils``

libs/engine/docs/nitpick-exceptions

@@ -48,6 +48,9 @@ py:class PandasObject
 py:class DataFrameGroupBy
 py:class ResultsFormatter
 py:class ResultsUpdater
+py:class PopulationFilter
+py:class _PopulationFilter
+py:class vivarium.engine.framework.results.interface._PopulationFilter
 py:class _NestedDict
 py:exc ResultsConfigurationError
 py:exc vivarium.engine.framework.results.exceptions.ResultsConfigurationError

libs/engine/docs/source/conf.py

@@ -4,7 +4,6 @@
 
 import importlib.metadata
 
-
 # -- Project information -----------------------------------------------------
 
 project = "vivarium.engine"

libs/engine/src/vivarium/engine/framework/results/context.py

@@ -26,7 +26,7 @@
 
 if TYPE_CHECKING:
     from vivarium.engine.framework.engine import Builder
-    from vivarium.engine.framework.results.interface import PopulationFilter
+    from vivarium.engine.framework.results.interface import _PopulationFilter
 
 
 class ResultsContext:
@@ -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: {PopulationFilter: {stratifications: list[Observation]}}}.
+        {lifecycle_state: {_PopulationFilter: {stratifications: list[Observation]}}}.
         Allowable lifecycle_states are "time_step__prepare", "time_step",
         "time_step__cleanup", and "collect_metrics".
     logger
@@ -67,7 +67,7 @@ def __init__(self) -> None:
         self.grouped_observations: defaultdict[
             str,
             defaultdict[
-                PopulationFilter,
+                _PopulationFilter,
                 defaultdict[tuple[str, ...] | None, list[Observation]],
             ],
         ] = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
@@ -223,7 +223,7 @@ def register_observation(
         self,
         observation_type: type[Observation],
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         stratifications: tuple[str, ...] | None,
@@ -430,14 +430,17 @@ def get_required_attributes(
         return list(required_attributes)
 
     def _filter_population(
-        self, population: pd.DataFrame, population_filter: PopulationFilter
+        self, population: pd.DataFrame, population_filter: _PopulationFilter
     ) -> pd.DataFrame:
         """Filter out simulants not to observe."""
         query = population_filter.query
         if not population_filter.include_untracked:
             # combine the tracking query with the population filter query
             query = pop_utils.combine_queries(query, self.get_tracked_query())
-        return population.query(query) if query else population.copy()
+        filtered = population.query(query) if query else population.copy()
+        if population_filter.index_filter is not None:
+            filtered = filtered.loc[population_filter.index_filter(filtered.index)]
+        return filtered
 
     def _drop_na_stratifications(
         self, population: pd.DataFrame, stratification_names: tuple[str, ...] | None

libs/engine/src/vivarium/engine/framework/results/interface.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, NamedTuple, Sequence, Union
+from typing import TYPE_CHECKING, Any, Sequence, Union
 
 import pandas as pd
 from pandas.core.groupby.generic import DataFrameGroupBy
@@ -40,6 +40,11 @@
 ResultsGatherer = Callable[[ResultsGathererInput], pd.DataFrame]
 """A Callable that optionally takes a possibly stratified population and returns
 new observation results."""
+PopulationIndexFilter = Callable[[pd.Index], pd.Index]  # type: ignore [type-arg]
+"""A Callable that takes a population index and returns the subset of indices to keep."""
+PopulationFilter = Union[str, PopulationIndexFilter, tuple[str, PopulationIndexFilter]]
+"""A population filter given as a Pandas query string, an index-filter callable, or a
+``(query, callable)`` tuple that applies the query first followed by the callable."""
 
 
 def _required_function_placeholder(
@@ -76,11 +81,44 @@ def _default_unstratified_observation_formatter(
     return results
 
 
-class PopulationFilter(NamedTuple):
-    """Container class for population query string and include_untracked flag."""
+class _PopulationFilter:
+    """Container for a query string, the untracked flag, and an optional index filter.
 
-    query: str = ""
-    include_untracked: bool = False
+    The ``pop_filter`` argument may be a Pandas query string, an index-filter
+    callable, or a ``(query, callable)`` tuple; it is resolved into the separate
+    ``query`` and ``index_filter`` attributes on construction.
+    """
+
+    def __init__(
+        self, pop_filter: PopulationFilter = "", include_untracked: bool = False
+    ) -> None:
+        self.query, self.index_filter = self._parse_pop_filter(pop_filter)
+        self.include_untracked = include_untracked
+
+    @staticmethod
+    def _parse_pop_filter(
+        pop_filter: PopulationFilter,
+    ) -> tuple[str, PopulationIndexFilter | None]:
+        """Resolve a pop_filter argument into a (query, index_filter) pair."""
+        if isinstance(pop_filter, str):
+            return pop_filter, None
+        if isinstance(pop_filter, tuple):
+            return pop_filter
+        return "", pop_filter
+
+    def __eq__(self, other: object) -> bool:
+        """Compare by field value so equal filters share a grouping key."""
+        if not isinstance(other, _PopulationFilter):
+            return NotImplemented
+        return (self.query, self.include_untracked, self.index_filter) == (
+            other.query,
+            other.include_untracked,
+            other.index_filter,
+        )
+
+    def __hash__(self) -> int:
+        """Hash by field value, consistent with __eq__, for use as a dict key."""
+        return hash((self.query, self.include_untracked, self.index_filter))
 
 
 class ResultsInterface(Interface):
@@ -199,7 +237,7 @@ def register_binned_stratification(
     def register_stratified_observation(
         self,
         name: str,
-        pop_filter: str = "",
+        pop_filter: PopulationFilter = "",
         include_untracked: bool = False,
         when: str = lifecycle_states.COLLECT_METRICS,
         requires_attributes: list[str] = [],
@@ -219,8 +257,9 @@ def register_stratified_observation(
             Name of the observation. It will also be the name of the output results file
             for this particular observation.
         pop_filter
-            A Pandas query filter string to filter the population down to the simulants who should
-            be considered for the observation.
+            A filter selecting which simulants to observe. Either a Pandas query string, a
+            callable mapping the population index to the subset of indices to keep, or a
+            ``(query, callable)`` tuple that applies the query first followed by the callable.
         include_untracked
             Whether to include simulants who are untracked from this observation.
         when
@@ -254,7 +293,7 @@ def register_stratified_observation(
         self._manager.register_observation(
             observation_type=StratifiedObservation,
             name=name,
-            population_filter=PopulationFilter(pop_filter, include_untracked),
+            population_filter=_PopulationFilter(pop_filter, include_untracked),
             when=when,
             requires_attributes=requires_attributes,
             results_updater=results_updater,
@@ -269,7 +308,7 @@ def register_stratified_observation(
     def register_unstratified_observation(
         self,
         name: str,
-        pop_filter: str = "",
+        pop_filter: PopulationFilter = "",
         include_untracked: bool = False,
         when: str = lifecycle_states.COLLECT_METRICS,
         requires_attributes: list[str] = [],
@@ -286,8 +325,9 @@ def register_unstratified_observation(
             Name of the observation. It will also be the name of the output results file
             for this particular observation.
         pop_filter
-            A Pandas query filter string to filter the population down to the simulants who should
-            be considered for the observation.
+            A filter selecting which simulants to observe. Either a Pandas query string, a
+            callable mapping the population index to the subset of indices to keep, or a
+            ``(query, callable)`` tuple that applies the query first followed by the callable.
         include_untracked
             Whether to include simulants who are untracked from this observation.
         when
@@ -317,7 +357,7 @@ def register_unstratified_observation(
         self._manager.register_observation(
             observation_type=UnstratifiedObservation,
             name=name,
-            population_filter=PopulationFilter(pop_filter, include_untracked),
+            population_filter=_PopulationFilter(pop_filter, include_untracked),
             when=when,
             requires_attributes=requires_attributes,
             results_updater=results_updater,
@@ -329,7 +369,7 @@ def register_unstratified_observation(
     def register_adding_observation(
         self,
         name: str,
-        pop_filter: str = "",
+        pop_filter: PopulationFilter = "",
         include_untracked: bool = False,
         when: str = lifecycle_states.COLLECT_METRICS,
         requires_attributes: list[str] = [],
@@ -355,8 +395,9 @@ def register_adding_observation(
             Name of the observation. It will also be the name of the output results file
             for this particular observation.
         pop_filter
-            A Pandas query filter string to filter the population down to the simulants who should
-            be considered for the observation.
+            A filter selecting which simulants to observe. Either a Pandas query string, a
+            callable mapping the population index to the subset of indices to keep, or a
+            ``(query, callable)`` tuple that applies the query first followed by the callable.
         include_untracked
             Whether to include simulants who are untracked from this observation.
         when
@@ -382,7 +423,7 @@ def register_adding_observation(
         self._manager.register_observation(
             observation_type=AddingObservation,
             name=name,
-            population_filter=PopulationFilter(pop_filter, include_untracked),
+            population_filter=_PopulationFilter(pop_filter, include_untracked),
             when=when,
             requires_attributes=requires_attributes,
             results_formatter=results_formatter,
@@ -396,7 +437,7 @@ def register_adding_observation(
     def register_concatenating_observation(
         self,
         name: str,
-        pop_filter: str = "",
+        pop_filter: PopulationFilter = "",
         include_untracked: bool = False,
         when: str = lifecycle_states.COLLECT_METRICS,
         requires_attributes: list[str] = [],
@@ -418,15 +459,16 @@ def register_concatenating_observation(
             Name of the observation. It will also be the name of the output results file
             for this particular observation.
         pop_filter
-            A Pandas query filter string to filter the population down to the simulants who should
-            be considered for the observation.
+            A filter selecting which simulants to observe. Either a Pandas query string, a
+            callable mapping the population index to the subset of indices to keep, or a
+            ``(query, callable)`` tuple that applies the query first followed by the callable.
         include_untracked
             Whether to include simulants who are untracked from this observation.
         when
             Name of the lifecycle phase the observation should happen. Valid values are:
             "time_step__prepare", "time_step", "time_step__cleanup", or "collect_metrics".
         requires_attributes
-            The population attributes that are required by the `aggregator`.
+            The population attributes to record for this observation.
         results_formatter
             Function that formats the raw observation results.
         to_observe
@@ -435,7 +477,7 @@ def register_concatenating_observation(
         self._manager.register_observation(
             observation_type=ConcatenatingObservation,
             name=name,
-            population_filter=PopulationFilter(pop_filter, include_untracked),
+            population_filter=_PopulationFilter(pop_filter, include_untracked),
             when=when,
             requires_attributes=requires_attributes,
             results_formatter=results_formatter,

libs/engine/src/vivarium/engine/framework/results/manager.py

@@ -25,7 +25,7 @@
 
 if TYPE_CHECKING:
     from vivarium.engine.framework.engine import Builder
-    from vivarium.engine.framework.results.interface import PopulationFilter
+    from vivarium.engine.framework.results.interface import _PopulationFilter
 
 
 class ResultsManager(Manager):
@@ -255,7 +255,7 @@ def register_observation(
         self,
         observation_type: type[Observation],
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         **kwargs: Any,

libs/engine/src/vivarium/engine/framework/results/observation.py

@@ -38,7 +38,7 @@
 )
 
 if TYPE_CHECKING:
-    from vivarium.engine.framework.results.interface import PopulationFilter
+    from vivarium.engine.framework.results.interface import _PopulationFilter
 
 VALUE_COLUMN = "value"
 
@@ -55,7 +55,7 @@ class Observation(ABC):
     name: str
     """Name of the observation. It will also be the name of the output results file
     for this particular observation."""
-    population_filter: PopulationFilter
+    population_filter: _PopulationFilter
     """A named tuple of population filtering details. The first item is a Pandas 
     query string to filter the population down to the simulants who should be 
     considered for the observation. The second item is a boolean indicating whether 
@@ -149,7 +149,7 @@ class UnstratifiedObservation(Observation):
     def __init__(
         self,
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         results_gatherer: Callable[[pd.DataFrame], pd.DataFrame],
@@ -233,7 +233,7 @@ class StratifiedObservation(Observation):
     def __init__(
         self,
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         results_updater: Callable[[pd.DataFrame, pd.DataFrame], pd.DataFrame],
@@ -441,7 +441,7 @@ class AddingObservation(StratifiedObservation):
     def __init__(
         self,
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         results_formatter: Callable[[str, pd.DataFrame], pd.DataFrame],
@@ -526,7 +526,7 @@ class ConcatenatingObservation(UnstratifiedObservation):
     def __init__(
         self,
         name: str,
-        population_filter: PopulationFilter,
+        population_filter: _PopulationFilter,
         when: str,
         requires_attributes: list[str],
         results_formatter: Callable[[str, pd.DataFrame], pd.DataFrame],