properly create the population_view in RateTransition (#482)

ihmeuw · Sep 10, 2024 · e6a49e1 · e6a49e1
1 parent 70d8dd8
commit e6a49e1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,3 +1,7 @@
+**3.0.7 - 09/10/24**
+
+ - Properly create the population_view in RateTransition
+
 **3.0.6 - 09/04/24**
 
  - Fix bug that was occurring when RiskEffect's rr_source was a float or DataFrame

diff --git a/src/vivarium_public_health/disease/transition.py b/src/vivarium_public_health/disease/transition.py
@@ -7,7 +7,7 @@
 
 """
 
-from typing import TYPE_CHECKING, Any, Callable, Dict, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import pandas as pd
 from vivarium.framework.engine import Builder
@@ -48,6 +48,10 @@ def configuration_defaults(self) -> Dict[str, Any]:
             },
         }
 
+    @property
+    def columns_required(self) -> Optional[List[str]]:
+        return ["alive"]
+
     @property
     def transition_rate_pipeline_name(self) -> str:
         if "incidence_rate" in self._get_data_functions:
@@ -100,8 +104,6 @@ def setup(self, builder: Builder) -> None:
             preferred_post_processor=union_post_processor,
         )
 
-        self.population_view = builder.population.get_view(["alive"])
-
     #################
     # Setup methods #
     #################