ENH: use pop_data.user_data.get pattern in examples.disease_model.Bas…

…ePopulation This makes it easier to demonstrate a minimal fertility model
ihmeuw · Jan 5, 2023 · 3cb1c9e · 3cb1c9e
1 parent 82f5027
commit 3cb1c9e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/vivarium/examples/disease_model/population.py b/src/vivarium/examples/disease_model/population.py
@@ -94,8 +94,9 @@ def on_initialize_simulants(self, pop_data: SimulantData):
 
         """
 
-        age_start = self.config.population.age_start
-        age_end = self.config.population.age_end
+        age_start = pop_data.user_data.get("age_start", self.config.population.age_start)
+        age_end = pop_data.user_data.get("age_end", self.config.population.age_end)
+
         if age_start == age_end:
             age_window = pop_data.creation_window / pd.Timedelta(days=365)
         else: