Change assertion to DataTransformationError

ihmeuw · Jan 27, 2025 · 9167806 · 9167806
1 parent b0dd623
commit 9167806
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/src/vivarium_inputs/core.py b/src/vivarium_inputs/core.py
@@ -24,6 +24,7 @@
     MEASURES,
     MINIMUM_EXPOSURE_VALUE,
     DataDoesNotExistError,
+    DataTransformationError,
     InvalidQueryError,
     Population,
     gbd,
@@ -759,12 +760,14 @@ def _filter_relative_risk_to_cause_restrictions(data: pd.DataFrame) -> pd.DataFr
         cause = causes_map[cause]
         if measure == "cause_specific_mortality_rate":
             start, end = utilities.get_age_group_ids_by_restriction(cause, "yll")
+        elif measure == "incidence_rate":  
+            start, end = utilities.get_age_group_ids_by_restriction(cause, "yld")
         else:
-            assert (
-                measure == "incidence_rate",
-                "Affected measure must be incidence_rate or cause_specific_mortality_rate to apply cause restrictions",
+            raise DataTransformationError(
+                "Affected measure must be incidence_rate or cause_specific_mortality_rate "
+                "to apply cause restrictions to relative risks. "
+                "Check the morbidity and mortality columns."
             )
-            start, end = utilities.get_age_group_ids_by_restriction(cause, "yld")
         start_index = list(ordered_age_ids).index(start)
         end_index = list(ordered_age_ids).index(end)
         allowed_ids = ordered_age_ids[start_index : (end_index + 1)]