PoC: Multiple-agent types scheduling and datacollection

mesa/datacollection.py

@@ -40,6 +40,7 @@
 from operator import attrgetter
 import pandas as pd
 import types
+import statistics
 
 
 class DataCollector:
@@ -95,20 +96,31 @@ class attributes of model
             {"Model_Function":[function, [param_1, param_2]]}
 
         """
-        self.model_reporters = {}
-        self.agent_reporters = {}
+        self.agent_types = agent_reporters.keys()
 
+        self.model_reporters = {}
         self.model_vars = {}
-        self._agent_records = {}
-        self.tables = {}
 
         if model_reporters is not None:
             for name, reporter in model_reporters.items():
                 self._new_model_reporter(name, reporter)
 
-        if agent_reporters is not None:
-            for name, reporter in agent_reporters.items():
-                self._new_agent_reporter(name, reporter)
+        self.agent_reporters = {}
+        self._agent_records = {}
+        self.agent_attr_index = {}
+        self.agent_name_index = {}
+
+        for agent_type in self.agent_types:
+            self.agent_reporters[agent_type] = {}
+            self._agent_records[agent_type] = {}
+            self.agent_name_index[agent_type] = {}
+        self.tables = {}
+
+        for agent_type in self.agent_types:
+            if agent_reporters[agent_type] is not None:
+                for name, reporter in agent_reporters[agent_type].items():
+                    self.agent_name_index[agent_type][name] = reporter
+                    self._new_agent_reporter(name, reporter, agent_type)
 
         if tables is not None:
             for name, columns in tables.items():
@@ -127,7 +139,7 @@ def _new_model_reporter(self, name, reporter):
         self.model_reporters[name] = reporter
         self.model_vars[name] = []
 
-    def _new_agent_reporter(self, name, reporter):
+    def _new_agent_reporter(self, name, reporter, agent_type):
         """Add a new agent-level reporter to collect.
 
         Args:
@@ -140,7 +152,7 @@ def _new_agent_reporter(self, name, reporter):
             attribute_name = reporter
             reporter = partial(self._getattr, reporter)
             reporter.attribute_name = attribute_name
-        self.agent_reporters[name] = reporter
+        self.agent_reporters[agent_type][name] = reporter
 
     def _new_table(self, table_name, table_columns):
         """Add a new table that objects can write to.
@@ -153,21 +165,24 @@ def _new_table(self, table_name, table_columns):
         new_table = {column: [] for column in table_columns}
         self.tables[table_name] = new_table
 
-    def _record_agents(self, model):
+    def _record_agents(self, model, agent_type):
         """Record agents data in a mapping of functions and agents."""
-        rep_funcs = self.agent_reporters.values()
+        rep_funcs = self.agent_reporters[agent_type].values()
         if all([hasattr(rep, "attribute_name") for rep in rep_funcs]):
             prefix = ["model.schedule.steps", "unique_id"]
             attributes = [func.attribute_name for func in rep_funcs]
+            self.agent_attr_index[agent_type] = {k: v for v, k in enumerate(prefix + attributes)}
+            print(self.agent_attr_index[agent_type])
             get_reports = attrgetter(*prefix + attributes)
+
         else:
 
             def get_reports(agent):
                 _prefix = (agent.model.schedule.steps, agent.unique_id)
                 reports = tuple(rep(agent) for rep in rep_funcs)
                 return _prefix + reports
 
-        agent_records = map(get_reports, model.schedule.agents)
+        agent_records = map(get_reports, model.schedule.agents_by_type[agent_type].values())
         return agent_records
 
     def _reporter_decorator(self, reporter):
@@ -190,9 +205,10 @@ def collect(self, model):
                 else:
                     self.model_vars[var].append(self._reporter_decorator(reporter))
 
-        if self.agent_reporters:
-            agent_records = self._record_agents(model)
-            self._agent_records[model.schedule.steps] = list(agent_records)
+        for agent_type in self.agent_types:
+            if self.agent_reporters[agent_type]:
+                agent_records = self._record_agents(model, agent_type)
+                self._agent_records[agent_type][model.schedule.steps] = list(agent_records)
 
     def add_table_row(self, table_name, row, ignore_missing=False):
         """Add a row dictionary to a specific table.
@@ -229,21 +245,21 @@ def get_model_vars_dataframe(self):
         """
         return pd.DataFrame(self.model_vars)
 
-    def get_agent_vars_dataframe(self):
+    def get_agent_vars_dataframe(self, agent_type):
         """Create a pandas DataFrame from the agent variables.
 
         The DataFrame has one column for each variable, with two additional
         columns for tick and agent_id.
 
         """
-        all_records = itertools.chain.from_iterable(self._agent_records.values())
-        rep_names = [rep_name for rep_name in self.agent_reporters]
+        all_records = itertools.chain.from_iterable(self._agent_records[agent_type].values())
+        rep_names = [rep_name for rep_name in self.agent_reporters[agent_type]]
 
         df = pd.DataFrame.from_records(
             data=all_records,
-            columns=["Step", "AgentID"] + rep_names,
+            columns=["Step", f"{agent_type}_AgentID"] + rep_names,
         )
-        df = df.set_index(["Step", "AgentID"])
+        df = df.set_index(["Step", f"{agent_type}_AgentID"])
         return df
 
     def get_table_dataframe(self, table_name):
@@ -256,3 +272,37 @@ def get_table_dataframe(self, table_name):
         if table_name not in self.tables:
             raise Exception("No such table.")
         return pd.DataFrame(self.tables[table_name])
+
+    def get_agent_metric(self, var_name, agent_type, metric="mean"):
+        """Get a single aggegrated value from an agent variable.
+
+        Args:
+            var_name: The name of the variable to aggegrate.
+            agent_type: Agent type class
+            metric: Statistics metric to be used (default: mean)
+                    all functions from built-in statistics module are supported
+                    as well as "min", "max", "sum" and "len"
+
+        """
+        # Get the reporter from the name
+        reporter = self.agent_name_index[agent_type][var_name]
+
+        # Get the index of the reporter
+        attr_index = self.agent_attr_index[agent_type][reporter]
+
+        # Create a dictionary with all attributes from all agents
+        attr_dict = self._agent_records[agent_type]
+
+        # Get the values from all agents in a list
+        values_tuples = list(attr_dict.values())[-1]
+
+        # Get the correct value using the attribute index
+        values = [value_tuple[attr_index] for value_tuple in values_tuples]
+
+        # Calculate the metric among all agents (mean by default)
+        if metric in ["min", "max", "sum", "len"]:
+            value = eval(f"{metric}(values)")
+        else:
+            stat_function = getattr(statistics, metric)
+            value = stat_function(values)
+        return value

mesa/time.py

@@ -21,7 +21,7 @@
     model has taken.
 """
 
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 
 # mypy
 from typing import Dict, Iterator, List, Optional, Union
@@ -198,3 +198,74 @@ def step(self) -> None:
             self.time += self.stage_time
 
         self.steps += 1
+
+class RandomActivationByType(RandomActivation):
+    """
+    A scheduler which activates each type of agent once per step, in random
+    order, with the order reshuffled every step.
+
+    This is equivalent to the NetLogo 'ask breed...' and is generally the
+    default behavior for an ABM.
+
+    Assumes that all agents have a step() method.
+    """
+
+    def __init__(self, model):
+        super().__init__(model)
+        self.agents_by_type = defaultdict(dict)
+
+    def add(self, agent):
+        """
+        Add an Agent object to the schedule
+
+        Args:
+            agent: An Agent to be added to the schedule.
+        """
+
+        self._agents[agent.unique_id] = agent
+        agent_class = type(agent)
+        self.agents_by_type[agent_class][agent.unique_id] = agent
+
+    def remove(self, agent):
+        """
+        Remove all instances of a given agent from the schedule.
+        """
+
+        del self._agents[agent.unique_id]
+
+        agent_class = type(agent)
+        del self.agents_by_type[agent_class][agent.unique_id]
+
+    def step(self, by_type=True):
+        """
+        Executes the step of each agent type, one at a time, in random order.
+
+        Args:
+            by_type: If True, run all agents of a single type before running
+                      the next one.
+        """
+        if by_type:
+            for agent_class in self.agents_by_type:
+                self.step_type(agent_class)
+            self.steps += 1
+            self.time += 1
+        else:
+            super().step()
+
+    def step_type(self, type):
+        """
+        Shuffle order and run all agents of a given type.
+
+        Args:
+            type: Class object of the type to run.
+        """
+        agent_keys = list(self.agents_by_type[type].keys())
+        self.model.random.shuffle(agent_keys)
+        for agent_key in agent_keys:
+            self.agents_by_type[type][agent_key].step()
+
+    def get_type_count(self, type_class):
+        """
+        Returns the current number of agents of certain type in the queue.
+        """
+        return len(self.agents_by_type[type_class].values())

mesa/visualization/modules/ChartVisualization.py

@@ -34,8 +34,6 @@ class ChartModule(VisualizationElement):
                                       data_collector_name="datacollector")
 
     TODO:
-        Have it be able to handle agent-level variables as well.
-
         More Pythonic customization; in particular, have both series-level and
         chart-level options settable in Python, and passed to the front-end
         the same way that "Color" is currently.
@@ -78,9 +76,22 @@ def render(self, model):
 
         for s in self.series:
             name = s["Label"]
-            try:
-                val = data_collector.model_vars[name][-1]  # Latest value
-            except (IndexError, KeyError):
-                val = 0
+            entity = s["Type"]
+            if entity == "Model":
+                try:
+                    val = data_collector.model_vars[name][-1]  # Latest value
+                except (IndexError, KeyError):
+                    val = 0
+            elif entity == "Agent":
+                agent_dict = {e.__name__: e for e in list(model.schedule.agents_by_type.keys())}
+                agent_type = agent_dict[s["Agent_type"]]
+                if "Metric" in s.keys():
+                    metric = s["Metric"]
+                else:
+                    metric = "mean"
+                try:
+                    val = data_collector.get_agent_metric(name, agent_type, metric)
+                except (IndexError, KeyError):
+                    val = 0
             current_values.append(val)
         return current_values