Merge pull request #66 from Princeton-CDH/collect-agent-totals

 Include agent totals per risk level in model data collection

simulatingrisk/hawkdove/model.py

@@ -259,6 +259,8 @@ def get_data_collector_options(self):
                 "percent_hawk": "percent_hawk",
                 "rolling_percent_hawk": "rolling_percent_hawk",
                 "status": "status",
+                # explicitly track total agents, instead of inferring from grid size
+                "total_agents": "num_agents",
             },
             "agent_reporters": {
                 "risk_level": "risk_level",

simulatingrisk/hawkdovemulti/model.py

@@ -287,18 +287,52 @@ def get_data_collector_options(self):
         # in addition to common hawk/dove data points,
         # we want to include population risk category
         opts = super().get_data_collector_options()
-        opts["model_reporters"]["population_risk_category"] = "population_risk_category"
+        model_reporters = {"population_risk_category": "population_risk_category"}
+        for risk_level in range(self.min_risk_level, self.max_risk_level + 1):
+            field = f"total_r{risk_level}"
+            model_reporters[field] = field
+
+        opts["model_reporters"].update(model_reporters)
         return opts
 
+    def step(self):
+        # delete cached property before the next round begins,
+        # so we recalcate values for current round before collecting data
+        try:
+            del self.total_per_risk_level
+        except AttributeError:
+            # property hasn't been set yet on the first round, ok to ignore
+            pass
+        super().step()
+
+    @cached_property
+    def total_per_risk_level(self):
+        # tally the number of agents for each risk level
+        return Counter([a.risk_level for a in self.schedule.agents])
+
+    def __getattr__(self, attr):
+        # support dynamic properties for data collection on total by risk level
+        if attr.startswith("total_r"):
+            try:
+                r = int(attr.replace("total_r", ""))
+                # only handle risk levels that are in bounds
+                if r > self.max_risk_level or r < self.min_risk_level:
+                    raise AttributeError
+                return self.total_per_risk_level[r]
+            except ValueError:
+                # ignore and throw attribute error
+                pass
+
+        raise AttributeError
+
     @property
     def population_risk_category(self):
         # calculate a category of risk distribution for the population
         # based on the proportion of agents in different risk categories
         # (categorization scheme defined by LB)
 
-        # tally the number of agents with each risk level
-        risk_counts = Counter([a.risk_level for a in self.schedule.agents])
         # count the number of agents in three groups:
+        risk_counts = self.total_per_risk_level
         #  Risk-inclined (RI) : r = 0, 1, 2
         #  Risk-moderate (RM): r = 3, 4, 5
         #  Risk-avoidant (RA): r = 6, 7, 8

tests/test_hawkdovemulti.py

@@ -82,6 +82,71 @@ def test_adjustment_round(params, expect_adjust_step):
             assert not model.adjustment_round
 
 
+def test_total_per_risk_level():
+    model = HawkDoveMultipleRiskModel(3)
+    model.schedule = Mock()
+    # add a few agents with different risk levels
+    mock_agents = [
+        Mock(risk_level=0),
+        Mock(risk_level=1),
+        Mock(risk_level=1),
+        Mock(risk_level=2),
+        Mock(risk_level=2),
+        Mock(risk_level=2),
+        Mock(risk_level=5),
+    ]
+    model.schedule.agents = mock_agents
+
+    totals = model.total_per_risk_level
+    assert totals[0] == 1
+    assert totals[1] == 2
+    assert totals[2] == 3
+    assert totals[4] == 0
+    assert totals[5] == 1
+    assert totals[8] == 0
+
+    # check caching works as desired
+    mock_agents.append(Mock(risk_level=8))
+    model.schedule.agents = mock_agents
+    # cached total should not change even though agents have changed
+    assert model.total_per_risk_level[8] == 0
+    # step should reset catched property
+    with patch("builtins.super"):
+        model.step()
+    # now the count should be updated
+    assert model.total_per_risk_level[8] == 1
+
+
+def test_total_rN_attr():
+    # dynamic attributes to get total per risk level, for data collection
+    model = HawkDoveMultipleRiskModel(3)
+    model.schedule = Mock()
+    # add a few agents with different risk levels
+    model.schedule.agents = [
+        Mock(risk_level=0),
+        Mock(risk_level=1),
+        Mock(risk_level=1),
+        Mock(risk_level=2),
+        Mock(risk_level=2),
+        Mock(risk_level=2),
+    ]
+    assert model.total_r0 == 1
+    assert model.total_r1 == 2
+    assert model.total_r2 == 3
+    assert model.total_r4 == 0
+
+    # error handling
+    # - non-numeric
+    with pytest.raises(AttributeError):
+        model.total_rfour
+    # - out of bounds
+    with pytest.raises(AttributeError):
+        model.total_r23
+    # - unsupported attribute
+    with pytest.raises(AttributeError):
+        model.some_other_total
+
+
 def test_population_risk_category():
     model = HawkDoveMultipleRiskModel(3)
     model.schedule = Mock()
@@ -90,15 +155,18 @@ def test_population_risk_category():
     model.schedule.agents = [Mock(risk_level=0), Mock(risk_level=1), Mock(risk_level=2)]
     assert model.population_risk_category == RiskState.c1
     # three risk-inclined agents and one risk moderate
+    del model.total_per_risk_level  #  reset cached property
     model.schedule.agents.append(Mock(risk_level=4))
     assert model.population_risk_category == RiskState.c2
 
     # majority risk moderate
     model.schedule.agents = [Mock(risk_level=4), Mock(risk_level=5), Mock(risk_level=6)]
+    del model.total_per_risk_level  #  reset cached property
     assert model.population_risk_category == RiskState.c7
 
     # majority risk avoidant
     model.schedule.agents = [Mock(risk_level=7), Mock(risk_level=8), Mock(risk_level=9)]
+    del model.total_per_risk_level  #  reset cached property
     assert model.population_risk_category == RiskState.c12