Merge pull request #35 from ihmeuw/feature/status_column

add the status column in population table
ihmeuw · Jul 11, 2018 · d7d324f · d7d324f
2 parents 8847437 + 97a9926
commit d7d324f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/tests/framework/test_engine.py b/tests/framework/test_engine.py
@@ -190,6 +190,7 @@ def test_SimulationContext_finalize(base_config, components):
 def test_SimulationContext_report(base_config, components):
     sim = SimulationContext(base_config, components)
     sim.setup()
+    sim.initialize_simulants()
     metrics = sim.report()
     assert metrics['test'] == len([c for c in sim.component_manager._components if isinstance(c, MockComponentB)])
 

diff --git a/vivarium/framework/population.py b/vivarium/framework/population.py
@@ -172,13 +172,14 @@ class PopulationManager:
 
     def __init__(self):
         self._population = pd.DataFrame()
-        self._population_initializers = []
+        self._population_initializers = [(self.create_status_column, ['tracked'], [])]
         self._initializers_ordered = False
         self.growing = False
 
     def setup(self, builder):
         self.clock = builder.time.clock()
         self.step_size = builder.time.step_size()
+        builder.value.register_value_modifier('metrics', modifier=self.metrics)
 
     def get_view(self, columns: Sequence[str], query: str=None) -> PopulationView:
         """Return a configured PopulationView
@@ -190,15 +191,22 @@ def get_view(self, columns: Sequence[str], query: str=None) -> PopulationView:
         generated column names that aren't known at definition time. Otherwise
         components should use ``uses_columns``.
         """
+        if 'tracked' not in columns:
+            query_with_track = query + 'and tracked == True' if query else 'tracked == True'
+            return PopulationView(self, columns, query_with_track)
         return PopulationView(self, columns, query)
 
     def register_simulant_initializer(self, initializer: Callable,
                                       creates_columns: Sequence[str]=(), requires_columns: Sequence[str]=()):
-        self._population_initializers.append((initializer, creates_columns, requires_columns))
+        self._population_initializers.append((initializer, creates_columns, tuple(requires_columns)+('tracked',)))
 
     def get_simulant_creator(self) -> Callable:
         return self._create_simulants
 
+    def create_status_column(self, pop_data):
+        status = pd.Series(True, index=pop_data.index)
+        self.get_view(['tracked']).update(status)
+
     def _order_initializers(self) -> None:
         unordered_initializers = deque(self._population_initializers)
         starting_length = -1
@@ -243,6 +251,16 @@ def _create_simulants(self, count: int, population_configuration: Mapping[str, A
         self.growing = False
         return index
 
+    def metrics(self, index, metrics):
+        population = self.get_view(['tracked']).get(index)
+        untracked = population[~population.tracked]
+        tracked = population[population.tracked]
+
+        metrics['total_population__untracked'] = len(untracked)
+        metrics['total_population__tracked'] = len(tracked)
+        metrics['total_population'] = len(untracked)+len(tracked)
+        return metrics
+
     @property
     def population(self) -> pd.DataFrame:
         return self._population.copy()

diff --git a/vivarium/test_util.py b/vivarium/test_util.py
@@ -76,8 +76,7 @@ def _build_population(core_population, location, randomness_stream):
         {'age': core_population['age'],
          'entrance_time': core_population['entrance_time'],
          'sex': randomness_stream.choice(index, ['Male', 'Female'], additional_key='sex_choice'),
-         'alive': pd.Series('alive', index=index).astype(
-             pd.api.types.CategoricalDtype(categories=['alive', 'dead', 'untracked'], ordered=False)),
+         'alive': pd.Series('alive', index=index),
          'location': location,
          'exit_time': pd.NaT, },
         index=index)
@@ -93,8 +92,7 @@ def _non_crn_build_population(index, age_start, age_end, location, creation_time
     population = pd.DataFrame(
         {'age': age,
          'sex': randomness_stream.choice(index, ['Male', 'Female'], additional_key='sex_choice'),
-         'alive': pd.Series('alive', index=index).astype(
-             pd.api.types.CategoricalDtype(categories=['alive', 'dead', 'untracked'], ordered=False)),
+         'alive': pd.Series('alive', index=index),
          'location': location,
          'entrance_time': creation_time,
          'exit_time': pd.NaT, },