Repository: ihmeuw/vivarium_public_health (language: python)

Issue:
throw an error if a mapper produces an invalid category
## Throw an error if mapper produces invalid category
<!-- Ideally, <=50 chars. 50 chars is here..: -->

### Description
<!-- For use in commit message, wrap at 72 chars. 72 chars is here: -->
- *Category*: bugfix
- *JIRA issue*: [MIC-3016](https://jira.ihme.washington.edu/browse/MIC-3016)

Wrap mappers so that that they throw an error if they produce an illegal value
Add docstrings.

### Testing
Ran test suite.
Ran toy CIFF model.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1560fea36..c4b2aafa4 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,11 @@
+**0.10.18 - 04/22/22**
+
+ - Improve ScaleUp component configuration
+ - Enable Mortality component to handle affected unmodeled causes
+ - Refactor RiskEffect calculation for clarity and extensibility
+ - Implement ResultsStratifier to stratify outputs
+ - Refactor all observers to be compatible with the ResultsStratifier component
+
 **0.10.17 - 02/15/22**
 
  - Autoformat code with black and isort.
diff --git a/docs/source/api_reference/metrics/population.rst b/docs/source/api_reference/metrics/population.rst
deleted file mode 100644
index 22a559501..000000000
--- a/docs/source/api_reference/metrics/population.rst
+++ /dev/null
@@ -1,1 +0,0 @@
-.. automodule:: vivarium_public_health.metrics.population
diff --git a/docs/source/api_reference/metrics/stratification.rst b/docs/source/api_reference/metrics/stratification.rst
new file mode 100644
index 000000000..95ab0c783
--- /dev/null
+++ b/docs/source/api_reference/metrics/stratification.rst
@@ -0,0 +1,1 @@
+.. automodule:: vivarium_public_health.metrics.stratification
diff --git a/docs/source/api_reference/metrics/utilities.rst b/docs/source/api_reference/metrics/utilities.rst
deleted file mode 100644
index 3a862f8ab..000000000
--- a/docs/source/api_reference/metrics/utilities.rst
+++ /dev/null
@@ -1,1 +0,0 @@
-.. automodule:: vivarium_public_health.metrics.utilities
diff --git a/setup.py b/setup.py
index b21c49ba3..b10569181 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@
         "scipy",
         "tables",
         "risk_distributions>=2.0.6",
+        "dataclasses; python_version < '3.7.0'",
     ]
 
     test_requirements = [
@@ -59,7 +60,7 @@
             "Operating System :: POSIX :: Linux",
             "Operating System :: Microsoft :: Windows",
             "Programming Language :: Python",
-            "Programming Language :: Python :: 3.6",
+            "Programming Language :: Python :: 3.8",
             "Programming Language :: Python :: Implementation :: CPython",
             "Topic :: Education",
             "Topic :: Scientific/Engineering",
diff --git a/src/vivarium_public_health/__about__.py b/src/vivarium_public_health/__about__.py
index fedf87082..0590e1d26 100644
--- a/src/vivarium_public_health/__about__.py
+++ b/src/vivarium_public_health/__about__.py
@@ -13,7 +13,7 @@
 __summary__ = "Components for modelling diseases, risks, and interventions with ``vivarium``"
 __uri__ = "https://github.com/ihmeuw/vivarium_public_health"
 
-__version__ = "0.10.17"
+__version__ = "0.10.18"
 
 __author__ = "The vivarium_public_health developers"
 __email__ = "vivarium.dev@gmail.com"
diff --git a/src/vivarium_public_health/disease/model.py b/src/vivarium_public_health/disease/model.py
index 79412441f..2420fcde3 100644
--- a/src/vivarium_public_health/disease/model.py
+++ b/src/vivarium_public_health/disease/model.py
@@ -8,12 +8,15 @@
 transitions at simulation initialization and during transitions.
 
 """
+from typing import List
+
 import numpy as np
 import pandas as pd
 from vivarium.exceptions import VivariumError
 from vivarium.framework.state_machine import Machine
 
 from vivarium_public_health.disease.state import SusceptibleState
+from vivarium_public_health.disease.transition import TransitionString
 
 
 class DiseaseModelError(VivariumError):
@@ -42,17 +45,17 @@ def name(self):
         return f"disease_model.{self.cause}"
 
     @property
-    def state_names(self):
+    def state_names(self) -> List[str]:
         return [s.name.split(".")[1] for s in self.states]
 
     @property
-    def transition_names(self):
+    def transition_names(self) -> List[TransitionString]:
         states = {s.name.split(".")[1]: s for s in self.states}
         transitions = []
         for state in states.values():
             for trans in state.transition_set.transitions:
                 _, _, init_state, _, end_state = trans.name.split(".")
-                transitions.append(f"{init_state}_TO_{end_state}")
+                transitions.append(TransitionString(f"{init_state}_TO_{end_state}"))
         return transitions
 
     def setup(self, builder):
diff --git a/src/vivarium_public_health/disease/special_disease.py b/src/vivarium_public_health/disease/special_disease.py
index 038b56acf..4d6063903 100644
--- a/src/vivarium_public_health/disease/special_disease.py
+++ b/src/vivarium_public_health/disease/special_disease.py
@@ -13,7 +13,8 @@
 import pandas as pd
 from vivarium.framework.values import list_combiner, union_post_processor
 
-from vivarium_public_health.utilities import EntityString
+from vivarium_public_health.disease.transition import TransitionString
+from vivarium_public_health.utilities import EntityString, is_non_zero
 
 
 class RiskAttributableDisease:
@@ -49,7 +50,7 @@ class RiskAttributableDisease:
     there is any mortality associated with this disease with the
     ``mortality`` configuration flag.
 
-    Finally, you may specify whether the someone should "recover"
+    Finally, you may specify whether an individual should "recover"
     from the disease if their exposure level falls outside the
     provided threshold.
 
@@ -57,7 +58,7 @@ class RiskAttributableDisease:
     ``protein_energy_malnutrition`` if their exposure drift out (or
     changes via an intervention) of the provided exposure categories.
     Having your ``fasting_plasma_glucose`` drop below a provided level
-    does not necessarily mean you're no longer diabetic however.
+    does not necessarily mean you're no longer diabetic.
 
     To add this component, you need to initialize it with full cause name
     and full risk name, e.g.,
@@ -101,7 +102,9 @@ def __init__(self, cause, risk):
             ]
         }
         self._state_names = [f"{self.cause.name}", f"susceptible_to_{self.cause.name}"]
-        self._transition_names = [f"susceptible_to_{self.cause.name}_TO_{self.cause.name}"]
+        self._transition_names = [
+            TransitionString(f"susceptible_to_{self.cause.name}_TO_{self.cause.name}")
+        ]
 
         self.excess_mortality_rate_pipeline_name = f"{self.cause.name}.excess_mortality_rate"
         self.excess_mortality_rate_paf_pipeline_name = (
@@ -127,6 +130,7 @@ def setup(self, builder):
         self.clock = builder.time.clock()
 
         disability_weight_data = builder.data.load(f"{self.cause}.disability_weight")
+        self.has_disability = is_non_zero(disability_weight_data)
         self.base_disability_weight = builder.lookup.build_table(
             disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
         )
@@ -152,6 +156,7 @@ def setup(self, builder):
         )
 
         excess_mortality_data = self.load_excess_mortality_rate_data(builder)
+        self.has_excess_mortality = is_non_zero(excess_mortality_data)
         self.base_excess_mortality_rate = builder.lookup.build_table(
             excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
         )
@@ -321,7 +326,7 @@ def continuous_filter(index):
     def adjust_state_and_transitions(self):
         if self.recoverable:
             self._transition_names.append(
-                f"{self.cause.name}_TO_susceptible_to_{self.cause.name}"
+                TransitionString(f"{self.cause.name}_TO_susceptible_to_{self.cause.name}")
             )
 
     def load_cause_specific_mortality_rate_data(self, builder):
diff --git a/src/vivarium_public_health/disease/state.py b/src/vivarium_public_health/disease/state.py
index eebfc93b3..681e9ca79 100644
--- a/src/vivarium_public_health/disease/state.py
+++ b/src/vivarium_public_health/disease/state.py
@@ -17,6 +17,7 @@
     ProportionTransition,
     RateTransition,
 )
+from vivarium_public_health.utilities import is_non_zero
 
 
 class BaseDiseaseState(State):
@@ -254,6 +255,7 @@ def setup(self, builder):
         )
 
         disability_weight_data = self.load_disability_weight_data(builder)
+        self.has_disability = is_non_zero(disability_weight_data)
         self.base_disability_weight = builder.lookup.build_table(
             disability_weight_data, key_columns=["sex"], parameter_columns=["age", "year"]
         )
@@ -267,6 +269,7 @@ def setup(self, builder):
         )
 
         excess_mortality_data = self.load_excess_mortality_rate_data(builder)
+        self.has_excess_mortality = is_non_zero(excess_mortality_data)
         self.base_excess_mortality_rate = builder.lookup.build_table(
             excess_mortality_data, key_columns=["sex"], parameter_columns=["age", "year"]
         )
diff --git a/src/vivarium_public_health/disease/transition.py b/src/vivarium_public_health/disease/transition.py
index ebdba24d2..54b86dff2 100644
--- a/src/vivarium_public_health/disease/transition.py
+++ b/src/vivarium_public_health/disease/transition.py
@@ -12,6 +12,14 @@
 from vivarium.framework.values import list_combiner, union_post_processor
 
 
+class TransitionString(str):
+    def __new__(cls, value):
+        # noinspection PyArgumentList
+        obj = str.__new__(cls, value.lower())
+        obj.from_state, obj.to_state = value.split("_TO_")
+        return obj
+
+
 class RateTransition(Transition):
     def __init__(self, input_state, output_state, get_data_functions=None, **kwargs):
         super().__init__(
diff --git a/src/vivarium_public_health/metrics/__init__.py b/src/vivarium_public_health/metrics/__init__.py
index 3aff3368f..b3adea18d 100644
--- a/src/vivarium_public_health/metrics/__init__.py
+++ b/src/vivarium_public_health/metrics/__init__.py
@@ -2,3 +2,4 @@
 from .disease import DiseaseObserver
 from .mortality import MortalityObserver
 from .risk import CategoricalRiskObserver
+from .stratification import ResultsStratifier
diff --git a/src/vivarium_public_health/metrics/disability.py b/src/vivarium_public_health/metrics/disability.py
index 8f52fdbb2..95cba6327 100644
--- a/src/vivarium_public_health/metrics/disability.py
+++ b/src/vivarium_public_health/metrics/disability.py
@@ -8,19 +8,24 @@
 
 """
 from collections import Counter
+from typing import Dict, List, Union
 
 import pandas as pd
+from vivarium.config_tree import ConfigTree
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.population import PopulationView
 from vivarium.framework.values import (
+    NumberLike,
+    Pipeline,
     list_combiner,
     rescale_post_processor,
     union_post_processor,
 )
 
 from vivarium_public_health.disease import DiseaseState, RiskAttributableDisease
-from vivarium_public_health.metrics.utilities import (
-    get_age_bins,
-    get_years_lived_with_disability,
-)
+from vivarium_public_health.metrics.stratification import ResultsStratifier
+from vivarium_public_health.utilities import to_years
 
 
 class DisabilityObserver:
@@ -28,8 +33,8 @@ class DisabilityObserver:
 
     By default, this counts both aggregate and cause-specific years lived
     with disability over the full course of the simulation. It can be
-    configured to bin the cause-specific YLDs into age groups, sexes, and years
-    by setting the ``by_age``, ``by_sex``, and ``by_year`` flags, respectively.
+    configured to add or remove stratification groups to the default groups
+    defined by a :class:ResultsStratifier.
 
     In the model specification, your configuration for this component should
     be specified as, e.g.:
@@ -37,99 +42,156 @@ class DisabilityObserver:
     .. code-block:: yaml
 
         configuration:
-            metrics:
+            observers:
                 disability:
-                    by_age: True
-                    by_year: False
-                    by_sex: True
-
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
     """
 
     configuration_defaults = {
-        "metrics": {
+        "observers": {
             "disability": {
-                "by_age": False,
-                "by_year": False,
-                "by_sex": False,
+                "exclude": [],
+                "include": [],
             }
         }
     }
 
+    def __init__(self):
+        self.ylds_column_name = "years_lived_with_disability"
+        self.metrics_pipeline_name = "metrics"
+        self.disability_weight_pipeline_name = "disability_weight"
+
+    def __repr__(self):
+        return "DisabilityObserver()"
+
+    ##############
+    # Properties #
+    ##############
+
     @property
     def name(self):
         return "disability_observer"
 
-    def setup(self, builder):
-        self.config = builder.configuration.metrics.disability
-        self.age_bins = get_age_bins(builder)
-        self.clock = builder.time.clock()
-        self.step_size = builder.time.step_size()
-        self.causes = [
-            c.state_id
-            for c in builder.components.get_components_by_type(
-                (DiseaseState, RiskAttributableDisease)
-            )
-        ]
-        self.years_lived_with_disability = Counter()
-        self.disability_weight_pipelines = {
-            cause: builder.value.get_value(f"{cause}.disability_weight")
-            for cause in self.causes
-        }
-
-        self.disability_weight = builder.value.register_value_producer(
-            "disability_weight",
+    #################
+    # Setup methods #
+    #################
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        self.config = self._get_stratification_configuration(builder)
+        self.stratifier = self._get_stratifier(builder)
+        self.disability_weight = self._get_disability_weight_pipeline(builder)
+        self._cause_components = self._get_cause_components(builder)
+        self.disability_pipelines = {}
+        self.population_view = self._get_population_view(builder)
+
+        self.counts = Counter()
+
+        self._register_post_setup_listener(builder)
+        self._register_simulant_initializer(builder)
+        self._register_time_step_prepare_listener(builder)
+        self._register_metrics_modifier(builder)
+
+    # noinspection PyMethodMayBeStatic
+    def _get_stratification_configuration(self, builder: Builder) -> ConfigTree:
+        return builder.configuration.observers.disability
+
+    # noinspection PyMethodMayBeStatic
+    def _get_stratifier(self, builder: Builder) -> ResultsStratifier:
+        return builder.components.get_component(ResultsStratifier.name)
+
+    def _get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
+        # todo observer should not be creating the disability weight pipeline
+        return builder.value.register_value_producer(
+            self.disability_weight_pipeline_name,
             source=lambda index: [pd.Series(0.0, index=index)],
             preferred_combiner=list_combiner,
             preferred_post_processor=_disability_post_processor,
         )
 
-        columns_required = ["tracked", "alive", "years_lived_with_disability"]
-        if self.config.by_age:
-            columns_required += ["age"]
-        if self.config.by_sex:
-            columns_required += ["sex"]
-        self.population_view = builder.population.get_view(columns_required)
+    # noinspection PyMethodMayBeStatic
+    def _get_cause_components(
+        self, builder: Builder
+    ) -> List[Union[DiseaseState, RiskAttributableDisease]]:
+        return builder.components.get_components_by_type(
+            (DiseaseState, RiskAttributableDisease)
+        )
+
+    # noinspection PyMethodMayBeStatic
+    def _get_population_view(self, builder: Builder) -> PopulationView:
+        columns_required = [
+            self.ylds_column_name,
+        ]
+        return builder.population.get_view(columns_required)
+
+    def _register_post_setup_listener(self, builder: Builder) -> None:
+        builder.event.register_listener("post_setup", self.on_post_setup)
+
+    def _register_simulant_initializer(self, builder: Builder) -> None:
+        # todo observer should not be modifying state table
         builder.population.initializes_simulants(
-            self.initialize_disability, creates_columns=["years_lived_with_disability"]
+            self.on_initialize_simulants,
+            creates_columns=[self.ylds_column_name],
         )
-        # FIXME: The state table is modified before the clock advances.
+
+    def _register_time_step_prepare_listener(self, builder: Builder) -> None:
         # In order to get an accurate representation of person time we need to look at
         # the state table before anything happens.
         builder.event.register_listener("time_step__prepare", self.on_time_step_prepare)
-        builder.value.register_value_modifier("metrics", modifier=self.metrics)
 
-    def initialize_disability(self, pop_data):
+    def _register_metrics_modifier(self, builder: Builder) -> None:
+        builder.value.register_value_modifier(
+            self.metrics_pipeline_name,
+            modifier=self.metrics,
+            requires_columns=[self.ylds_column_name],
+        )
+
+    ########################
+    # Event-driven methods #
+    ########################
+
+    # noinspection PyUnusedLocal
+    def on_post_setup(self, event: Event) -> None:
+        for cause in self._cause_components:
+            if cause.has_disability:
+                self.disability_pipelines[cause.state_id] = cause.disability_weight
+
+    def on_initialize_simulants(self, pop_data: pd.DataFrame) -> None:
         self.population_view.update(
-            pd.Series(0.0, index=pop_data.index, name="years_lived_with_disability")
+            pd.Series(0.0, index=pop_data.index, name=self.ylds_column_name)
         )
 
-    def on_time_step_prepare(self, event):
+    def on_time_step_prepare(self, event: Event) -> None:
+        step_size_in_years = to_years(event.step_size)
         pop = self.population_view.get(
             event.index, query='tracked == True and alive == "alive"'
         )
-        ylds_this_step = get_years_lived_with_disability(
-            pop,
-            self.config.to_dict(),
-            self.clock().year,
-            self.step_size(),
-            self.age_bins,
-            self.disability_weight_pipelines,
-            self.causes,
-        )
-        self.years_lived_with_disability.update(ylds_this_step)
-
-        pop.loc[:, "years_lived_with_disability"] += self.disability_weight(pop.index)
+        groups = self.stratifier.group(pop.index, self.config.include, self.config.exclude)
+        for label, group_mask in groups:
+            group_index = pop[group_mask].index
+            for cause, disability_weight in self.disability_pipelines.items():
+                disability_weight = disability_weight(group_index).sum() * step_size_in_years
+                new_observations = {
+                    f"ylds_due_to_{cause}_{label}": disability_weight,
+                }
+                self.counts.update(new_observations)
+
+        pop.loc[:, self.ylds_column_name] += self.disability_weight(pop.index)
         self.population_view.update(pop)
 
-    def metrics(self, index, metrics):
-        total_ylds = self.population_view.get(index)["years_lived_with_disability"].sum()
+    ##################################
+    # Pipeline sources and modifiers #
+    ##################################
+
+    def metrics(self, index: pd.Index, metrics: Dict) -> Dict:
+        total_ylds = self.population_view.get(index)[self.ylds_column_name].sum()
         metrics["years_lived_with_disability"] = total_ylds
-        metrics.update(self.years_lived_with_disability)
+        metrics.update(self.counts)
         return metrics
 
-    def __repr__(self):
-        return "DisabilityObserver()"
-
 
-def _disability_post_processor(value, step_size):
+def _disability_post_processor(value: NumberLike, step_size: pd.Timedelta) -> NumberLike:
     return rescale_post_processor(union_post_processor(value, step_size), step_size)
diff --git a/src/vivarium_public_health/metrics/disease.py b/src/vivarium_public_health/metrics/disease.py
index b8d13c403..df1e7da8d 100644
--- a/src/vivarium_public_health/metrics/disease.py
+++ b/src/vivarium_public_health/metrics/disease.py
@@ -8,30 +8,26 @@
 
 """
 from collections import Counter
+from typing import Dict, List
 
 import pandas as pd
+from vivarium.config_tree import ConfigTree
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.population import PopulationView
 
-from vivarium_public_health.metrics.utilities import (
-    TransitionString,
-    get_age_bins,
-    get_prevalent_cases,
-    get_state_person_time,
-    get_transition_count,
-)
+from vivarium_public_health.disease.transition import TransitionString
+from vivarium_public_health.metrics.stratification import ResultsStratifier
+from vivarium_public_health.utilities import to_years
 
 
 class DiseaseObserver:
-    """Observes disease counts, person time, and prevalent cases for a cause.
+    """Observes disease counts and person time for a cause.
 
-    By default, this observer computes aggregate susceptible person time
-    and counts of disease cases over the entire simulation.  It can be
-    configured to bin these into age_groups, sexes, and years by setting
-    the ``by_age``, ``by_sex``, and ``by_year`` flags, respectively.
-
-    It also can record prevalent cases on a particular sample date each year,
-    though by default this is disabled. These will also be binned based on the
-    flags set for the observer. Additionally, the sample date is configurable
-    and defaults to July 1st of each year.
+    By default, this observer computes aggregate susceptible person time and
+    counts of disease cases over the full course of the simulation. It can be
+    configured to add or remove stratification groups to the default groups
+    defined by a ResultsStratifier.
 
     In the model specification, your configuration for this component should
     be specified as, e.g.:
@@ -39,142 +35,170 @@ class DiseaseObserver:
     .. code-block:: yaml
 
         configuration:
-            metrics:
-                {YOUR_DISEASE_NAME}_observer:
-                    by_age: True
-                    by_year: False
-                    by_sex: True
-                    sample_prevalence:
-                        sample: True
-                        date:
-                            month: 4
-                            day: 10
+            observers:
+                cause_name:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
 
     """
 
     configuration_defaults = {
-        "metrics": {
-            "disease_observer": {
-                "by_age": False,
-                "by_year": False,
-                "by_sex": False,
-                "sample_prevalence": {
-                    "sample": False,
-                    "date": {
-                        "month": 7,
-                        "day": 1,
-                    },
-                },
+        "observers": {
+            "disease": {
+                "exclude": [],
+                "include": [],
             }
         }
     }
 
     def __init__(self, disease: str):
         self.disease = disease
-        self.configuration_defaults = {
-            "metrics": {
-                f"{disease}_observer": DiseaseObserver.configuration_defaults["metrics"][
-                    "disease_observer"
-                ]
+        self.configuration_defaults = self._get_configuration_defaults()
+
+        self.disease_component_name = f"disease_model.{self.disease}"
+        self.current_state_column_name = self.disease
+        self.previous_state_column_name = f"previous_{self.disease}"
+        self.metrics_pipeline_name = "metrics"
+
+    def __repr__(self):
+        return f"DiseaseObserver({self.disease})"
+
+    ##########################
+    # Initialization methods #
+    ##########################
+
+    # noinspection PyMethodMayBeStatic
+    def _get_configuration_defaults(self) -> Dict[str, Dict]:
+        return {
+            "observers": {
+                self.disease: DiseaseObserver.configuration_defaults["observers"]["disease"]
             }
         }
 
+    ##############
+    # Properties #
+    ##############
+
     @property
     def name(self):
         return f"disease_observer.{self.disease}"
 
-    def setup(self, builder):
-        self.config = builder.configuration["metrics"][f"{self.disease}_observer"]
-        self.clock = builder.time.clock()
-        self.age_bins = get_age_bins(builder)
+    #################
+    # Setup methods #
+    #################
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        self.config = self._get_stratification_configuration(builder)
+        self.stratifier = self._get_stratifier(builder)
+        self.states = self._get_states(builder)
+        self.transitions = self._get_transitions(builder)
+        self.population_view = self._get_population_view(builder)
+
         self.counts = Counter()
-        self.person_time = Counter()
-        self.prevalence = Counter()
 
-        comp = builder.components.get_component(f"disease_model.{self.disease}")
-        self.states = comp.state_names
-        self.transitions = comp.transition_names
+        self._register_simulant_initializer(builder)
+        self._register_time_step_prepare_listener(builder)
+        self._register_collect_metrics_listener(builder)
+        self._register_metrics_modifier(builder)
+
+    def _get_stratification_configuration(self, builder: Builder) -> ConfigTree:
+        return builder.configuration.observers[self.disease]
+
+    # noinspection PyMethodMayBeStatic
+    def _get_stratifier(self, builder: Builder) -> ResultsStratifier:
+        return builder.components.get_component(ResultsStratifier.name)
+
+    def _get_states(self, builder: Builder) -> List[str]:
+        return builder.components.get_component(self.disease_component_name).state_names
+
+    def _get_transitions(self, builder: Builder) -> List[TransitionString]:
+        return builder.components.get_component(self.disease_component_name).transition_names
 
-        self.previous_state_column = f"previous_{self.disease}"
+    # noinspection PyMethodMayBeStatic
+    def _get_population_view(self, builder: Builder) -> PopulationView:
+        columns_required = [
+            self.current_state_column_name,
+            self.previous_state_column_name,
+        ]
+        return builder.population.get_view(columns_required)
+
+    def _register_simulant_initializer(self, builder: Builder) -> None:
+        # todo observer should not be modifying state table
         builder.population.initializes_simulants(
-            self.on_initialize_simulants, creates_columns=[self.previous_state_column]
+            self.on_initialize_simulants,
+            creates_columns=[self.previous_state_column_name],
         )
 
-        columns_required = ["alive", f"{self.disease}", self.previous_state_column]
-        if self.config.by_age:
-            columns_required += ["age"]
-        if self.config.by_sex:
-            columns_required += ["sex"]
-        self.population_view = builder.population.get_view(columns_required)
-
-        builder.value.register_value_modifier("metrics", self.metrics)
-        # FIXME: The state table is modified before the clock advances.
+    def _register_time_step_prepare_listener(self, builder: Builder) -> None:
         # In order to get an accurate representation of person time we need to look at
         # the state table before anything happens.
         builder.event.register_listener("time_step__prepare", self.on_time_step_prepare)
-        builder.event.register_listener("collect_metrics", self.on_collect_metrics)
 
-    def on_initialize_simulants(self, pop_data):
-        self.population_view.update(
-            pd.Series("", index=pop_data.index, name=self.previous_state_column)
+    def _register_collect_metrics_listener(self, builder: Builder) -> None:
+        builder.event.register_listener("time_step", self.on_collect_metrics)
+
+    def _register_metrics_modifier(self, builder: Builder) -> None:
+        builder.value.register_value_modifier(
+            self.metrics_pipeline_name,
+            modifier=self.metrics,
         )
 
-    def on_time_step_prepare(self, event):
-        pop = self.population_view.get(event.index)
-
-        for state in self.states:
-            # noinspection PyTypeChecker
-            state_person_time_this_step = get_state_person_time(
-                pop,
-                self.config,
-                self.disease,
-                state,
-                self.clock().year,
-                event.step_size,
-                self.age_bins,
-            )
-            self.person_time.update(state_person_time_this_step)
+    ########################
+    # Event-driven methods #
+    ########################
+
+    def on_initialize_simulants(self, pop_data: pd.DataFrame) -> None:
+        self.population_view.update(
+            pd.Series("", index=pop_data.index, name=self.previous_state_column_name)
+        )
 
+    def on_time_step_prepare(self, event: Event) -> None:
+        step_size_in_years = to_years(event.step_size)
+        pop = self.population_view.get(
+            event.index, query='tracked == True and alive == "alive"'
+        )
+        groups = self.stratifier.group(pop.index, self.config.include, self.config.exclude)
+        for label, group_mask in groups:
+            for state in self.states:
+                state_in_group_mask = group_mask & (
+                    pop[self.current_state_column_name] == state
+                )
+                person_time_in_group = state_in_group_mask.sum() * step_size_in_years
+                new_observations = {
+                    f"{self.disease}_{state}_person_time_{label}": person_time_in_group
+                }
+                self.counts.update(new_observations)
+
+        # todo observer should not be maintaining this column
         # This enables tracking of transitions between states
-        prior_state_pop = self.population_view.get(event.index)
-        prior_state_pop[self.previous_state_column] = prior_state_pop[self.disease]
-        self.population_view.update(prior_state_pop)
-
-        if self._should_sample(event.time):
-            point_prevalence = get_prevalent_cases(
-                pop, self.config.to_dict(), self.disease, event.time, self.age_bins
-            )
-            self.prevalence.update(point_prevalence)
-
-    def on_collect_metrics(self, event):
-        pop = self.population_view.get(event.index)
-        for transition in self.transitions:
-            # noinspection PyTypeChecker
-            transition_counts_this_step = get_transition_count(
-                pop,
-                self.config,
-                self.disease,
-                TransitionString(transition),
-                event.time,
-                self.age_bins,
-            )
-            self.counts.update(transition_counts_this_step)
-
-    def _should_sample(self, event_time: pd.Timestamp) -> bool:
-        """Returns true if we should sample on this time step."""
-        should_sample = self.config.sample_prevalence.sample
-        if should_sample:
-            sample_date = pd.Timestamp(
-                year=event_time.year, **self.config.sample_prevalence.date.to_dict()
-            )
-            should_sample &= self.clock() <= sample_date < event_time
-        return should_sample
-
-    def metrics(self, index, metrics):
+        pop[self.previous_state_column_name] = pop[self.disease]
+        self.population_view.update(pop)
+
+    def on_collect_metrics(self, event: Event) -> None:
+        pop = self.population_view.get(
+            event.index, query='tracked == True and alive == "alive"'
+        )
+        groups = self.stratifier.group(pop.index, self.config.include, self.config.exclude)
+        for label, group_mask in groups:
+            for transition in self.transitions:
+                transition_mask = (
+                    group_mask
+                    & (pop[self.previous_state_column_name] == transition.from_state)
+                    & (pop[self.current_state_column_name] == transition.to_state)
+                )
+                new_observations = {
+                    f"{self.disease}_{transition}_event_count_{label}": transition_mask.sum()
+                }
+                self.counts.update(new_observations)
+
+    ##################################
+    # Pipeline sources and modifiers #
+    ##################################
+
+    # noinspection PyUnusedLocal
+    def metrics(self, index: pd.Index, metrics: Dict) -> Dict:
         metrics.update(self.counts)
-        metrics.update(self.person_time)
-        metrics.update(self.prevalence)
         return metrics
-
-    def __repr__(self):
-        return "DiseaseObserver()"
diff --git a/src/vivarium_public_health/metrics/mortality.py b/src/vivarium_public_health/metrics/mortality.py
index 0909db913..b33f8cf9d 100644
--- a/src/vivarium_public_health/metrics/mortality.py
+++ b/src/vivarium_public_health/metrics/mortality.py
@@ -7,22 +7,25 @@
 excess mortality in the simulation.
 
 """
+from collections import Counter
+from typing import Dict, List, Union
+
+import pandas as pd
+from vivarium.framework.engine import Builder, ConfigTree
+from vivarium.framework.event import Event
+from vivarium.framework.population import PopulationView
+
 from vivarium_public_health.disease import DiseaseState, RiskAttributableDisease
-from vivarium_public_health.metrics.utilities import (
-    get_age_bins,
-    get_deaths,
-    get_person_time,
-    get_years_of_life_lost,
-)
+from vivarium_public_health.metrics.stratification import ResultsStratifier
 
 
 class MortalityObserver:
     """An observer for cause-specific deaths, ylls, and total person time.
 
-    By default, this counts cause-specific deaths, years of life lost, and
-    total person time over the full course of the simulation. It can be
-    configured to bin these measures into age groups, sexes, and years
-    by setting the ``by_age``, ``by_sex``, and ``by_year`` flags, respectively.
+    By default, this counts cause-specific deaths and years of life lost over
+    the full course of the simulation. It can be configured to add or remove
+    stratification groups to the default groups defined by a
+    :class:ResultsStratifier.
 
     In the model specification, your configuration for this component should
     be specified as, e.g.:
@@ -30,98 +33,137 @@ class MortalityObserver:
     .. code-block:: yaml
 
         configuration:
-            metrics:
+            observers:
                 mortality:
-                    by_age: True
-                    by_year: False
-                    by_sex: True
-
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
     """
 
     configuration_defaults = {
-        "metrics": {
+        "observers": {
             "mortality": {
-                "by_age": False,
-                "by_year": False,
-                "by_sex": False,
+                "exclude": [],
+                "include": [],
             }
         }
     }
 
+    def __init__(self):
+        self.metrics_pipeline_name = "metrics"
+        self.tmrle_key = "population.theoretical_minimum_risk_life_expectancy"
+
+    def __repr__(self):
+        return "MortalityObserver()"
+
+    ##############
+    # Properties #
+    ##############
+
     @property
     def name(self):
         return "mortality_observer"
 
-    def setup(self, builder):
-        self.config = builder.configuration.metrics.mortality
-        self.clock = builder.time.clock()
-        self.step_size = builder.time.step_size()
-        self.start_time = self.clock()
-        self.initial_pop_entrance_time = self.start_time - self.step_size()
-        self.age_bins = get_age_bins(builder)
-        diseases = builder.components.get_components_by_type(
+    #################
+    # Setup methods #
+    #################
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder):
+        self.config = self._get_stratification_configuration(builder)
+        self.stratifier = self._get_stratifier(builder)
+        self._cause_components = self._get_cause_components(builder)
+        self.causes_of_death = ["other_causes"]
+        self.population_view = self._get_population_view(builder)
+
+        self.counts = Counter()
+
+        self._register_post_setup_listener(builder)
+        self._register_collect_metrics_listener(builder)
+        self._register_metrics_modifier(builder)
+
+    # noinspection PyMethodMayBeStatic
+    def _get_stratification_configuration(self, builder: Builder) -> ConfigTree:
+        return builder.configuration.observers.mortality
+
+    # noinspection PyMethodMayBeStatic
+    def _get_stratifier(self, builder: Builder) -> ResultsStratifier:
+        return builder.components.get_component(ResultsStratifier.name)
+
+    # noinspection PyMethodMayBeStatic
+    def _get_cause_components(
+        self, builder: Builder
+    ) -> List[Union[DiseaseState, RiskAttributableDisease]]:
+        return builder.components.get_components_by_type(
             (DiseaseState, RiskAttributableDisease)
         )
-        self.causes = [c.state_id for c in diseases] + ["other_causes"]
-
-        life_expectancy_data = builder.data.load(
-            "population.theoretical_minimum_risk_life_expectancy"
-        )
-        self.life_expectancy = builder.lookup.build_table(
-            life_expectancy_data, key_columns=[], parameter_columns=["age"]
-        )
 
+    # noinspection PyMethodMayBeStatic
+    def _get_population_view(self, builder: Builder) -> PopulationView:
         columns_required = [
             "tracked",
             "alive",
-            "entrance_time",
-            "exit_time",
-            "cause_of_death",
             "years_of_life_lost",
-            "age",
+            "cause_of_death",
+            "exit_time",
         ]
-        if self.config.by_sex:
-            columns_required += ["sex"]
-        self.population_view = builder.population.get_view(columns_required)
+        return builder.population.get_view(columns_required)
 
-        builder.value.register_value_modifier("metrics", self.metrics)
+    def _register_post_setup_listener(self, builder: Builder) -> None:
+        builder.event.register_listener("post_setup", self.on_post_setup)
 
-    def metrics(self, index, metrics):
-        pop = self.population_view.get(index)
-        pop.loc[pop.exit_time.isnull(), "exit_time"] = self.clock()
+    def _register_collect_metrics_listener(self, builder: Builder) -> None:
+        builder.event.register_listener("time_step", self.on_collect_metrics)
 
-        person_time = get_person_time(
-            pop, self.config.to_dict(), self.start_time, self.clock(), self.age_bins
-        )
-        deaths = get_deaths(
-            pop,
-            self.config.to_dict(),
-            self.start_time,
-            self.clock(),
-            self.age_bins,
-            self.causes,
-        )
-        ylls = get_years_of_life_lost(
-            pop,
-            self.config.to_dict(),
-            self.start_time,
-            self.clock(),
-            self.age_bins,
-            self.life_expectancy,
-            self.causes,
+    def _register_metrics_modifier(self, builder: Builder) -> None:
+        builder.value.register_value_modifier(
+            self.metrics_pipeline_name,
+            modifier=self.metrics,
+            requires_columns=["age", "exit_time", "alive"],
         )
 
-        metrics.update(person_time)
-        metrics.update(deaths)
-        metrics.update(ylls)
+    ########################
+    # Event-driven methods #
+    ########################
+
+    # noinspection PyUnusedLocal
+    def on_post_setup(self, event: Event) -> None:
+        self.causes_of_death += [
+            cause.state_id for cause in self._cause_components if cause.has_excess_mortality
+        ]
+
+    def on_collect_metrics(self, event: Event) -> None:
+        pop = self.population_view.get(event.index)
+        pop_died = pop[(pop["alive"] == "dead") & (pop["exit_time"] == event.time)]
+
+        groups = self.stratifier.group(
+            pop_died.index, self.config.include, self.config.exclude
+        )
+        for label, group_mask in groups:
+            for cause in self.causes_of_death:
+                cause_mask = pop_died["cause_of_death"] == cause
+                pop_died_of_cause = pop_died[group_mask & cause_mask]
+                new_observations = {
+                    f"death_due_to_{cause}_{label}": pop_died_of_cause.index.size,
+                    f"ylls_due_to_{cause}_{label}": pop_died_of_cause[
+                        "years_of_life_lost"
+                    ].sum(),
+                }
+                self.counts.update(new_observations)
+
+    ##################################
+    # Pipeline sources and modifiers #
+    ##################################
+
+    def metrics(self, index: pd.Index, metrics: Dict) -> Dict:
+        pop = self.population_view.get(index)
 
         the_living = pop[(pop.alive == "alive") & pop.tracked]
         the_dead = pop[pop.alive == "dead"]
-        metrics["years_of_life_lost"] = self.life_expectancy(the_dead.index).sum()
-        metrics["total_population_living"] = len(the_living)
-        metrics["total_population_dead"] = len(the_dead)
+        metrics["years_of_life_lost"] = the_dead["years_of_life_lost"].sum()
+        metrics["total_population_living"] = the_living.size
+        metrics["total_population_dead"] = the_dead.size
+        metrics.update(self.counts)
 
         return metrics
-
-    def __repr__(self):
-        return "MortalityObserver()"
diff --git a/src/vivarium_public_health/metrics/population.py b/src/vivarium_public_health/metrics/population.py
deleted file mode 100644
index 4b73a9f28..000000000
--- a/src/vivarium_public_health/metrics/population.py
+++ /dev/null
@@ -1,100 +0,0 @@
-"""
-===================
-Population Observer
-===================
-
-This module contains tools for observing population counts during the
-simulation.
-
-"""
-from collections import Counter
-
-import pandas as pd
-
-from vivarium_public_health.metrics.utilities import get_age_bins, get_population_counts
-
-
-class PopulationObserver:
-    """An observer for population counts.
-
-    By default, this counts the population at a particular sample date
-    annually. It can be configured to bin the population into age groups and
-    sexes by setting the ``by_age`` and ``by_sex`` flags, respectively. The
-    sample date defaults to July 1st.  This is also configurable.
-
-    The configuration defaults also contain a ``by_year`` flag, but it has
-    no behavior associated with it.
-
-    In the model specification, your configuration for this component should
-    be specified as, e.g.:
-
-    .. code-block:: yaml
-
-        configuration:
-            metrics:
-                population:
-                    by_age: True
-                    by_sex: True
-                    sample_date:
-                        month: 4
-                        day: 10
-
-    """
-
-    configuration_defaults = {
-        "metrics": {
-            "population": {
-                "by_age": False,
-                "by_year": False,
-                "by_sex": False,
-                "sample_date": {
-                    "month": 7,
-                    "day": 1,
-                },
-            }
-        }
-    }
-
-    @property
-    def name(self):
-        return "population_observer"
-
-    def setup(self, builder):
-        self.config = builder.configuration.metrics.population
-        self.clock = builder.time.clock()
-        self.age_bins = get_age_bins(builder)
-        self.population = Counter()
-
-        columns_required = ["tracked", "alive"]
-        if self.config.by_age:
-            columns_required += ["age"]
-        if self.config.by_sex:
-            columns_required += ["sex"]
-        self.population_view = builder.population.get_view(columns_required)
-
-        builder.event.register_listener("time_step__prepare", self.on_time_step_prepare)
-
-        builder.value.register_value_modifier("metrics", self.metrics)
-
-    def on_time_step_prepare(self, event):
-        pop = self.population_view.get(event.index)
-
-        if self.should_sample(event.time):
-            population_counts = get_population_counts(
-                pop, self.config.to_dict(), event.time, self.age_bins
-            )
-            self.population.update(population_counts)
-
-    def should_sample(self, event_time: pd.Timestamp) -> bool:
-        """Returns true if we should sample on this time step."""
-        sample_date = pd.Timestamp(
-            event_time.year, self.config.sample_date.month, self.config.sample_date.day
-        )
-        return self.clock() <= sample_date < event_time
-
-    def metrics(self, index, metrics):
-        metrics.update(self.population)
-        return metrics
-
-    def __repr__(self):
-        return "PopulationObserver()"
diff --git a/src/vivarium_public_health/metrics/risk.py b/src/vivarium_public_health/metrics/risk.py
index 3885f54ca..6692d670d 100644
--- a/src/vivarium_public_health/metrics/risk.py
+++ b/src/vivarium_public_health/metrics/risk.py
@@ -7,17 +7,17 @@
 
 """
 from collections import Counter
-from typing import Dict
+from typing import Dict, List
 
 import pandas as pd
+from vivarium.config_tree import ConfigTree
 from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
+from vivarium.framework.population import PopulationView
+from vivarium.framework.values import Pipeline
 
-from vivarium_public_health.metrics.utilities import (
-    get_age_bins,
-    get_prevalent_cases,
-    get_state_person_time,
-)
+from vivarium_public_health.metrics.stratification import ResultsStratifier
+from vivarium_public_health.utilities import to_years
 
 
 class CategoricalRiskObserver:
@@ -26,38 +26,29 @@ class CategoricalRiskObserver:
     Observes category person time for a risk factor.
 
     By default, this observer computes aggregate categorical person time
-    over the entire simulation.  It can be configured to bin these into
-    age_groups, sexes, and years by setting the ``by_age``, ``by_sex``,
-    and ``by_year`` flags, respectively.
+    over the full course of the simulation. It can be configured to add or
+    remove stratification groups to the default groups defined by a
+    ResultsStratifier.
 
-    This component can also observe the number of simulants in each age
-    group who are alive and in each category of risk at the specified
-    sample date each year (the sample date defaults to July, 1, and can
-    be set in the configuration).
-
-    Here is an example configuration to change the sample date to Dec. 31:
+    In the model specification, your configuration for this component should
+    be specified as, e.g.:
 
     .. code-block:: yaml
 
-        {risk_name}_observer:
-            sample_date:
-                month: 12
-                day: 31
+        configuration:
+            observers:
+                risk_name:
+                    exclude:
+                        - "sex"
+                    include:
+                        - "sample_stratification"
     """
 
     configuration_defaults = {
-        "metrics": {
+        "observers": {
             "risk": {
-                "by_age": False,
-                "by_year": False,
-                "by_sex": False,
-                "sample_exposure": {
-                    "sample": False,
-                    "date": {
-                        "month": 7,
-                        "day": 1,
-                    },
-                },
+                "exclude": [],
+                "include": [],
             }
         }
     }
@@ -67,85 +58,115 @@ def __init__(self, risk: str):
         Parameters
         ----------
         risk :
-        the type and name of a risk, specified as "type.name". Type is singular.
+        name of a risk
 
         """
         self.risk = risk
-        self.configuration_defaults = {
-            "metrics": {
-                f"{self.risk}": CategoricalRiskObserver.configuration_defaults["metrics"][
+        self.configuration_defaults = self._get_configuration_defaults()
+
+        self.exposure_pipeline_name = f"{self.risk}.exposure"
+        self.metrics_pipeline_name = "metrics"
+
+    def __repr__(self):
+        return f"CategoricalRiskObserver({self.risk})"
+
+    ##########################
+    # Initialization methods #
+    ##########################
+
+    # noinspection PyMethodMayBeStatic
+    def _get_configuration_defaults(self) -> Dict[str, Dict]:
+        return {
+            "observers": {
+                f"{self.risk}": CategoricalRiskObserver.configuration_defaults["observers"][
                     "risk"
                 ]
             }
         }
 
+    ##############
+    # Properties #
+    ##############
+
     @property
     def name(self):
         return f"categorical_risk_observer.{self.risk}"
 
+    #################
+    # Setup methods #
+    #################
+
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder):
-        self.data = {}
-        self.config = builder.configuration[f"metrics"][f"{self.risk}"]
-        self.clock = builder.time.clock()
-        self.categories = builder.data.load(f"risk_factor.{self.risk}.categories")
-        self.age_bins = get_age_bins(builder)
-        self.person_time = Counter()
-        self.sampled_exposure = Counter()
+        self.config = self._get_stratification_configuration(builder)
+        self.stratifier = self._get_stratifier(builder)
+        self.categories = self._get_categories(builder)
+        self.pipelines = self._get_pipelines(builder)
+        self.population_view = self._get_population_view(builder)
+
+        self.counts = Counter()
+
+        self._register_time_step_prepare_listener(builder)
+        self._register_metrics_modifier(builder)
+
+    def _get_stratification_configuration(self, builder: Builder) -> ConfigTree:
+        return builder.configuration.observers[self.risk]
 
+    # noinspection PyMethodMayBeStatic
+    def _get_stratifier(self, builder: Builder) -> ResultsStratifier:
+        return builder.components.get_component(ResultsStratifier.name)
+
+    def _get_categories(self, builder: Builder) -> List[str]:
+        return builder.data.load(f"risk_factor.{self.risk}.categories")
+
+    def _get_pipelines(self, builder: Builder) -> Dict[str, Pipeline]:
+        return {
+            self.exposure_pipeline_name: builder.value.get_value(self.exposure_pipeline_name)
+        }
+
+    # noinspection PyMethodMayBeStatic
+    def _get_population_view(self, builder: Builder) -> PopulationView:
         columns_required = ["alive"]
-        if self.config.by_age:
-            columns_required += ["age"]
-        if self.config.by_sex:
-            columns_required += ["sex"]
-        self.population_view = builder.population.get_view(columns_required)
-
-        self.exposure = builder.value.get_value(f"{self.risk}.exposure")
-        builder.value.register_value_modifier("metrics", self.metrics)
+        return builder.population.get_view(columns_required)
+
+    def _register_time_step_prepare_listener(self, builder: Builder) -> None:
+        # In order to get an accurate representation of person time we need to look at
+        # the state table before anything happens.
         builder.event.register_listener("time_step__prepare", self.on_time_step_prepare)
 
-    def on_time_step_prepare(self, event: Event):
-        pop = pd.concat(
-            [
-                self.population_view.get(event.index),
-                pd.Series(self.exposure(event.index), name=self.risk),
-            ],
-            axis=1,
+    def _register_metrics_modifier(self, builder: Builder) -> None:
+        builder.value.register_value_modifier(
+            self.metrics_pipeline_name,
+            modifier=self.metrics,
         )
 
-        for category in self.categories:
-            state_person_time_this_step = get_state_person_time(
-                pop,
-                self.config,
-                self.risk,
-                category,
-                self.clock().year,
-                event.step_size,
-                self.age_bins,
-            )
-            self.person_time.update(state_person_time_this_step)
-
-        if self._should_sample(event.time):
-            sampled_exposure = get_prevalent_cases(
-                pop, self.config.to_dict(), self.risk, event.time, self.age_bins
-            )
-            self.sampled_exposure.update(sampled_exposure)
-
-    def _should_sample(self, event_time: pd.Timestamp) -> bool:
-        """Returns true if we should sample on this time step."""
-        should_sample = self.config.sample_exposure.sample
-        if should_sample:
-            sample_date = pd.Timestamp(
-                year=event_time.year, **self.config.sample_prevalence.date.to_dict()
-            )
-            should_sample &= self.clock() <= sample_date < event_time
-        return should_sample
+    ########################
+    # Event-driven methods #
+    ########################
+
+    def on_time_step_prepare(self, event: Event) -> None:
+        step_size_in_years = to_years(event.step_size)
+        pop = self.population_view.get(
+            event.index, query="tracked == True and alive == 'alive'"
+        )
+        exposures = self.pipelines[self.exposure_pipeline_name](pop.index)
+        groups = self.stratifier.group(
+            exposures.index, self.config.include, self.config.exclude
+        )
+        for label, group_mask in groups:
+            for category in self.categories:
+                category_in_group_mask = group_mask & (exposures == category)
+                person_time_in_group = category_in_group_mask.sum() * step_size_in_years
+                new_observations = {
+                    f"{self.risk}_{category}_person_time_{label}": person_time_in_group
+                }
+                self.counts.update(new_observations)
+
+    ##################################
+    # Pipeline sources and modifiers #
+    ##################################
 
     # noinspection PyUnusedLocal
     def metrics(self, index: pd.Index, metrics: Dict) -> Dict:
-        metrics.update(self.person_time)
-        metrics.update(self.sampled_exposure)
+        metrics.update(self.counts)
         return metrics
-
-    def __repr__(self):
-        return f"CategoricalRiskObserver({self.risk})"
diff --git a/src/vivarium_public_health/metrics/stratification.py b/src/vivarium_public_health/metrics/stratification.py
new file mode 100644
index 000000000..baa0b8763
--- /dev/null
+++ b/src/vivarium_public_health/metrics/stratification.py
@@ -0,0 +1,362 @@
+"""
+==================
+Results Stratifier
+==================
+
+This module contains tools for stratifying observed quantities
+by specified characteristics.
+
+"""
+import itertools
+from dataclasses import dataclass
+from enum import Enum
+from typing import Callable, Dict, Iterable, List, Set, Tuple
+
+import pandas as pd
+from vivarium.framework.engine import Builder
+from vivarium.framework.event import Event
+from vivarium.framework.population import PopulationView
+
+
+class SourceType(Enum):
+    COLUMN = "column"
+    PIPELINE = "pipeline"
+    CLOCK = "clock"
+
+
+@dataclass
+class Source:
+    """
+    A source of information about simulants used to determine which category
+    they belong to for a given stratification level. The source name should be
+    the name of the column or pipeline being used as the source. If the source
+    is of type clock, the name should be something descriptive and unique.
+    """
+
+    name: str
+    type: SourceType
+
+
+@dataclass
+class StratificationLevel:
+    """
+    A level of stratification. Each StratificationLevel represents a set of
+    mutually exclusive and collectively exhaustive categories into which
+    simulants can be assigned.
+
+    The mapper is a function which is applied to the list of sources and
+    assigns each simulant to specific category. It will throw an error if its
+    output is not one of the stratification's categories. By default, the
+    mapper assumes the StratificationLevel has a single source and returns its
+    value as the category.
+
+    A current category getter can be defined if it is known that certain
+    categories are not possible during specific times. Removing these
+    categories from iteration can provide a performance improvement in these
+    cases. The primary use case for the current category getter is for
+    stratification by time, and in particular stratification by year. By
+    default, this will return all categories.
+    """
+
+    name: str
+    sources: List[Source]
+    categories: Set[str]
+    mapper: Callable[[pd.Series], str] = None
+    current_categories_getter: Callable[[], Set[str]] = None
+
+    def __post_init__(self):
+        self._set_mapper()
+        self._set_current_categories_getter()
+
+    @property
+    def current_categories(self) -> Set[str]:
+        return self.current_categories_getter()
+
+    def _set_mapper(self) -> None:
+        """
+        Sets the mapper to be the identity mapper if no mapper has been
+        provided. Additionally, wraps the mapper so that it will throw a
+        ValueError if the mapper produces an output that is not in the set of
+        this StratificationLevel's categories.
+        """
+        name = self.name
+        categories = self.categories
+        mapper = self.mapper if self.mapper else self._default_mapper
+
+        def wrapped_mapper(row: pd.Series) -> str:
+            category = mapper(row)
+            if category not in categories:
+                raise ValueError(f"Invalid value '{category}' found in {name}.")
+            return category
+
+        self.mapper = wrapped_mapper
+
+    def _set_current_categories_getter(self) -> None:
+        self.current_categories_getter = (
+            self.current_categories_getter
+            if self.current_categories_getter
+            else self._default_current_categories_getter
+        )
+
+    ###############################
+    # Default getters and mappers #
+    ###############################
+
+    # noinspection PyMethodMayBeStatic
+    def _default_mapper(self, row: pd.Series) -> str:
+        return str(row[0])
+
+    def _default_current_categories_getter(self) -> Set[str]:
+        return self.categories
+
+
+class ResultsStratifier:
+    """Centralized component for handling results stratification.
+
+    This component manages the assignment of simulants to groups for the
+    purpose of stratification. Each observer component will get a reference to
+    this component so that it can properly stratify its output.
+    """
+
+    name = "results_stratifier"
+
+    configuration_defaults = {
+        "observers": {
+            "default": [],
+        }
+    }
+
+    def __init__(self):
+        self.metrics_pipeline_name = "metrics"
+        self.tmrle_key = "population.theoretical_minimum_risk_life_expectancy"
+
+    #################
+    # Setup methods #
+    #################
+
+    # noinspection PyAttributeOutsideInit
+    def setup(self, builder: Builder) -> None:
+        """Perform this component's setup."""
+        self.clock = builder.time.clock()
+        self.default_stratification_levels = self._get_default_stratification_levels(builder)
+        self.pipelines = {}
+        self.columns_required = {"tracked"}
+        self.clock_sources = set()
+        self.stratification_levels: Dict[str, StratificationLevel] = {}
+        self.stratification_groups: pd.DataFrame = None
+
+        self.age_bins = self._get_age_bins(builder)
+
+        self.register_stratifications(builder)
+        self.population_view = self._get_population_view(builder)
+
+        self._register_timestep_prepare_listener(builder)
+
+    # noinspection PyMethodMayBeStatic
+    def _get_default_stratification_levels(self, builder: Builder) -> Set[str]:
+        return set(builder.configuration.observers.default)
+
+    def register_stratifications(self, builder: Builder) -> None:
+        """Register each desired stratification with calls to _setup_stratification"""
+
+        start_year = builder.configuration.time.start.year
+        end_year = builder.configuration.time.end.year
+
+        self.setup_stratification(
+            builder,
+            name=ResultsStratifier.YEAR,
+            sources=[ResultsStratifier.YEAR_SOURCE],
+            categories={str(year) for year in range(start_year, end_year + 1)},
+            mapper=self.year_stratification_mapper,
+            current_category_getter=self.year_current_categories_getter,
+        )
+
+        self.setup_stratification(
+            builder,
+            name=ResultsStratifier.SEX,
+            sources=[ResultsStratifier.SEX_SOURCE],
+            categories=ResultsStratifier.SEX_CATEGORIES,
+        )
+
+        self.setup_stratification(
+            builder,
+            name=ResultsStratifier.AGE,
+            sources=[ResultsStratifier.AGE_SOURCE],
+            categories={age_bin for age_bin in self.age_bins["age_group_name"]},
+            mapper=self.age_stratification_mapper,
+        )
+
+    # noinspection PyMethodMayBeStatic
+    def _get_population_view(self, builder: Builder) -> PopulationView:
+        return builder.population.get_view(list(self.columns_required))
+
+    # noinspection PyMethodMayBeStatic
+    def _get_age_bins(self, builder: Builder) -> pd.DataFrame:
+        raw_age_bins = builder.data.load("population.age_bins")
+        age_start = builder.configuration.population.age_start
+        exit_age = builder.configuration.population.exit_age
+
+        age_start_mask = age_start < raw_age_bins["age_end"]
+        exit_age_mask = raw_age_bins["age_start"] < exit_age if exit_age else True
+
+        age_bins = raw_age_bins.loc[age_start_mask & exit_age_mask, :]
+        return age_bins
+
+    def _register_timestep_prepare_listener(self, builder: Builder) -> None:
+        builder.event.register_listener(
+            "time_step__prepare", self.on_time_step_prepare, priority=0
+        )
+
+    ########################
+    # Event-driven methods #
+    ########################
+
+    def on_time_step_prepare(self, event: Event) -> None:
+        """Determine each simulant's category for each stratification level"""
+        pop = self.population_view.get(
+            event.index, query='tracked == True and alive == "alive"'
+        )
+        pipeline_values = [
+            pd.Series(pipeline(pop.index), name=name)
+            for name, pipeline in self.pipelines.items()
+        ]
+        clock_values = [
+            pd.Series(self.clock(), index=pop.index, name=name) for name in self.clock_sources
+        ]
+        sources = pd.concat([pop] + pipeline_values + clock_values, axis=1)
+
+        stratification_groups = [
+            sources[[source.name for source in stratification_level.sources]]
+            .apply(stratification_level.mapper, axis=1)
+            .rename(stratification_level.name)
+            for stratification_level in self.stratification_levels.values()
+        ]
+        # noinspection PyAttributeOutsideInit
+        self.stratification_groups = pd.concat(stratification_groups, axis=1)
+
+    ##################
+    # Public methods #
+    ##################
+
+    # todo add caching of stratifications
+    def group(
+        self, index: pd.Index, include: Iterable[str], exclude: Iterable[str]
+    ) -> Iterable[Tuple[str, pd.Series]]:
+        """Takes a full population index and yields stratified subgroups.
+
+        Parameters
+        ----------
+        index
+            The index of the population to stratify.
+        include
+            List of stratifications to add to the default stratifications
+        exclude
+            List of stratifications to remove from the default stratifications
+
+        Yields
+        ------
+            A tuple of stratification labels and the population subgroup
+            corresponding to those labels.
+
+        """
+        index = index.intersection(self.stratification_groups.index)
+        stratification_groups = self.stratification_groups.loc[index]
+
+        for stratification in self._get_current_stratifications(include, exclude):
+            stratification_key = self._get_stratification_key(stratification)
+            group_mask = True
+            if not index.empty:
+                for level, category in stratification:
+                    group_mask &= stratification_groups[level.name] == category
+            yield stratification_key, group_mask
+
+    ##################
+    # Helper methods #
+    ##################
+
+    def setup_stratification(
+        self,
+        builder: Builder,
+        name: str,
+        sources: List[Source],
+        categories: Set[str],
+        mapper: Callable[[pd.Series], str] = None,
+        current_category_getter: Callable[[], Set[str]] = None,
+    ) -> None:
+        """
+        Defines the characteristics of a stratification level and ensures that
+        each of its sources is available to the ResultsStratifier
+        """
+        stratification_level = StratificationLevel(
+            name, sources, categories, mapper, current_category_getter
+        )
+        self.stratification_levels[name] = stratification_level
+
+        for source in sources:
+            if source.type == SourceType.PIPELINE:
+                self.pipelines[source.name] = builder.value.get_value(source.name)
+            elif source.type == SourceType.COLUMN:
+                self.columns_required.add(source.name)
+            elif source.type == SourceType.CLOCK:
+                self.clock_sources.add(source.name)
+            else:
+                raise ValueError(f"Invalid stratification source type '{source.type}'.")
+
+    def _get_current_stratifications(
+        self, include: Iterable[str], exclude: Iterable[str]
+    ) -> List[Tuple[Tuple[StratificationLevel, str], ...]]:
+        """
+        Gets all stratification combinations. Returns a List of
+        Stratifications. Each Stratification is represented as a Tuple of
+        Levels. Each Level is represented as a Tuple of a StratificationLevel
+        object and string referring to the specific stratification category.
+
+        If no stratification levels are defined, returns a List with a single empty Tuple
+        """
+        include = set(include)
+        exclude = set(exclude)
+        level_names = (self.default_stratification_levels | include) - exclude
+        groups = [
+            [(level, category) for category in level.current_categories]
+            for level_name, level in self.stratification_levels.items()
+            if level_name in level_names
+        ]
+        # Get product of all stratification combinations
+        return list(itertools.product(*groups))
+
+    @staticmethod
+    def _get_stratification_key(
+        stratification: Iterable[Tuple[StratificationLevel, str]]
+    ) -> str:
+        return (
+            "_".join([f"{level[0].name}_{level[1]}" for level in stratification])
+            .replace(" ", "_")
+            .lower()
+        )
+
+    ##########################
+    # Stratification Details #
+    ##########################
+
+    AGE = "age"
+    AGE_SOURCE = Source("age", SourceType.COLUMN)
+
+    def age_stratification_mapper(self, row: pd.Series) -> str:
+        age_group_mask = (
+            self.age_bins["age_start"] <= row[ResultsStratifier.AGE_SOURCE.name]
+        ) & (row[ResultsStratifier.AGE_SOURCE.name] < self.age_bins["age_end"])
+        return str(self.age_bins.loc[age_group_mask, "age_group_name"].squeeze())
+
+    SEX = "sex"
+    SEX_SOURCE = Source("sex", SourceType.COLUMN)
+    SEX_CATEGORIES = {"Female", "Male"}
+
+    YEAR = "year"
+    YEAR_SOURCE = Source("year", SourceType.CLOCK)
+
+    # noinspection PyMethodMayBeStatic
+    def year_stratification_mapper(self, row: pd.Series) -> str:
+        return str(row[0].year)
+
+    def year_current_categories_getter(self) -> Set[str]:
+        return {str(self.clock().year)}
diff --git a/src/vivarium_public_health/metrics/utilities.py b/src/vivarium_public_health/metrics/utilities.py
deleted file mode 100644
index 71e9c385f..000000000
--- a/src/vivarium_public_health/metrics/utilities.py
+++ /dev/null
@@ -1,762 +0,0 @@
-"""
-=================
-Metrics Utilities
-=================
-
-This module contains shared utilities for querying, parsing, and transforming
-simulation data to support particular observations during the simulation.
-
-"""
-from collections import ChainMap
-from string import Template
-from typing import Callable, Dict, List, Tuple, Union
-
-import numpy as np
-import pandas as pd
-from vivarium.framework.lookup import LookupTable
-from vivarium.framework.values import Pipeline
-
-from vivarium_public_health.utilities import to_years
-
-_MIN_AGE = 0.0
-_MAX_AGE = 150.0
-_MIN_YEAR = 1900
-_MAX_YEAR = 2100
-
-
-class QueryString(str):
-    """Convenience string that forms logical conjunctions using addition.
-
-    This class is meant to be used to create a logical statement for
-    use with pandas ``query`` functions. It hides away the management
-    of conjunctions and the fence posting problems that management creates.
-
-    Examples
-    --------
-    >>> from vivarium_public_health.metrics.utilities import QueryString
-    >>> s = QueryString('')
-    >>> s
-    ''
-    >>> s + ''
-    ''
-    >>> s + 'abc'
-    'abc'
-    >>> s += 'abc'
-    >>> s + 'def'
-    'abc and def'
-
-    """
-
-    def __add__(self, other: Union[str, "QueryString"]) -> "QueryString":
-        if self:
-            if other:
-                return QueryString(str(self) + " and " + str(other))
-            else:
-                return self
-        else:
-            return QueryString(other)
-
-    def __radd__(self, other: Union[str, "QueryString"]) -> "QueryString":
-        return QueryString(other) + self
-
-
-class SubstituteString(str):
-    """Normal string plus a no-op substitute method.
-
-    Meant to be used with the OutputTemplate.
-
-    """
-
-    def substitute(self, *_, **__) -> "SubstituteString":
-        """No-op method for consistency with OutputTemplate."""
-        return self
-
-
-class OutputTemplate(Template):
-    """Output string template that enforces standardized formatting."""
-
-    @staticmethod
-    def format_template_value(value):
-        """Formatting helper method for substituting values into a template."""
-        return str(value).replace(" ", "_").lower()
-
-    @staticmethod
-    def get_mapping(*args, **kws):
-        """Gets a consistent mapping from args passed to substitute."""
-        # This is copied directly from the first part of Template.substitute
-        if not args:
-            raise TypeError(
-                "descriptor 'substitute' of 'Template' object " "needs an argument"
-            )
-        self, *args = args  # allow the "self" keyword be passed
-        if len(args) > 1:
-            raise TypeError("Too many positional arguments")
-        if not args:
-            mapping = dict(kws)
-        elif kws:
-            mapping = ChainMap(kws, args[0])
-        else:
-            mapping = args[0]
-        return self, mapping
-
-    def substitute(*args, **kws) -> Union[SubstituteString, "OutputTemplate"]:
-        """Substitutes provided values into the template.
-
-        Users are allowed to pass any dictionary like object whose keys match
-        placeholders in the template. Alternatively, they can provide
-        keyword arguments where the keywords are the placeholders. If both
-        are provided, the keywords take precedence.
-
-        Returns
-        -------
-        Union[SubstituteString, OutputTemplate]
-            Another output template with the provided values substituted in
-            to the template placeholders if any placeholders remain,
-            otherwise a completed ``SubstituteString``.
-        """
-        self, mapping = OutputTemplate.get_mapping(*args, **kws)
-        mapping = {key: self.format_template_value(value) for key, value in mapping.items()}
-        try:
-            return SubstituteString(super(OutputTemplate, self).substitute(mapping))
-        except KeyError:
-            return OutputTemplate(super(OutputTemplate, self).safe_substitute(mapping))
-
-    def safe_substitute(*args, **kws):
-        """Alias to OutputTemplate.substitute."""
-        self, mapping = OutputTemplate.get_mapping(*args, **kws)
-        return self.substitute(mapping)
-
-    def __repr__(self):
-        return super(OutputTemplate, self).safe_substitute()
-
-
-def get_age_bins(builder) -> pd.DataFrame:
-    """Retrieves age bins relevant to the current simulation.
-
-    Parameters
-    ----------
-    builder
-        The simulation builder.
-
-    Returns
-    -------
-        DataFrame with columns ``age_group_name``, ``age_start``,
-        and ``age_end``.
-
-    """
-    age_bins = builder.data.load("population.age_bins")
-
-    # Works based on the fact that currently only models with age_start = 0 can include fertility
-    age_start = builder.configuration.population.age_start
-    min_bin_start = age_bins.age_start[np.asscalar(np.digitize(age_start, age_bins.age_end))]
-    age_bins = age_bins[age_bins.age_start >= min_bin_start]
-    age_bins.loc[age_bins.age_start < age_start, "age_start"] = age_start
-
-    exit_age = builder.configuration.population.exit_age
-    if exit_age:
-        age_bins = age_bins[age_bins.age_start < exit_age]
-        age_bins.loc[age_bins.age_end > exit_age, "age_end"] = exit_age
-    return age_bins
-
-
-def get_output_template(by_age: bool, by_sex: bool, by_year: bool, **_) -> OutputTemplate:
-    """Gets a template string for output metrics.
-
-    The template string should be filled in using filter criteria for
-    measure, age, sex, and year in the observer using this function.
-
-    Parameters
-    ----------
-    by_age
-        Whether the template should include age criteria.
-    by_sex
-        Whether the template should include sex criteria.
-    by_year
-        Whether the template should include year criteria.
-
-    Returns
-    -------
-    OutputTemplate
-        A template string with measure and possibly additional criteria.
-
-    """
-    template = "${measure}"
-    if by_year:
-        template += "_in_${year}"
-    if by_sex:
-        template += "_among_${sex}"
-    if by_age:
-        template += "_in_age_group_${age_group}"
-    return OutputTemplate(template)
-
-
-def get_age_sex_filter_and_iterables(
-    config: Dict[str, bool], age_bins: pd.DataFrame, in_span: bool = False
-) -> Tuple[QueryString, Tuple[List[Tuple[str, pd.Series]], List[str]]]:
-    """Constructs a filter and a set of iterables for age and sex.
-
-    The constructed filter and iterables are based on configuration for the
-    observer component.
-
-    Parameters
-    ----------
-    config
-        A mapping with 'by_age' and 'by_sex' keys and boolean values
-        indicating whether the observer is binning data by the respective
-        categories.
-    age_bins
-        A table containing bin names and bin edges.
-    in_span
-        Whether the filter for age corresponds to living through an age
-        group in a time span or uses a point estimate of age at a particular
-        point in time.
-
-    Returns
-    -------
-    QueryString
-        A filter on age and sex for use with DataFrame.query
-    Tuple[List[Tuple[str, pd.Series]], List[str]]
-        Iterables for the age and sex groups partially defining the bins
-        for the observers.
-
-    """
-    age_sex_filter = QueryString("")
-    if config["by_age"]:
-        ages = list(age_bins.set_index("age_group_name").iterrows())
-        if in_span:
-            age_sex_filter += (
-                "{age_start} < age_at_span_end and age_at_span_start < {age_end}"
-            )
-        else:
-            age_sex_filter += "{age_start} <= age and age < {age_end}"
-    else:
-        ages = [("all_ages", pd.Series({"age_start": _MIN_AGE, "age_end": _MAX_AGE}))]
-
-    if config["by_sex"]:
-        sexes = ["Male", "Female"]
-        age_sex_filter += 'sex == "{sex}"'
-    else:
-        sexes = ["Both"]
-
-    return age_sex_filter, (ages, sexes)
-
-
-def get_time_iterable(
-    config: Dict[str, bool], sim_start: pd.Timestamp, sim_end: pd.Timestamp
-) -> List[Tuple[str, Tuple[pd.Timestamp, pd.Timestamp]]]:
-    """Constructs an iterable for time bins.
-
-    The constructed iterable are based on configuration for the observer
-    component.
-
-    Parameters
-    ----------
-    config
-        A mapping with 'by_year' and a boolean value indicating whether
-        the observer is binning data by year.
-    sim_start
-        The time the simulation starts.
-    sim_end
-        The time the simulation ends.
-
-    Returns
-    -------
-    List[Tuple[str, Tuple[pandas.Timestamp, pandas.Timestamp]]]
-        Iterable for the time groups partially defining the bins
-        for the observers.
-
-    """
-    if config["by_year"]:
-        time_spans = [
-            (year, (pd.Timestamp(f"1-1-{year}"), pd.Timestamp(f"1-1-{year + 1}")))
-            for year in range(sim_start.year, sim_end.year + 1)
-        ]
-    else:
-        time_spans = [
-            (
-                "all_years",
-                (pd.Timestamp(f"1-1-{_MIN_YEAR}"), pd.Timestamp(f"1-1-{_MAX_YEAR}")),
-            )
-        ]
-    return time_spans
-
-
-def get_group_counts(
-    pop: pd.DataFrame,
-    base_filter: str,
-    base_key: OutputTemplate,
-    config: Dict[str, bool],
-    age_bins: pd.DataFrame,
-    aggregate: Callable[[pd.DataFrame], Union[int, float]] = len,
-) -> Dict[Union[SubstituteString, OutputTemplate], Union[int, float]]:
-    """Gets a count of people in a custom subgroup.
-
-    The user is responsible for providing a default filter (e.g. only alive
-    people, or people susceptible to a particular disease).  Demographic
-    filters will be applied based on standardized configuration.
-
-    Parameters
-    ----------
-    pop
-        The population dataframe to be counted.  It must contain sufficient
-        columns for any necessary filtering (e.g. the ``age`` column if
-        filtering by age).
-    base_filter
-        A base filter term (e.g.: alive, susceptible to a particular disease)
-        formatted to work with the query method of the provided population
-        dataframe.
-    base_key
-        A template string with replaceable fields corresponding to the
-        requested filters.
-    config
-        A dict with ``by_age`` and ``by_sex`` keys and boolean values.
-    age_bins
-        A dataframe with ``age_start`` and ``age_end`` columns.
-    aggregate
-        Transformation used to produce the aggregate.
-
-    Returns
-    -------
-    Dict[Union[SubstituteString, OutputTemplate], Union[int, float]]
-        A dictionary of output_key, count pairs where the output key is a
-        string or template representing the sub groups.
-    """
-    age_sex_filter, (ages, sexes) = get_age_sex_filter_and_iterables(config, age_bins)
-    base_filter += age_sex_filter
-
-    group_counts = {}
-
-    for group, age_group in ages:
-        start, end = age_group.age_start, age_group.age_end
-        for sex in sexes:
-            filter_kwargs = {
-                "age_start": start,
-                "age_end": end,
-                "sex": sex,
-                "age_group": group,
-            }
-            group_key = base_key.substitute(**filter_kwargs)
-            group_filter = base_filter.format(**filter_kwargs)
-            in_group = pop.query(group_filter) if group_filter and not pop.empty else pop
-
-            group_counts[group_key] = aggregate(in_group)
-
-    return group_counts
-
-
-def get_susceptible_person_time(pop, config, disease, current_year, step_size, age_bins):
-    base_key = get_output_template(**config).substitute(
-        measure=f"{disease}_susceptible_person_time", year=current_year
-    )
-    base_filter = QueryString(f'alive == "alive" and {disease} == "susceptible_to_{disease}"')
-    person_time = get_group_counts(
-        pop,
-        base_filter,
-        base_key,
-        config,
-        age_bins,
-        aggregate=lambda x: len(x) * to_years(step_size),
-    )
-    return person_time
-
-
-def get_disease_event_counts(pop, config, disease, event_time, age_bins):
-    base_key = get_output_template(**config).substitute(
-        measure=f"{disease}_counts", year=event_time.year
-    )
-    # Can't use query with time stamps, so filter
-    pop = pop.loc[pop[f"{disease}_event_time"] == event_time]
-    base_filter = QueryString("")
-    return get_group_counts(pop, base_filter, base_key, config, age_bins)
-
-
-def get_prevalent_cases(pop, config, disease, event_time, age_bins):
-    config = config.copy()
-    config["by_year"] = True  # This is always an annual point estimate
-    base_key = get_output_template(**config).substitute(
-        measure=f"{disease}_prevalent_cases", year=event_time.year
-    )
-    base_filter = QueryString(f'alive == "alive" and {disease} != "susceptible_to_{disease}"')
-    return get_group_counts(pop, base_filter, base_key, config, age_bins)
-
-
-def get_population_counts(pop, config, event_time, age_bins):
-    config = config.copy()
-    config["by_year"] = True  # This is always an annual point estimate
-    base_key = get_output_template(**config).substitute(
-        measure=f"population_count", year=event_time.year
-    )
-    base_filter = QueryString(f'alive == "alive"')
-    return get_group_counts(pop, base_filter, base_key, config, age_bins)
-
-
-def get_person_time(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    sim_start: pd.Timestamp,
-    sim_end: pd.Timestamp,
-    age_bins: pd.DataFrame,
-) -> Dict[str, float]:
-    base_key = get_output_template(**config).substitute(measure="person_time")
-    base_filter = QueryString("")
-    time_spans = get_time_iterable(config, sim_start, sim_end)
-
-    person_time = {}
-    for year, (t_start, t_end) in time_spans:
-        year_key = base_key.substitute(year=year)
-        lived_in_span = get_lived_in_span(pop, t_start, t_end)
-        person_time_in_span = get_person_time_in_span(
-            lived_in_span, base_filter, year_key, config, age_bins
-        )
-        person_time.update(person_time_in_span)
-    return person_time
-
-
-def get_lived_in_span(
-    pop: pd.DataFrame, t_start: pd.Timestamp, t_end: pd.Timestamp
-) -> pd.DataFrame:
-    """Gets a subset of the population that lived in the time span.
-
-    Parameters
-    ----------
-    pop
-        A table representing the population with columns for 'entrance_time',
-        'exit_time' and 'age'.
-    t_start
-        The date and time at the start of the span.
-    t_end
-        The date and time at the end of the span.
-
-    Returns
-    -------
-    pandas.DataFrame
-        A table representing the population who lived some amount of time
-        within the time span with all columns provided in the original
-        table and additional columns 'age_at_span_start' and 'age_at_span_end'
-        indicating the age of the individual at the start and end of the time
-        span, respectively. 'age_at_span_start' will never be lower than the
-        age at the simulant's entrance time and 'age_at_span_end' will never
-        be greater than the age at the simulant's exit time.
-
-    """
-    lived_in_span = pop.loc[(t_start < pop["exit_time"]) & (pop["entrance_time"] < t_end)]
-
-    span_entrance_time = lived_in_span.entrance_time.copy()
-    span_entrance_time.loc[t_start > span_entrance_time] = t_start
-    span_exit_time = lived_in_span.exit_time.copy()
-    span_exit_time.loc[t_end < span_exit_time] = t_end
-
-    lived_in_span.loc[:, "age_at_span_end"] = lived_in_span.age - to_years(
-        lived_in_span.exit_time - span_exit_time
-    )
-    lived_in_span.loc[:, "age_at_span_start"] = lived_in_span.age - to_years(
-        lived_in_span.exit_time - span_entrance_time
-    )
-    return lived_in_span
-
-
-def get_person_time_in_span(
-    lived_in_span: pd.DataFrame,
-    base_filter: QueryString,
-    span_key: OutputTemplate,
-    config: Dict[str, bool],
-    age_bins: pd.DataFrame,
-) -> Dict[Union[SubstituteString, OutputTemplate], float]:
-    """Counts the amount of person time lived in a particular time span.
-
-    Parameters
-    ----------
-    lived_in_span
-        A table representing the subset of the population who lived in a
-        particular time span. Must have columns for 'age_at_span_start' and
-        'age_at_span_end'.
-    base_filter
-        A base filter term (e.g.: alive, susceptible to a particular disease)
-        formatted to work with the query method of the provided population
-        dataframe.
-    span_key
-        A template string with replaceable fields corresponding to the
-        requested filters.
-    config
-        A dict with ``by_age`` and ``by_sex`` keys and boolean values.
-    age_bins
-        A dataframe with ``age_start`` and ``age_end`` columns.
-
-    Returns
-    -------
-    Dict[Union[SubstituteString, OutputTemplate], float]
-        A dictionary of output_key, person_time pairs where the output key
-        corresponds to a particular demographic group.
-    """
-    person_time = {}
-    age_sex_filter, (ages, sexes) = get_age_sex_filter_and_iterables(
-        config, age_bins, in_span=True
-    )
-    base_filter += age_sex_filter
-
-    for group, age_bin in ages:
-        a_start, a_end = age_bin.age_start, age_bin.age_end
-        for sex in sexes:
-            filter_kwargs = {
-                "sex": sex,
-                "age_start": a_start,
-                "age_end": a_end,
-                "age_group": group,
-            }
-
-            key = span_key.substitute(**filter_kwargs)
-            group_filter = base_filter.format(**filter_kwargs)
-
-            in_group = (
-                lived_in_span.query(group_filter) if group_filter else lived_in_span.copy()
-            )
-            age_start = np.maximum(in_group.age_at_span_start, a_start)
-            age_end = np.minimum(in_group.age_at_span_end, a_end)
-
-            person_time[key] = (age_end - age_start).sum()
-
-    return person_time
-
-
-def get_deaths(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    sim_start: pd.Timestamp,
-    sim_end: pd.Timestamp,
-    age_bins: pd.DataFrame,
-    causes: List[str],
-) -> Dict[str, int]:
-    """Counts the number of deaths by cause.
-
-    Parameters
-    ----------
-    pop
-        The population dataframe to be counted. It must contain sufficient
-        columns for any necessary filtering (e.g. the ``age`` column if
-        filtering by age).
-    config
-        A dict with ``by_age``, ``by_sex``, and ``by_year`` keys and
-        boolean values.
-    sim_start
-        The simulation start time.
-    sim_end
-        The simulation end time.
-    age_bins
-        A dataframe with ``age_start`` and ``age_end`` columns.
-    causes
-        List of causes present in the simulation.
-
-    Returns
-    -------
-    Dict[str, int]
-        A dictionary of output_key, death_count pairs where the output_key
-        represents a particular demographic subgroup.
-
-    """
-    base_filter = QueryString('alive == "dead" and cause_of_death == "death_due_to_{cause}"')
-    base_key = get_output_template(**config)
-    pop = clean_cause_of_death(pop)
-
-    time_spans = get_time_iterable(config, sim_start, sim_end)
-
-    deaths = {}
-    for year, (t_start, t_end) in time_spans:
-        died_in_span = pop[(t_start <= pop.exit_time) & (pop.exit_time < t_end)]
-        for cause in causes:
-            cause_year_key = base_key.substitute(measure=f"death_due_to_{cause}", year=year)
-            cause_filter = base_filter.format(cause=cause)
-            group_deaths = get_group_counts(
-                died_in_span, cause_filter, cause_year_key, config, age_bins
-            )
-            deaths.update(group_deaths)
-    return deaths
-
-
-def get_years_of_life_lost(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    sim_start: pd.Timestamp,
-    sim_end: pd.Timestamp,
-    age_bins: pd.DataFrame,
-    life_expectancy: LookupTable,
-    causes: List[str],
-) -> Dict[str, float]:
-    """Counts the years of life lost by cause.
-
-    Parameters
-    ----------
-    pop
-        The population dataframe to be counted. It must contain sufficient
-        columns for any necessary filtering (e.g. the ``age`` column if
-        filtering by age).
-    config
-        A dict with ``by_age``, ``by_sex``, and ``by_year`` keys and
-        boolean values.
-    sim_start
-        The simulation start time.
-    sim_end
-        The simulation end time.
-    age_bins
-        A dataframe with ``age_start`` and ``age_end`` columns.
-    life_expectancy
-        A lookup table that takes in a pandas index and returns the life
-        expectancy of the each individual represented by the index.
-    causes
-        List of causes present in the simulation.
-
-    Returns
-    -------
-    Dict[str, float]
-        A dictionary of output_key, yll_count pairs where the output_key
-        represents a particular demographic subgroup.
-
-    """
-    base_filter = QueryString('alive == "dead" and cause_of_death == "death_due_to_{cause}"')
-    base_key = get_output_template(**config)
-    pop = clean_cause_of_death(pop)
-
-    time_spans = get_time_iterable(config, sim_start, sim_end)
-
-    years_of_life_lost = {}
-    for year, (t_start, t_end) in time_spans:
-        died_in_span = pop[(t_start <= pop.exit_time) & (pop.exit_time < t_end)]
-        for cause in causes:
-            cause_year_key = base_key.substitute(measure=f"ylls_due_to_{cause}", year=year)
-            cause_filter = base_filter.format(cause=cause)
-            group_ylls = get_group_counts(
-                died_in_span,
-                cause_filter,
-                cause_year_key,
-                config,
-                age_bins,
-                aggregate=lambda subgroup: sum(life_expectancy(subgroup.index)),
-            )
-            years_of_life_lost.update(group_ylls)
-    return years_of_life_lost
-
-
-def get_years_lived_with_disability(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    current_year: int,
-    step_size: pd.Timedelta,
-    age_bins: pd.DataFrame,
-    disability_weights: Dict[str, Pipeline],
-    causes: List[str],
-) -> Dict[str, float]:
-    """Counts the years lived with disability by cause in the time step.
-
-    Parameters
-    ----------
-    pop
-        The population dataframe to be counted. It must contain sufficient
-        columns for any necessary filtering (e.g. the ``age`` column if
-        filtering by age).
-    config
-        A dict with ``by_age``, ``by_sex``, and ``by_year`` keys and
-        boolean values.
-    current_year
-        The current year in the simulation.
-    step_size
-        The size of the current time step.
-    age_bins
-        A dataframe with ``age_start`` and ``age_end`` columns.
-    disability_weights
-        A mapping between causes and their disability weight pipelines.
-    causes
-        List of causes present in the simulation.
-
-    Returns
-    -------
-    Dict[str, float]
-        A dictionary of output_key, yld_count pairs where the output_key
-        represents a particular demographic subgroup.
-
-    """
-    base_key = get_output_template(**config).substitute(year=current_year)
-    base_filter = QueryString('alive == "alive"')
-
-    years_lived_with_disability = {}
-    for cause in causes:
-        cause_key = base_key.substitute(measure=f"ylds_due_to_{cause}")
-
-        def count_ylds(sub_group):
-            """Counts ylds attributable to a cause in the time step."""
-            return sum(disability_weights[cause](sub_group.index) * to_years(step_size))
-
-        group_ylds = get_group_counts(
-            pop, base_filter, cause_key, config, age_bins, aggregate=count_ylds
-        )
-        years_lived_with_disability.update(group_ylds)
-
-    return years_lived_with_disability
-
-
-def clean_cause_of_death(pop: pd.DataFrame) -> pd.DataFrame:
-    """Standardizes cause of death names to all read ``death_due_to_cause``."""
-
-    def _clean(cod: str) -> str:
-        if "death" in cod or "dead" in cod:
-            pass
-        else:
-            cod = f"death_due_to_{cod}"
-        return cod
-
-    pop.cause_of_death = pop.cause_of_death.apply(_clean)
-    return pop
-
-
-def get_state_person_time(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    state_machine: str,
-    state: str,
-    current_year: Union[str, int],
-    step_size: pd.Timedelta,
-    age_bins: pd.DataFrame,
-) -> Dict[str, float]:
-    """Custom person time getter that handles state column name assumptions"""
-    base_key = get_output_template(**config).substitute(
-        measure=f"{state}_person_time", year=current_year
-    )
-    base_filter = QueryString(f'alive == "alive" and {state_machine} == "{state}"')
-    person_time = get_group_counts(
-        pop,
-        base_filter,
-        base_key,
-        config,
-        age_bins,
-        aggregate=lambda x: len(x) * to_years(step_size),
-    )
-    return person_time
-
-
-class TransitionString(str):
-    def __new__(cls, value):
-        # noinspection PyArgumentList
-        obj = str.__new__(cls, value.lower())
-        obj.from_state, obj.to_state = value.split("_TO_")
-        return obj
-
-
-def get_transition_count(
-    pop: pd.DataFrame,
-    config: Dict[str, bool],
-    state_machine: str,
-    transition: TransitionString,
-    event_time: pd.Timestamp,
-    age_bins: pd.DataFrame,
-) -> Dict[str, float]:
-    """Counts transitions that occurred this step."""
-    event_this_step = (pop[f"previous_{state_machine}"] == transition.from_state) & (
-        pop[state_machine] == transition.to_state
-    )
-    transitioned_pop = pop.loc[event_this_step]
-    base_key = get_output_template(**config).substitute(
-        measure=f"{transition}_event_count", year=event_time.year
-    )
-    base_filter = QueryString("")
-    transition_count = get_group_counts(
-        transitioned_pop, base_filter, base_key, config, age_bins
-    )
-    return transition_count
diff --git a/src/vivarium_public_health/population/mortality.py b/src/vivarium_public_health/population/mortality.py
index 598fc430c..1e589ce60 100644
--- a/src/vivarium_public_health/population/mortality.py
+++ b/src/vivarium_public_health/population/mortality.py
@@ -3,11 +3,48 @@
 The Core Mortality Model
 ========================
 
-This module contains tools modeling all cause mortality and hooks for
-disease models to contribute cause-specific and excess mortality.
+Summary
+=======
+
+The mortality component models all cause mortality and allows for disease
+models to contribute cause specific mortality. At each timestep the currently
+"alive" population is subjected to a mortality event using the mortality rate to
+determine probabilities of death for each simulant. A weighted probable cause of
+death is used to choose the cause of death. The years of life lost are
+calculated by subtracting a simulant's age from the population TMRLE and the
+population is updated.
+
+Columns Created
+===============
+
+ - cause_of_death
+ - years_of_life_lost
+
+Pipelines Exposed
+=================
+
+ - cause_specific_mortality_rate
+ - mortality_rate
+ - affected_unmodeled.cause_specific_mortality_rate
+ - affected_unmodeled.cause_specific_mortality_rate.paf
+
+
+All cause mortality is read from the artifact (GBD). At setup cause specific
+mortality is initialized to an empty table. As disease models are registered,
+they affect cause specific mortality by means of
+the cause_specific_mortality_rate pipeline. This is population level data.
+
+If there are causes of death which are unmodeled, but may be impacted by some
+modeled entity, they can be specified using the configuration key
+"unmodeled_causes".
+
+The mortality component's mortality_rate pipeline reflects the
+cause deleted mortality rate (ACMR - CSMR). Then the impact of unmodeled causes
+on mortality is calculated, by subtracting the raw unmodeled csmr before adding
+back the modified unmodeled csmr.
 
 """
-from typing import Callable
+from typing import Dict, List, Union
 
 import pandas as pd
 from vivarium.framework.engine import Builder
@@ -15,22 +52,35 @@
 from vivarium.framework.lookup import LookupTable
 from vivarium.framework.population import PopulationView, SimulantData
 from vivarium.framework.randomness import RandomnessStream
-from vivarium.framework.time import Time
-from vivarium.framework.utilities import rate_to_probability
-from vivarium.framework.values import Pipeline
+from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
 
 
 class Mortality:
+
+    configuration_defaults = {"unmodeled_causes": []}
+
     def __init__(self):
         self._randomness_stream_name = "mortality_handler"
-        self.cause_specific_mortality_rate_pipeline_name = "cause_specific_mortality_rate"
-        self.mortality_rate_pipeline_name = "mortality_rate"
+
         self.cause_of_death_column_name = "cause_of_death"
         self.years_of_life_lost_column_name = "years_of_life_lost"
 
+        self.cause_specific_mortality_rate_pipeline_name = "cause_specific_mortality_rate"
+        self.mortality_rate_pipeline_name = "mortality_rate"
+        self.unmodeled_csmr_pipeline_name = "affected_unmodeled.cause_specific_mortality_rate"
+        self.unmodeled_csmr_paf_pipeline_name = f"{self.unmodeled_csmr_pipeline_name}.paf"
+
     def __repr__(self) -> str:
         return f"Mortality()"
 
+    ##########################
+    # Initialization methods #
+    ##########################
+
+    # noinspection PyMethodMayBeStatic
+    def _get_configuration_defaults(self) -> Dict[str, List]:
+        return {"unmodeled_causes": Mortality.configuration_defaults["unmodeled_causes"]}
+
     ##############
     # Properties #
     ##############
@@ -46,11 +96,18 @@ def name(self) -> str:
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
         self.random = self._get_randomness_stream(builder)
-        self.clock = self._get_clock(builder)
-        self.all_cause_mortality_rate = self._get_all_cause_mortality_rate(builder)
+        self.clock = builder.time.clock()
+
         self.cause_specific_mortality_rate = self._get_cause_specific_mortality_rate(builder)
         self.mortality_rate = self._get_mortality_rate(builder)
+
+        self.all_cause_mortality_rate = self._get_all_cause_mortality_rate(builder)
         self.life_expectancy = self._get_life_expectancy(builder)
+
+        self._raw_unmodeled_csmr = self._get_raw_unmodeled_csmr(builder)
+        self.unmodeled_csmr = self._get_unmodeled_csmr(builder)
+        self.unmodeled_csmr_paf = self._get_unmodeled_csmr_paf(builder)
+
         self.population_view = self._get_population_view(builder)
 
         self._register_simulant_initializer(builder)
@@ -59,17 +116,6 @@ def setup(self, builder: Builder) -> None:
     def _get_randomness_stream(self, builder) -> RandomnessStream:
         return builder.randomness.get_stream(self._randomness_stream_name)
 
-    # noinspection PyMethodMayBeStatic
-    def _get_clock(self, builder: Builder) -> Callable[[], Time]:
-        return builder.time.clock()
-
-    # noinspection PyMethodMayBeStatic
-    def _get_all_cause_mortality_rate(self, builder: Builder) -> LookupTable:
-        acmr_data = builder.data.load("cause.all_causes.cause_specific_mortality_rate")
-        return builder.lookup.build_table(
-            acmr_data, key_columns=["sex"], parameter_columns=["age", "year"]
-        )
-
     def _get_cause_specific_mortality_rate(self, builder: Builder) -> Pipeline:
         return builder.value.register_value_producer(
             self.cause_specific_mortality_rate_pipeline_name,
@@ -84,12 +130,54 @@ def _get_mortality_rate(self, builder: Builder) -> Pipeline:
         )
 
     # noinspection PyMethodMayBeStatic
-    def _get_life_expectancy(self, builder: Builder) -> LookupTable:
+    def _get_all_cause_mortality_rate(self, builder: Builder) -> Union[LookupTable, Pipeline]:
+        acmr_data = builder.data.load("cause.all_causes.cause_specific_mortality_rate")
+        return builder.lookup.build_table(
+            acmr_data, key_columns=["sex"], parameter_columns=["age", "year"]
+        )
+
+    # noinspection PyMethodMayBeStatic
+    def _get_life_expectancy(self, builder: Builder) -> Union[LookupTable, Pipeline]:
         life_expectancy_data = builder.data.load(
             "population.theoretical_minimum_risk_life_expectancy"
         )
         return builder.lookup.build_table(life_expectancy_data, parameter_columns=["age"])
 
+    # noinspection PyMethodMayBeStatic
+    def _get_raw_unmodeled_csmr(self, builder: Builder) -> Union[LookupTable, Pipeline]:
+        unmodeled_causes = builder.configuration.unmodeled_causes
+        raw_csmr = 0.0
+        for idx, cause in enumerate(unmodeled_causes):
+            csmr = f"cause.{cause}.cause_specific_mortality_rate"
+            if 0 == idx:
+                raw_csmr = builder.data.load(csmr)
+            else:
+                raw_csmr.loc[:, "value"] += builder.data.load(csmr).value
+
+        additional_parameters = (
+            {"key_columns": ["sex"], "parameter_columns": ["age", "year"]}
+            if unmodeled_causes
+            else {}
+        )
+
+        return builder.lookup.build_table(raw_csmr, **additional_parameters)
+
+    def _get_unmodeled_csmr(self, builder: Builder) -> Pipeline:
+        return builder.value.register_value_producer(
+            self.unmodeled_csmr_pipeline_name,
+            source=self._get_unmodeled_csmr_source,
+            requires_columns=["age", "sex"],
+        )
+
+    def _get_unmodeled_csmr_paf(self, builder: Builder) -> Pipeline:
+        unmodeled_csmr_paf = builder.lookup.build_table(0)
+        return builder.value.register_value_producer(
+            self.unmodeled_csmr_paf_pipeline_name,
+            source=lambda index: [unmodeled_csmr_paf(index)],
+            preferred_combiner=list_combiner,
+            preferred_post_processor=union_post_processor,
+        )
+
     def _get_population_view(self, builder: Builder) -> PopulationView:
         return builder.population.get_view(
             [
@@ -99,7 +187,6 @@ def _get_population_view(self, builder: Builder) -> PopulationView:
                 "exit_time",
                 "age",
                 "sex",
-                "location",
             ]
         )
 
@@ -131,20 +218,26 @@ def on_initialize_simulants(self, pop_data: SimulantData) -> None:
 
     def on_time_step(self, event: Event) -> None:
         pop = self.population_view.get(event.index, query="alive =='alive'")
-        prob_df = rate_to_probability(pd.DataFrame(self.mortality_rate(pop.index)))
-        prob_df["no_death"] = 1 - prob_df.sum(axis=1)
-        prob_df["cause_of_death"] = self.random.choice(
-            prob_df.index, prob_df.columns, prob_df
+        mortality_rates = self.mortality_rate(pop.index)
+        mortality_hazard = mortality_rates.sum(axis=1)
+        deaths = self.random.filter_for_rate(
+            pop.index, mortality_hazard, additional_key="death"
         )
-        dead_pop = prob_df.query('cause_of_death != "no_death"').copy()
-
-        if not dead_pop.empty:
-            dead_pop["alive"] = pd.Series("dead", index=dead_pop.index)
-            dead_pop["exit_time"] = event.time
-            dead_pop["years_of_life_lost"] = self.life_expectancy(dead_pop.index)
-            self.population_view.update(
-                dead_pop[["alive", "exit_time", "cause_of_death", "years_of_life_lost"]]
+        if not deaths.empty:
+            cause_of_death_weights = mortality_rates.loc[deaths, :].divide(
+                mortality_hazard.loc[deaths], axis=0
             )
+            cause_of_death = self.random.choice(
+                deaths,
+                cause_of_death_weights.columns,
+                cause_of_death_weights,
+                additional_key="cause_of_death",
+            )
+            pop.loc[deaths, "alive"] = "dead"
+            pop.loc[deaths, "exit_time"] = event.time
+            pop.loc[deaths, "years_of_life_lost"] = self.life_expectancy(deaths)
+            pop.loc[deaths, "cause_of_death"] = cause_of_death
+            self.population_view.update(pop)
 
     ##################################
     # Pipeline sources and modifiers #
@@ -152,6 +245,15 @@ def on_time_step(self, event: Event) -> None:
 
     def _calculate_mortality_rate(self, index: pd.Index) -> pd.DataFrame:
         acmr = self.all_cause_mortality_rate(index)
-        csmr = self.cause_specific_mortality_rate(index, skip_post_processor=True)
-        cause_deleted_mortality_rate = acmr - csmr
+        modeled_csmr = self.cause_specific_mortality_rate(index)
+        unmodeled_csmr_raw = self._raw_unmodeled_csmr(index)
+        unmodeled_csmr = self.unmodeled_csmr(index)
+        cause_deleted_mortality_rate = (
+            acmr - modeled_csmr - unmodeled_csmr_raw + unmodeled_csmr
+        )
         return pd.DataFrame({"other_causes": cause_deleted_mortality_rate})
+
+    def _get_unmodeled_csmr_source(self, index: pd.Index) -> pd.Series:
+        raw_csmr = self._raw_unmodeled_csmr(index)
+        paf = self.unmodeled_csmr_paf(index)
+        return raw_csmr * (1 - paf)
diff --git a/src/vivarium_public_health/risks/effect.py b/src/vivarium_public_health/risks/effect.py
index 6e69bc085..4792a30a5 100644
--- a/src/vivarium_public_health/risks/effect.py
+++ b/src/vivarium_public_health/risks/effect.py
@@ -10,12 +10,13 @@
 
 from typing import Callable, Dict
 
+import numpy as np
 import pandas as pd
 from vivarium.framework.engine import Builder
 from vivarium.framework.lookup import LookupTable
 
 from vivarium_public_health.risks.data_transformations import (
-    get_exposure_effect,
+    get_distribution_type,
     get_population_attributable_fraction_data,
     get_relative_risk_data,
 )
@@ -68,6 +69,7 @@ def __init__(self, risk: str, target: str):
         self.target = TargetString(target)
         self.configuration_defaults = self._get_configuration_defaults()
 
+        self.exposure_pipeline_name = f"{self.risk.name}.exposure"
         self.target_pipeline_name = f"{self.target.name}.{self.target.measure}"
         self.target_paf_pipeline_name = f"{self.target_pipeline_name}.paf"
 
@@ -101,15 +103,24 @@ def name(self) -> str:
 
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
+        self.exposure_distribution_type = self._get_distribution_type(builder)
+        self.exposure = self._get_risk_exposure(builder)
         self.relative_risk = self._get_relative_risk_source(builder)
         self.population_attributable_fraction = (
             self._get_population_attributable_fraction_source(builder)
         )
+
         self.target_modifier = self._get_target_modifier(builder)
 
         self._register_target_modifier(builder)
         self._register_paf_modifier(builder)
 
+    def _get_distribution_type(self, builder: Builder) -> str:
+        return get_distribution_type(builder, self.risk)
+
+    def _get_risk_exposure(self, builder: Builder) -> Callable[[pd.Index], pd.Series]:
+        return builder.value.get_value(self.exposure_pipeline_name)
+
     def _get_relative_risk_source(self, builder: Builder) -> LookupTable:
         relative_risk_data = get_relative_risk_data(builder, self.risk, self.target)
         return builder.lookup.build_table(
@@ -125,10 +136,33 @@ def _get_population_attributable_fraction_source(self, builder: Builder) -> Look
     def _get_target_modifier(
         self, builder: Builder
     ) -> Callable[[pd.Index, pd.Series], pd.Series]:
-        exposure_effect = get_exposure_effect(builder, self.risk)
-
-        def adjust_target(index: pd.Index, target: pd.Series) -> pd.Series:
-            return exposure_effect(target, self.relative_risk(index))
+        if self.exposure_distribution_type in ["normal", "lognormal", "ensemble"]:
+            tmred = builder.data.load(f"{self.risk}.tmred")
+            tmrel = 0.5 * (tmred["min"] + tmred["max"])
+            scale = builder.data.load(f"{self.risk}.relative_risk_scalar")
+
+            def adjust_target(index: pd.Index, target: pd.Series) -> pd.Series:
+                rr = self.relative_risk(index)
+                exposure = self.exposure(index)
+                relative_risk = np.maximum(rr.values ** ((exposure - tmrel) / scale), 1)
+                return target * relative_risk
+
+        else:
+            index_columns = ["index", self.risk.name]
+
+            def adjust_target(index: pd.Index, target: pd.Series) -> pd.Series:
+                rr = self.relative_risk(index)
+                exposure = self.exposure(index).reset_index()
+                exposure.columns = index_columns
+                exposure = exposure.set_index(index_columns)
+
+                relative_risk = rr.stack().reset_index()
+                relative_risk.columns = index_columns + ["value"]
+                relative_risk = relative_risk.set_index(index_columns)
+
+                effect = relative_risk.loc[exposure.index, "value"].droplevel(self.risk.name)
+                affected_rates = target * effect
+                return affected_rates
 
         return adjust_target
 
diff --git a/src/vivarium_public_health/treatment/scale_up.py b/src/vivarium_public_health/treatment/scale_up.py
index 577bb17dc..b79540a2f 100644
--- a/src/vivarium_public_health/treatment/scale_up.py
+++ b/src/vivarium_public_health/treatment/scale_up.py
@@ -52,13 +52,13 @@ class LinearScaleUp:
 
     configuration_defaults = {
         "treatment": {
-            "start": {
-                "date": "start",
-                "value": "data",
+            "date": {
+                "start": "start",
+                "end": "end",
             },
-            "end": {
-                "date": "end",
-                "value": "data",
+            "value": {
+                "start": "data",
+                "end": "data",
             },
         }
     }
@@ -101,13 +101,10 @@ def setup(self, builder: Builder) -> None:
         """Perform this component's setup."""
         self.is_intervention_scenario = self._get_is_intervention_scenario(builder)
         self.clock = self._get_clock(builder)
-        self.scale_up_start_date, self.scale_up_end_date = self._get_scale_up_date_endpoints(
+        self.scale_up_start_date, self.scale_up_end_date = self._get_scale_up_dates(builder)
+        self.scale_up_start_value, self.scale_up_end_value = self._get_scale_up_values(
             builder
         )
-        (
-            self.scale_up_start_value,
-            self.scale_up_end_value,
-        ) = self._get_scale_up_value_endpoints(builder)
 
         required_columns = self._get_required_columns()
         self.pipelines = self._get_required_pipelines(builder)
@@ -126,28 +123,26 @@ def _get_clock(self, builder: Builder) -> Callable[[], Time]:
         return builder.time.clock()
 
     # noinspection PyMethodMayBeStatic
-    def _get_scale_up_date_endpoints(self, builder: Builder) -> Tuple[datetime, datetime]:
-        scale_up_config = builder.configuration[self.configuration_key]
+    def _get_scale_up_dates(self, builder: Builder) -> Tuple[datetime, datetime]:
+        scale_up_config = builder.configuration[self.configuration_key]["date"]
 
         def get_endpoint(endpoint_type: str) -> datetime:
-            if scale_up_config[endpoint_type]["date"] == endpoint_type:
+            if scale_up_config[endpoint_type] == endpoint_type:
                 endpoint = get_time_stamp(builder.configuration.time[endpoint_type])
             else:
-                endpoint = get_time_stamp(scale_up_config[endpoint_type]["date"])
+                endpoint = get_time_stamp(scale_up_config[endpoint_type])
             return endpoint
 
         return get_endpoint("start"), get_endpoint("end")
 
-    def _get_scale_up_value_endpoints(
-        self, builder: Builder
-    ) -> Tuple[LookupTable, LookupTable]:
-        scale_up_config = builder.configuration[self.configuration_key]
+    def _get_scale_up_values(self, builder: Builder) -> Tuple[LookupTable, LookupTable]:
+        scale_up_config = builder.configuration[self.configuration_key]["value"]
 
         def get_endpoint_value(endpoint_type: str) -> LookupTable:
-            if scale_up_config[endpoint_type]["value"] == "data":
+            if scale_up_config[endpoint_type] == "data":
                 endpoint = self._get_endpoint_value_from_data(builder, endpoint_type)
             else:
-                endpoint = builder.lookup.build_table(scale_up_config[endpoint_type]["value"])
+                endpoint = builder.lookup.build_table(scale_up_config[endpoint_type])
             return endpoint
 
         return get_endpoint_value("start"), get_endpoint_value("end")
@@ -178,19 +173,15 @@ def _get_population_view(
 
     def _coverage_effect(self, idx: pd.Index, target: pd.Series) -> pd.Series:
         if not self.is_intervention_scenario or self.clock() < self.scale_up_start_date:
-            scale_up_progress = 0.0
+            progress = 0.0
         elif self.scale_up_start_date <= self.clock() < self.scale_up_end_date:
-            scale_up_progress = (self.clock() - self.scale_up_start_date) / (
+            progress = (self.clock() - self.scale_up_start_date) / (
                 self.scale_up_end_date - self.scale_up_start_date
             )
         else:
-            scale_up_progress = 1.0
+            progress = 1.0
 
-        target = (
-            self._apply_scale_up(idx, target, scale_up_progress)
-            if scale_up_progress
-            else target
-        )
+        target = self._apply_scale_up(idx, target, progress) if progress else target
         return target
 
     ##################
diff --git a/src/vivarium_public_health/utilities.py b/src/vivarium_public_health/utilities.py
index a720ee898..518bf7555 100644
--- a/src/vivarium_public_health/utilities.py
+++ b/src/vivarium_public_health/utilities.py
@@ -7,9 +7,10 @@
 vivarium_public_health components.
 
 """
-from typing import Union
+from typing import Iterable, Union
 
 import pandas as pd
+from vivarium.framework.lookup import ScalarValue
 
 
 class EntityString(str):
@@ -80,3 +81,14 @@ def to_time_delta(span_in_days: Union[int, float, str]):
 def to_years(time: pd.Timedelta) -> float:
     """Converts a time delta to a float for years."""
     return time / pd.Timedelta(days=DAYS_PER_YEAR)
+
+
+def is_non_zero(data: Union[Iterable[ScalarValue], ScalarValue, pd.DataFrame]) -> bool:
+    if isinstance(data, pd.DataFrame):
+        attribute_sum = data.value.sum()
+    elif isinstance(data, Iterable):
+        attribute_sum = sum(data)
+    else:
+        attribute_sum = data
+
+    return attribute_sum != 0.0


Test patch:
diff --git a/tests/disease/test_disease.py b/tests/disease/test_disease.py
index b58fdf35d..faec575c9 100644
--- a/tests/disease/test_disease.py
+++ b/tests/disease/test_disease.py
@@ -3,7 +3,6 @@
 import pytest
 from vivarium import InteractiveContext
 from vivarium.framework.utilities import from_yearly
-from vivarium.framework.values import rescale_post_processor
 from vivarium.testing_utilities import TestPopulation, build_table, metadata
 
 from vivarium_public_health.disease import (
@@ -13,6 +12,7 @@
     RateTransition,
 )
 from vivarium_public_health.disease.state import SusceptibleState
+from vivarium_public_health.disease.transition import TransitionString
 from vivarium_public_health.population import Mortality
 
 
@@ -497,12 +497,11 @@ def test_state_transition_names(disease):
     healthy.add_transition(with_condition)
     with_condition.add_transition(healthy)
     model = DiseaseModel(disease, initial_state=healthy, states=[healthy, with_condition])
-    assert set(model.state_names) == set(
-        ["diarrheal_diseases", "susceptible_to_diarrheal_diseases"]
-    )
-    assert set(model.transition_names) == set(
-        [
-            "diarrheal_diseases_TO_susceptible_to_diarrheal_diseases",
-            "susceptible_to_diarrheal_diseases_TO_diarrheal_diseases",
-        ]
-    )
+    assert set(model.state_names) == {
+        "diarrheal_diseases",
+        "susceptible_to_diarrheal_diseases",
+    }
+    assert set(model.transition_names) == {
+        TransitionString("diarrheal_diseases_TO_susceptible_to_diarrheal_diseases"),
+        TransitionString("susceptible_to_diarrheal_diseases_TO_diarrheal_diseases"),
+    }
diff --git a/tests/disease/test_special_disease.py b/tests/disease/test_special_disease.py
index e563e870b..10ad7afa6 100644
--- a/tests/disease/test_special_disease.py
+++ b/tests/disease/test_special_disease.py
@@ -5,6 +5,7 @@
 import pytest
 
 from vivarium_public_health.disease import RiskAttributableDisease
+from vivarium_public_health.disease.transition import TransitionString
 
 
 @pytest.fixture
@@ -128,10 +129,7 @@ def test_mortality_rate_pandas_dataframe(disease_mock):
     assert np.all(expected == disease.adjust_mortality_rate(test_index, rates_df))
 
 
-test_data = [
-    ("disease_no_recovery", False),
-    ("disease_with_recovery", True),
-]
+test_data = [("disease_no_recovery", False), ("disease_with_recovery", True)]
 
 
 @pytest.mark.parametrize("disease, recoverable", test_data)
@@ -140,10 +138,8 @@ def test_state_transition_names(disease, recoverable):
     model.recoverable = recoverable
     model.adjust_state_and_transitions()
     states = [disease, f"susceptible_to_{disease}"]
-    transitions = [
-        f"susceptible_to_{disease}_TO_{disease}",
-    ]
+    transitions = [TransitionString(f"susceptible_to_{disease}_TO_{disease}")]
     if recoverable:
-        transitions.append(f"{disease}_TO_susceptible_to_{disease}")
+        transitions.append(TransitionString(f"{disease}_TO_susceptible_to_{disease}"))
     assert set(model.state_names) == set(states)
     assert set(model.transition_names) == set(transitions)
diff --git a/tests/metrics/test_stratification.py b/tests/metrics/test_stratification.py
new file mode 100644
index 000000000..b56171723
--- /dev/null
+++ b/tests/metrics/test_stratification.py
@@ -0,0 +1,426 @@
+from datetime import datetime
+from math import floor
+from typing import Callable, List, Set
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+import pytest
+
+from vivarium_public_health.metrics.stratification import (
+    ResultsStratifier,
+    Source,
+    SourceType,
+    StratificationLevel,
+)
+
+AGE_BINS = pd.DataFrame(
+    [
+        {"age_start": age, "age_end": age + 5, "age_group_name": f"{age}_to_{age + 5}"}
+        for i, age in enumerate(range(10, 65, 5))
+    ]
+)
+DEFAULT_CONFIGS = ["year", "sex", "age"]
+START_YEAR = 2020
+END_YEAR = 2030
+STRATIFICATION_LEVELS = {
+    "age": StratificationLevel(
+        name="age",
+        sources=[Source("age", SourceType.COLUMN)],
+        categories={"0", "1", "2", "3", "4"},
+        mapper=lambda row: str(floor(row[0])),
+    ),
+    "sex": StratificationLevel(
+        name="l2",
+        sources=[Source("sex", SourceType.COLUMN)],
+        categories={"Male", "Female"},
+    ),
+    "multi": StratificationLevel(
+        name="multi",
+        sources=[
+            Source("part_1", SourceType.COLUMN),
+            Source("pipeline_1", SourceType.PIPELINE),
+        ],
+        categories={"c2", "c30", "c31", "c10", "c11"},
+        mapper=lambda row: "c2"
+        if row["part_1"] == "c2"
+        else row["part_1"] + row["pipeline_1"],
+    ),
+    "p2": StratificationLevel(
+        name="p2",
+        sources=[Source("pipeline_2", SourceType.PIPELINE)],
+        categories={"2", "3"},
+    ),
+    "year": StratificationLevel(
+        name="year",
+        sources=[Source("clock_1", SourceType.CLOCK)],
+        categories={"2020", "2021", "2022", "2023"},
+        mapper=lambda row: str(row[0].year),
+        current_categories_getter=lambda: {str(mock_clock().year)},
+    ),
+    "month": StratificationLevel(
+        name="month",
+        sources=[Source("clock_2", SourceType.CLOCK)],
+        categories={str(i + 1) for i in range(12)},
+        mapper=lambda row: str(row[0].month),
+        current_categories_getter=lambda: {str(mock_clock().month)},
+    ),
+}
+
+
+def unmock_instance_method(instance, method: Callable) -> Callable:
+    return lambda *args: method(instance, *args)
+
+
+@pytest.fixture
+def mock_stratification_level_init(mocker):
+    return mocker.patch(
+        "vivarium_public_health.metrics.stratification.StratificationLevel.__init__",
+        return_value=None,
+    )
+
+
+@pytest.fixture()
+def mock_stratifier(mocker) -> MagicMock:
+    stratifier = mocker.MagicMock()
+    stratifier.setup = mocker.MagicMock(
+        side_effect=unmock_instance_method(stratifier, ResultsStratifier.setup)
+    )
+    return stratifier
+
+
+@pytest.fixture()
+def mock_builder(mocker) -> MagicMock:
+    builder = mocker.MagicMock()
+
+    builder.configuration.observers.default = DEFAULT_CONFIGS
+    builder.configuration.time.start.year = START_YEAR
+    builder.configuration.time.end.year = END_YEAR
+
+    builder.data.load = mocker.MagicMock(side_effect=load)
+    builder.value.get_value = mocker.MagicMock(side_effect=get_value)
+
+    builder.time.clock.return_value = mock_clock
+    return builder
+
+
+def mock_clock() -> datetime:
+    return datetime(2022, 5, 3, 12)
+
+
+def load(key: str) -> pd.DataFrame:
+    return {
+        "population.age_bins": pd.DataFrame(
+            [
+                {
+                    "age_start": age,
+                    "age_end": age + 5,
+                    "age_group_name": f"{age}_to_{age + 5}",
+                }
+                for i, age in enumerate(range(0, 100, 5))
+            ]
+        )
+    }[key]
+
+
+def get_value(pipeline_name: str) -> str:
+    return f"{pipeline_name}_pipeline"
+
+
+def mock_mapper(row: pd.Series) -> str:
+    return str(row[0]) + str(row[0])
+
+
+def mock_current_categories_getter() -> Set[str]:
+    return {"a", "b", "c"}
+
+
+def test_setting_default_stratification(mock_stratifier: MagicMock, mock_builder: MagicMock):
+    # Setup mocks
+    mock_stratifier._get_default_stratification_levels = MagicMock(
+        side_effect=unmock_instance_method(
+            mock_stratifier, ResultsStratifier._get_default_stratification_levels
+        )
+    )
+
+    # Code to test
+    mock_stratifier.setup(mock_builder)
+
+    # Assertions
+    assert mock_stratifier.default_stratification_levels == set(DEFAULT_CONFIGS)
+
+
+@pytest.mark.parametrize(
+    "age_start, age_end, expected_age_start, expected_age_end",
+    [(20, 55, 20, 55), (18, 47, 15, 50)],
+    ids=["on_bound", "between_bounds"],
+)
+def test_setting_age_bins(
+    mock_stratifier: MagicMock,
+    mock_builder: MagicMock,
+    age_start: int,
+    age_end: int,
+    expected_age_start: int,
+    expected_age_end: int,
+):
+    # Setup mocks
+    mock_builder.configuration.population.age_start = age_start
+    mock_builder.configuration.population.exit_age = age_end
+
+    mock_stratifier._get_age_bins = MagicMock(
+        side_effect=unmock_instance_method(mock_stratifier, ResultsStratifier._get_age_bins)
+    )
+
+    # Code to test
+    mock_stratifier.setup(mock_builder)
+
+    # Assertions
+    expected_outputs = pd.DataFrame(
+        [
+            {"age_start": age, "age_end": age + 5, "age_group_name": f"{age}_to_{age + 5}"}
+            for i, age in enumerate(range(expected_age_start, expected_age_end, 5))
+        ]
+    )
+    assert (mock_stratifier.age_bins.values == expected_outputs.values).all()
+
+
+def test_stratifications_and_listeners_registered(
+    mock_stratifier: MagicMock,
+    mock_builder: MagicMock,
+):
+    # Setup mocks
+    mock_stratifier._register_timestep_prepare_listener = MagicMock(
+        side_effect=unmock_instance_method(
+            mock_stratifier, ResultsStratifier._register_timestep_prepare_listener
+        )
+    )
+
+    # Code to test
+    mock_stratifier.setup(mock_builder)
+
+    # Assertions
+    mock_stratifier.register_stratifications.assert_called_once_with(mock_builder)
+    mock_builder.event.register_listener.assert_called_once_with(
+        "time_step__prepare", mock_stratifier.on_time_step_prepare, priority=0
+    )
+
+
+def test_registering_stratifications(mock_stratifier: MagicMock, mock_builder: MagicMock):
+    # Setup mocks
+    mock_stratifier.age_bins = AGE_BINS
+    mock_stratifier.register_stratifications = MagicMock(
+        side_effect=unmock_instance_method(
+            mock_stratifier, ResultsStratifier.register_stratifications
+        )
+    )
+
+    # Code to test
+    mock_stratifier.register_stratifications(mock_builder)
+
+    # Assertions
+    mock_stratifier.setup_stratification.assert_any_call(
+        mock_builder,
+        name="year",
+        sources=[Source("year", SourceType.CLOCK)],
+        categories={str(year) for year in range(START_YEAR, END_YEAR + 1)},
+        mapper=mock_stratifier.year_stratification_mapper,
+        current_category_getter=mock_stratifier.year_current_categories_getter,
+    )
+
+    mock_stratifier.setup_stratification.assert_any_call(
+        mock_builder,
+        name="sex",
+        sources=[Source("sex", SourceType.COLUMN)],
+        categories={"Female", "Male"},
+    )
+
+    mock_stratifier.setup_stratification.assert_any_call(
+        mock_builder,
+        name="age",
+        sources=[Source("age", SourceType.COLUMN)],
+        categories={age_bin for age_bin in AGE_BINS["age_group_name"]},
+        mapper=mock_stratifier.age_stratification_mapper,
+    )
+
+
+@pytest.mark.parametrize(
+    "sources",
+    [
+        [Source("age", SourceType.COLUMN)],
+        [Source("exposure", SourceType.PIPELINE)],
+        [Source("year", SourceType.CLOCK)],
+        [Source("age", SourceType.COLUMN), Source("exposure", SourceType.PIPELINE)],
+    ],
+    ids=["column", "pipeline", "clock", "multiple"],
+)
+def test_setup_stratification(
+    mock_stratifier: MagicMock,
+    mock_builder: MagicMock,
+    mock_stratification_level_init: patch,
+    sources: List[Source],
+):
+    # Setup mocks
+    mock_stratifier.stratification_levels = {}
+    mock_stratifier.pipelines = {}
+    mock_stratifier.columns_required = {"tracked"}
+    mock_stratifier.clock_sources = set()
+
+    mock_stratifier.setup_stratification = MagicMock(
+        side_effect=unmock_instance_method(
+            mock_stratifier, ResultsStratifier.setup_stratification
+        )
+    )
+    name = "test_name"
+    categories = {"a", "b", "c"}
+
+    # Code to test
+    mock_stratifier.setup_stratification(
+        mock_builder,
+        name,
+        sources,
+        categories,
+        mock_mapper,
+        mock_current_categories_getter,
+    )
+
+    # Assertions
+    mock_stratification_level_init.assert_called_once_with(
+        name, sources, categories, mock_mapper, mock_current_categories_getter
+    )
+    assert name in mock_stratifier.stratification_levels
+
+    for source in sources:
+        if source.type == SourceType.PIPELINE:
+            assert source.name in mock_stratifier.pipelines
+            assert mock_stratifier.pipelines[source.name] == get_value(source.name)
+        if source.type == SourceType.COLUMN:
+            assert source.name in mock_stratifier.columns_required
+        if source.type == SourceType.CLOCK:
+            assert source.name in mock_stratifier.clock_sources
+
+
+def test_setting_stratification_groups_on_time_step_prepare(
+    mock_stratifier: MagicMock, mock_builder: MagicMock
+):
+    # Setup mocks
+    def mock_pipeline_1(index: pd.Index) -> pd.Series:
+        return pd.Series([str(i % 2) for i in range(index.size)], index=index)
+
+    def mock_pipeline_2(index: pd.Index) -> pd.Series:
+        return pd.Series([str(i % 2 + 2) for i in range(index.size)], index=index)
+
+    mock_pipelines = {
+        "pipeline_1": mock_pipeline_1,
+        "pipeline_2": mock_pipeline_2,
+    }
+
+    mock_population_view_data = pd.DataFrame(
+        {
+            "age": [4.5, 3.4, 0.2, 0.3],
+            "sex": ["Male", "Female", "Female", "Male"],
+            "part_1": ["c2", "c2", "c3", "c1"],
+        }
+    )
+
+    mock_clock_sources = {"clock_1", "clock_2"}
+
+    mock_event = MagicMock()
+    mock_event.index = mock_population_view_data.index
+
+    mock_stratifier.population_view.get.return_value = mock_population_view_data
+    mock_stratifier.pipelines = mock_pipelines
+    mock_stratifier.clock_sources = mock_clock_sources
+    mock_stratifier.clock = mock_clock
+    mock_stratifier.stratification_levels = STRATIFICATION_LEVELS
+
+    mock_stratifier.on_time_step_prepare = MagicMock(
+        side_effect=unmock_instance_method(
+            mock_stratifier, ResultsStratifier.on_time_step_prepare
+        )
+    )
+
+    # Code to test
+    mock_stratifier.on_time_step_prepare(mock_event)
+
+    # Assertions
+    expected_groups = pd.DataFrame(
+        {
+            "age": ["4", "3", "0", "0"],
+            "sex": ["Male", "Female", "Female", "Male"],
+            "multi": ["c2", "c2", "c30", "c11"],
+            "p2": ["2", "3", "2", "3"],
+            "year": ["2022", "2022", "2022", "2022"],
+            "month": ["5", "5", "5", "5"],
+        },
+        index=mock_population_view_data.index,
+    )
+    assert (mock_stratifier.stratification_groups.values == expected_groups.values).all()
+
+
+# todo get test_group to work
+# TEST_GROUP_INDEX = pd.Index(list(range(4)))
+#
+#
+# @pytest.mark.parametrize(
+#     "include, exclude, expected",
+#     [
+#         ([], DEFAULT_CONFIGS, {"": True}),
+#         ([], ["age", "year"], {
+#             "sex_male": pd.Series([True, False, False, True], index=TEST_GROUP_INDEX),
+#             "sex_female": pd.Series([False, True, True, False], index=TEST_GROUP_INDEX)
+#         }),
+#     ],
+#     ids=[
+#         'no stratification',
+#         "sex_only",
+#     ]
+# )
+# def test_group(
+#         mock_stratifier: MagicMock,
+#         include: Set[str],
+#         exclude: Set[str],
+#         expected: Dict[str, pd.Series]
+# ):
+#     # Setup mocks
+#     mock_stratifier.default_stratification_levels = set(DEFAULT_CONFIGS)
+#     mock_stratifier.stratification_levels = STRATIFICATION_LEVELS
+#     mock_stratifier.stratification_groups = pd.DataFrame(
+#         {
+#             "age": ["4", "3", "0", "0", "3"],
+#             "sex": ["Male", "Female", "Female", "Male", "Male"],
+#             "multi": ["c2", "c2", "c30", "c11", "c11"],
+#             "p2": ["2", "3", "2", "3", "3"],
+#             "year": ["2022", "2022", "2022", "2022", "2022"],
+#             "month": ["5", "5", "5", "5", "5"],
+#         }
+#     )
+#
+#     mock_stratifier.group = MagicMock(
+#         side_effect=unmock_instance_method(mock_stratifier, ResultsStratifier.group)
+#     )
+#     mock_stratifier._get_current_stratifications = MagicMock(
+#         side_effect=unmock_instance_method(
+#             mock_stratifier, ResultsStratifier._get_current_stratifications
+#         )
+#     )
+#     mock_stratifier._get_stratification_key = MagicMock(
+#         side_effect=ResultsStratifier._get_stratification_key
+#     )
+#
+#     # Code to test
+#     groups = mock_stratifier.group(TEST_GROUP_INDEX, include, exclude)
+#     number_of_groups = 0
+#
+#     # Assertions
+#     for label, group_mask in groups:
+#         assert (
+#             group_mask == expected[label] if isinstance(group_mask, bool)
+#             else (group_mask == expected[label]).all()
+#         )
+#         number_of_groups += 1
+#
+#     assert number_of_groups == len(expected)
+
+# todo test age_stratification_mapper
+# todo test year_stratification_mapper
+# todo test year_current_categories_getter
+# todo test StratificationLevel init
diff --git a/tests/metrics/test_utilities.py b/tests/metrics/test_utilities.py
deleted file mode 100644
index 6490bd204..000000000
--- a/tests/metrics/test_utilities.py
+++ /dev/null
@@ -1,691 +0,0 @@
-from itertools import combinations, product
-
-import numpy as np
-import pandas as pd
-import pytest
-from vivarium.testing_utilities import metadata
-
-from vivarium_public_health.metrics.utilities import (
-    _MAX_AGE,
-    _MAX_YEAR,
-    _MIN_AGE,
-    _MIN_YEAR,
-    OutputTemplate,
-    QueryString,
-    get_age_bins,
-    get_age_sex_filter_and_iterables,
-    get_deaths,
-    get_disease_event_counts,
-    get_lived_in_span,
-    get_output_template,
-    get_person_time_in_span,
-    get_susceptible_person_time,
-    get_time_iterable,
-    get_years_lived_with_disability,
-    get_years_of_life_lost,
-    to_years,
-)
-
-
-@pytest.fixture(params=((0, 100, 5, 1000), (20, 100, 5, 1000)))
-def ages_and_bins(request):
-    age_min = request.param[0]
-    age_max = request.param[1]
-    age_groups = request.param[2]
-    num_ages = request.param[3]
-
-    ages = np.linspace(age_min, age_max - age_groups / num_ages, num_ages)
-    bin_ages, step = np.linspace(age_min, age_max, age_groups, endpoint=False, retstep=True)
-    age_bins = pd.DataFrame(
-        {
-            "age_start": bin_ages,
-            "age_end": bin_ages + step,
-            "age_group_name": [str(name) for name in range(len(bin_ages))],
-        }
-    )
-
-    return ages, age_bins
-
-
-@pytest.fixture
-def sexes():
-    return ["Male", "Female"]
-
-
-@pytest.fixture(params=list(product((True, False), repeat=3)))
-def observer_config(request):
-    c = {"by_age": request.param[0], "by_sex": request.param[1], "by_year": request.param[2]}
-    return c
-
-
-@pytest.fixture()
-def builder(mocker):
-    builder = mocker.MagicMock()
-    df = pd.DataFrame(
-        {
-            "age_start": [0, 1, 4],
-            "age_group_name": ["youngest", "younger", "young"],
-            "age_end": [1, 4, 6],
-        }
-    )
-    builder.data.load.return_value = df
-    return builder
-
-
-@pytest.mark.parametrize(
-    "reference, test", product([QueryString(""), QueryString("abc")], [QueryString(""), ""])
-)
-def test_query_string_empty(reference, test):
-    result = str(reference)
-    assert reference + test == result
-    assert reference + test == QueryString(result)
-    assert isinstance(reference + test, QueryString)
-
-    assert test + reference == result
-    assert test + reference == QueryString(result)
-    assert isinstance(test + reference, QueryString)
-
-    reference += test
-    assert reference == result
-    assert reference == QueryString(result)
-    assert isinstance(reference, QueryString)
-
-    test += reference
-    assert test == result
-    assert test == QueryString(result)
-    assert isinstance(test, QueryString)
-
-
-@pytest.mark.parametrize("a, b", product([QueryString("a")], [QueryString("b"), "b"]))
-def test_query_string(a, b):
-    assert a + b == "a and b"
-    assert a + b == QueryString("a and b")
-    assert isinstance(a + b, QueryString)
-
-    assert b + a == "b and a"
-    assert b + a == QueryString("b and a")
-    assert isinstance(b + a, QueryString)
-
-    a += b
-    assert a == "a and b"
-    assert a == QueryString("a and b")
-    assert isinstance(a, QueryString)
-
-    b += a
-    assert b == "b and a and b"
-    assert b == QueryString("b and a and b")
-    assert isinstance(b, QueryString)
-
-
-def test_get_output_template(observer_config):
-    template = get_output_template(**observer_config)
-
-    assert isinstance(template, OutputTemplate)
-    assert "${measure}" in template.template
-
-    if observer_config["by_year"]:
-        assert "_in_${year}" in template.template
-    if observer_config["by_sex"]:
-        assert "_among_${sex}" in template.template
-    if observer_config["by_age"]:
-        assert "_in_age_group_${age_group}" in template.template
-
-
-@pytest.mark.parametrize(
-    "measure, sex, age, year",
-    product(
-        ["test", "Test"], ["female", "Female"], [1.0, 1, "Early Neonatal"], [2011, "2011"]
-    ),
-)
-def test_output_template(observer_config, measure, sex, age, year):
-    template = get_output_template(**observer_config)
-
-    out1 = template.substitute(measure=measure, sex=sex, age_group=age, year=year)
-    out2 = (
-        template.substitute(measure=measure)
-        .substitute(sex=sex)
-        .substitute(age_group=age)
-        .substitute(year=year)
-    )
-    assert out1 == out2
-
-
-def test_output_template_exact():
-    template = get_output_template(by_age=True, by_sex=True, by_year=True)
-
-    out = template.substitute(measure="Test", sex="Female", age_group=1.0, year=2011)
-    expected = "test_in_2011_among_female_in_age_group_1.0"
-    assert out == expected
-
-    out = template.substitute(
-        measure="Test", sex="Female", age_group="Early Neonatal", year=2011
-    )
-    expected = "test_in_2011_among_female_in_age_group_early_neonatal"
-
-    assert out == expected
-
-
-def test_get_age_sex_filter_and_iterables(ages_and_bins, observer_config):
-    _, age_bins = ages_and_bins
-    age_sex_filter, (ages, sexes) = get_age_sex_filter_and_iterables(
-        observer_config, age_bins
-    )
-
-    assert isinstance(age_sex_filter, QueryString)
-    if observer_config["by_age"] and observer_config["by_sex"]:
-        assert age_sex_filter == '{age_start} <= age and age < {age_end} and sex == "{sex}"'
-
-        for (g1, s1), (g2, s2) in zip(ages, age_bins.set_index("age_group_name").iterrows()):
-            assert g1 == g2
-            assert s1.equals(s2)
-
-        assert sexes == ["Male", "Female"]
-
-    elif observer_config["by_age"]:
-        assert age_sex_filter == "{age_start} <= age and age < {age_end}"
-
-        for (g1, s1), (g2, s2) in zip(ages, age_bins.set_index("age_group_name").iterrows()):
-            assert g1 == g2
-            assert s1.equals(s2)
-
-        assert sexes == ["Both"]
-    elif observer_config["by_sex"]:
-        assert age_sex_filter == 'sex == "{sex}"'
-
-        assert len(ages) == 1
-        group, data = ages[0]
-        assert group == "all_ages"
-        assert data["age_start"] == _MIN_AGE
-        assert data["age_end"] == _MAX_AGE
-
-        assert sexes == ["Male", "Female"]
-
-    else:
-        assert age_sex_filter == ""
-
-        assert len(ages) == 1
-        group, data = ages[0]
-        assert group == "all_ages"
-        assert data["age_start"] == _MIN_AGE
-        assert data["age_end"] == _MAX_AGE
-
-        assert sexes == ["Both"]
-
-
-def test_get_age_sex_filter_and_iterables_with_span(ages_and_bins, observer_config):
-    _, age_bins = ages_and_bins
-    age_sex_filter, (ages, sexes) = get_age_sex_filter_and_iterables(
-        observer_config, age_bins, in_span=True
-    )
-
-    assert isinstance(age_sex_filter, QueryString)
-    if observer_config["by_age"] and observer_config["by_sex"]:
-        expected = '{age_start} < age_at_span_end and age_at_span_start < {age_end} and sex == "{sex}"'
-        assert age_sex_filter == expected
-
-        for (g1, s1), (g2, s2) in zip(ages, age_bins.set_index("age_group_name").iterrows()):
-            assert g1 == g2
-            assert s1.equals(s2)
-
-        assert sexes == ["Male", "Female"]
-
-    elif observer_config["by_age"]:
-        assert (
-            age_sex_filter
-            == "{age_start} < age_at_span_end and age_at_span_start < {age_end}"
-        )
-
-        for (g1, s1), (g2, s2) in zip(ages, age_bins.set_index("age_group_name").iterrows()):
-            assert g1 == g2
-            assert s1.equals(s2)
-
-        assert sexes == ["Both"]
-    elif observer_config["by_sex"]:
-        assert age_sex_filter == 'sex == "{sex}"'
-
-        assert len(ages) == 1
-        group, data = ages[0]
-        assert group == "all_ages"
-        assert data["age_start"] == _MIN_AGE
-        assert data["age_end"] == _MAX_AGE
-
-        assert sexes == ["Male", "Female"]
-
-    else:
-        assert age_sex_filter == ""
-
-        assert len(ages) == 1
-        group, data = ages[0]
-        assert group == "all_ages"
-        assert data["age_start"] == _MIN_AGE
-        assert data["age_end"] == _MAX_AGE
-
-        assert sexes == ["Both"]
-
-
-@pytest.mark.parametrize("year_start, year_end", [(2011, 2017), (2011, 2011)])
-def test_get_time_iterable_no_year(year_start, year_end):
-    config = {"by_year": False}
-    sim_start = pd.Timestamp(f"7-2-{year_start}")
-    sim_end = pd.Timestamp(f"3-15-{year_end}")
-
-    time_spans = get_time_iterable(config, sim_start, sim_end)
-
-    assert len(time_spans) == 1
-    name, (start, end) = time_spans[0]
-    assert name == "all_years"
-    assert start == pd.Timestamp(f"1-1-{_MIN_YEAR}")
-    assert end == pd.Timestamp(f"1-1-{_MAX_YEAR}")
-
-
-@pytest.mark.parametrize("year_start, year_end", [(2011, 2017), (2011, 2011)])
-def test_get_time_iterable_with_year(year_start, year_end):
-    config = {"by_year": True}
-    sim_start = pd.Timestamp(f"7-2-{year_start}")
-    sim_end = pd.Timestamp(f"3-15-{year_end}")
-
-    time_spans = get_time_iterable(config, sim_start, sim_end)
-
-    years = list(range(year_start, year_end + 1))
-    assert len(time_spans) == len(years)
-    for year, time_span in zip(years, time_spans):
-        name, (start, end) = time_span
-        assert name == year
-        assert start == pd.Timestamp(f"1-1-{year}")
-        assert end == pd.Timestamp(f"1-1-{year+1}")
-
-
-def test_get_susceptible_person_time(ages_and_bins, sexes, observer_config):
-    ages, age_bins = ages_and_bins
-    disease = "test_disease"
-    states = [f"susceptible_to_{disease}", disease]
-    pop = pd.DataFrame(list(product(ages, sexes, states)), columns=["age", "sex", disease])
-    pop["alive"] = "alive"
-    # Shuffle the rows
-    pop = pop.sample(frac=1).reset_index(drop=True)
-
-    year = 2017
-    step_size = pd.Timedelta(days=7)
-
-    person_time = get_susceptible_person_time(
-        pop, observer_config, disease, year, step_size, age_bins
-    )
-
-    values = set(person_time.values())
-    assert len(values) == 1
-    expected_value = to_years(step_size) * len(pop) / 2
-    if observer_config["by_sex"]:
-        expected_value /= 2
-    if observer_config["by_age"]:
-        expected_value /= len(age_bins)
-    assert np.isclose(values.pop(), expected_value)
-
-    # Doubling pop should double person time
-    pop = pd.concat([pop, pop], axis=0, ignore_index=True)
-
-    person_time = get_susceptible_person_time(
-        pop, observer_config, disease, year, step_size, age_bins
-    )
-
-    values = set(person_time.values())
-    assert len(values) == 1
-    assert np.isclose(values.pop(), 2 * expected_value)
-
-
-def test_get_disease_event_counts(ages_and_bins, sexes, observer_config):
-    ages, age_bins = ages_and_bins
-    disease = "test_disease"
-    event_time = pd.Timestamp("1-1-2017")
-    states = [event_time, pd.NaT]
-    pop = pd.DataFrame(
-        list(product(ages, sexes, states)), columns=["age", "sex", f"{disease}_event_time"]
-    )
-    # Shuffle the rows
-    pop = pop.sample(frac=1).reset_index(drop=True)
-
-    counts = get_disease_event_counts(pop, observer_config, disease, event_time, age_bins)
-
-    values = set(counts.values())
-    assert len(values) == 1
-    expected_value = len(pop) / len(states)
-    if observer_config["by_sex"]:
-        expected_value /= 2
-    if observer_config["by_age"]:
-        expected_value /= len(age_bins)
-    assert np.isclose(values.pop(), expected_value)
-
-    # Doubling pop should double counts
-    pop = pd.concat([pop, pop], axis=0, ignore_index=True)
-
-    counts = get_disease_event_counts(pop, observer_config, disease, event_time, age_bins)
-
-    values = set(counts.values())
-    assert len(values) == 1
-    assert np.isclose(values.pop(), 2 * expected_value)
-
-
-def test_get_lived_in_span():
-    dt = pd.Timedelta(days=5)
-    reference_t = pd.Timestamp("1-10-2010")
-
-    early_1 = reference_t - 2 * dt
-    early_2 = reference_t - dt
-
-    t_start = reference_t
-
-    mid_1 = reference_t + dt
-    mid_2 = reference_t + 2 * dt
-
-    t_end = reference_t + 3 * dt
-
-    late_1 = reference_t + 4 * dt
-    late_2 = reference_t + 5 * dt
-
-    # 28 combinations, six of which are entirely out of the time span
-    times = [early_1, early_2, t_start, mid_1, mid_2, t_end, late_1, late_2]
-    starts, ends = zip(*combinations(times, 2))
-    pop = pd.DataFrame({"age": to_years(10 * dt), "entrance_time": starts, "exit_time": ends})
-
-    lived_in_span = get_lived_in_span(pop, t_start, t_end)
-    # Indices here are from the combinatorics math. They represent
-    # 0: (early_1, early_2)
-    # 1: (early_1, t_start)
-    # 7: (early_2, t_start)
-    # 25: (t_end, late_1)
-    # 26: (t_end, late_2)
-    # 27: (late_1, late_2)
-    assert {0, 1, 7, 25, 26, 27}.intersection(lived_in_span.index) == set()
-
-    exit_before_span_end = lived_in_span.exit_time <= t_end
-    assert np.all(
-        lived_in_span.loc[exit_before_span_end, "age_at_span_end"]
-        == lived_in_span.loc[exit_before_span_end, "age"]
-    )
-
-    exit_after_span_end = ~exit_before_span_end
-    age_at_end = lived_in_span.age - to_years(lived_in_span.exit_time - t_end)
-    assert np.all(
-        lived_in_span.loc[exit_after_span_end, "age_at_span_end"]
-        == age_at_end.loc[exit_after_span_end]
-    )
-
-    enter_after_span_start = lived_in_span.entrance_time >= t_start
-    age_at_start = lived_in_span.age - to_years(
-        lived_in_span.exit_time - lived_in_span.entrance_time
-    )
-    assert np.all(
-        lived_in_span.loc[enter_after_span_start, "age_at_span_start"]
-        == age_at_start.loc[enter_after_span_start]
-    )
-
-    enter_before_span_start = ~enter_after_span_start
-    age_at_start = lived_in_span.age - to_years(lived_in_span.exit_time - t_start)
-    assert np.all(
-        lived_in_span.loc[enter_before_span_start, "age_at_span_start"]
-        == age_at_start.loc[enter_before_span_start]
-    )
-
-
-def test_get_lived_in_span_no_one_in_span():
-    dt = pd.Timedelta(days=365.25)
-    t_start = pd.Timestamp("1-1-2010")
-    t_end = t_start + dt
-
-    pop = pd.DataFrame(
-        {"entrance_time": t_start - 2 * dt, "exit_time": t_start - dt, "age": range(100)}
-    )
-    lived_in_span = get_lived_in_span(pop, t_start, t_end)
-    assert lived_in_span.empty
-
-    pop = pd.DataFrame(
-        {"entrance_time": t_end + dt, "exit_time": t_end + 2 * dt, "age": range(100)}
-    )
-    lived_in_span = get_lived_in_span(pop, t_start, t_end)
-    assert lived_in_span.empty
-
-
-def test_get_person_time_in_span(ages_and_bins, observer_config):
-    _, age_bins = ages_and_bins
-    start = int(age_bins.age_start.min())
-    end = int(age_bins.age_end.max())
-    n_ages = len(list(range(start, end)))
-    n_bins = len(age_bins)
-    segments_per_age = [(i + 1) * (n_ages - i) for i in range(n_ages)]
-    ages_per_bin = n_ages // n_bins
-    age_bins["expected_time"] = [
-        sum(segments_per_age[ages_per_bin * i : ages_per_bin * (i + 1)])
-        for i in range(n_bins)
-    ]
-
-    age_starts, age_ends = zip(*combinations(range(start, end + 1), 2))
-    women = pd.DataFrame(
-        {"age_at_span_start": age_starts, "age_at_span_end": age_ends, "sex": "Female"}
-    )
-    men = women.copy()
-    men.loc[:, "sex"] = "Male"
-
-    lived_in_span = (
-        pd.concat([women, men], ignore_index=True).sample(frac=1).reset_index(drop=True)
-    )
-    base_filter = QueryString("")
-    span_key = get_output_template(**observer_config).substitute(
-        measure="person_time", year=2019
-    )
-
-    pt = get_person_time_in_span(
-        lived_in_span, base_filter, span_key, observer_config, age_bins
-    )
-
-    if observer_config["by_age"]:
-        for group, age_bin in age_bins.iterrows():
-            group_pt = sum(set([v for k, v in pt.items() if f"in_age_group_{group}" in k]))
-            if observer_config["by_sex"]:
-                assert group_pt == age_bin.expected_time
-            else:
-                assert group_pt == 2 * age_bin.expected_time
-    else:
-        group_pt = sum(set(pt.values()))
-        if observer_config["by_sex"]:
-            assert group_pt == age_bins.expected_time.sum()
-        else:
-            assert group_pt == 2 * age_bins.expected_time.sum()
-
-
-def test_get_deaths(ages_and_bins, sexes, observer_config):
-    alive = ["dead", "alive"]
-    ages, age_bins = ages_and_bins
-    exit_times = [pd.Timestamp("1-1-2012"), pd.Timestamp("1-1-2013")]
-    causes = ["cause_a", "cause_b"]
-
-    pop = pd.DataFrame(
-        list(product(alive, ages, sexes, exit_times, causes)),
-        columns=["alive", "age", "sex", "exit_time", "cause_of_death"],
-    )
-    # Shuffle the rows
-    pop = pop.sample(frac=1).reset_index(drop=True)
-
-    deaths = get_deaths(
-        pop,
-        observer_config,
-        pd.Timestamp("1-1-2010"),
-        pd.Timestamp("1-1-2015"),
-        age_bins,
-        causes,
-    )
-    values = set(deaths.values())
-
-    expected_value = len(pop) / (len(causes) * len(alive))
-    if observer_config["by_year"]:
-        assert (
-            len(values) == 2
-        )  # Uniform across bins with deaths, 0 in year bins without deaths
-        expected_value /= 2
-    else:
-        assert len(values) == 1
-    value = max(values)
-    if observer_config["by_sex"]:
-        expected_value /= 2
-    if observer_config["by_age"]:
-        expected_value /= len(age_bins)
-    assert np.isclose(value, expected_value)
-
-    # Doubling pop should double counts
-    pop = pd.concat([pop, pop], axis=0, ignore_index=True)
-
-    deaths = get_deaths(
-        pop,
-        observer_config,
-        pd.Timestamp("1-1-2010"),
-        pd.Timestamp("1-1-2015"),
-        age_bins,
-        causes,
-    )
-    values = set(deaths.values())
-
-    if observer_config["by_year"]:
-        assert (
-            len(values) == 2
-        )  # Uniform across bins with deaths, 0 in year bins without deaths
-    else:
-        assert len(values) == 1
-    value = max(values)
-    assert np.isclose(value, 2 * expected_value)
-
-
-def test_get_years_of_life_lost(ages_and_bins, sexes, observer_config):
-    alive = ["dead", "alive"]
-    ages, age_bins = ages_and_bins
-    exit_times = [pd.Timestamp("1-1-2012"), pd.Timestamp("1-1-2013")]
-    causes = ["cause_a", "cause_b"]
-
-    pop = pd.DataFrame(
-        list(product(alive, ages, sexes, exit_times, causes)),
-        columns=["alive", "age", "sex", "exit_time", "cause_of_death"],
-    )
-    # Shuffle the rows
-    pop = pop.sample(frac=1).reset_index(drop=True)
-
-    def life_expectancy(index):
-        return pd.Series(1, index=index)
-
-    ylls = get_years_of_life_lost(
-        pop,
-        observer_config,
-        pd.Timestamp("1-1-2010"),
-        pd.Timestamp("1-1-2015"),
-        age_bins,
-        life_expectancy,
-        causes,
-    )
-    values = set(ylls.values())
-
-    expected_value = len(pop) / (len(causes) * len(alive))
-    if observer_config["by_year"]:
-        assert (
-            len(values) == 2
-        )  # Uniform across bins with deaths, 0 in year bins without deaths
-        expected_value /= 2
-    else:
-        assert len(values) == 1
-    value = max(values)
-    if observer_config["by_sex"]:
-        expected_value /= 2
-    if observer_config["by_age"]:
-        expected_value /= len(age_bins)
-    assert np.isclose(value, expected_value)
-
-    # Doubling pop should double counts
-    pop = pd.concat([pop, pop], axis=0, ignore_index=True)
-
-    ylls = get_years_of_life_lost(
-        pop,
-        observer_config,
-        pd.Timestamp("1-1-2010"),
-        pd.Timestamp("1-1-2015"),
-        age_bins,
-        life_expectancy,
-        causes,
-    )
-    values = set(ylls.values())
-
-    if observer_config["by_year"]:
-        assert (
-            len(values) == 2
-        )  # Uniform across bins with deaths, 0 in year bins without deaths
-    else:
-        assert len(values) == 1
-    value = max(values)
-    assert np.isclose(value, 2 * expected_value)
-
-
-def test_get_years_lived_with_disability(ages_and_bins, sexes, observer_config):
-    alive = ["dead", "alive"]
-    ages, age_bins = ages_and_bins
-    causes = ["cause_a", "cause_b"]
-    cause_a = ["susceptible_to_cause_a", "cause_a"]
-    cause_b = ["susceptible_to_cause_b", "cause_b"]
-    year = 2010
-    step_size = pd.Timedelta(days=7)
-
-    pop = pd.DataFrame(
-        list(product(alive, ages, sexes, cause_a, cause_b)),
-        columns=["alive", "age", "sex"] + causes,
-    )
-    # Shuffle the rows
-    pop = pop.sample(frac=1).reset_index(drop=True)
-
-    def disability_weight(cause):
-        def inner(index):
-            sub_pop = pop.loc[index]
-            return pd.Series(1, index=index) * (sub_pop[cause] == cause)
-
-        return inner
-
-    disability_weights = {cause: disability_weight(cause) for cause in causes}
-
-    ylds = get_years_lived_with_disability(
-        pop, observer_config, year, step_size, age_bins, disability_weights, causes
-    )
-
-    values = set(ylds.values())
-    assert len(values) == 1
-    states_per_cause = len(cause_a)
-    expected_value = len(pop) / (len(alive) * states_per_cause) * to_years(step_size)
-    if observer_config["by_sex"]:
-        expected_value /= 2
-    if observer_config["by_age"]:
-        expected_value /= len(age_bins)
-    assert np.isclose(values.pop(), expected_value)
-
-    # Doubling pop should double person time
-    pop = pd.concat([pop, pop], axis=0, ignore_index=True)
-
-    ylds = get_years_lived_with_disability(
-        pop, observer_config, year, step_size, age_bins, disability_weights, causes
-    )
-
-    values = set(ylds.values())
-    assert len(values) == 1
-    assert np.isclose(values.pop(), 2 * expected_value)
-
-
-@pytest.mark.parametrize(
-    "age_start, exit_age, result_age_end_values, result_age_start_values",
-    [
-        (2, 5, {4, 5}, {2, 4}),
-        (0, None, {1, 4, 6}, {0, 1, 4}),
-        (1, 4, {4}, {1}),
-        (1, 3, {3}, {1}),
-        (0.8, 6, {1, 4, 6}, {0.8, 1, 4}),
-    ],
-)
-def test_get_age_bins(
-    builder, base_config, age_start, exit_age, result_age_end_values, result_age_start_values
-):
-    base_config.update(
-        {"population": {"age_start": age_start, "exit_age": exit_age}}, **metadata(__file__)
-    )
-    builder.configuration = base_config
-    df = get_age_bins(builder)
-    assert set(df.age_end) == result_age_end_values
-    assert set(df.age_start) == result_age_start_values