Fix #1419. DataCollector accepts an arbitrary schedule at creation (d…

docs/useful-snippets/snippets.rst

@@ -58,3 +58,67 @@ code as follows.
 If you would still like to run your code in Jupyter you will need to adjust the cell as noted above. Then you can
 you can add the `nbmultitask library <(https://nbviewer.org/github/micahscopes/nbmultitask/blob/39b6f31b047e8a51a0fcb5c93ae4572684f877ce/examples.ipynb)>`__
 or look at this `stackoverflow <https://stackoverflow.com/questions/50937362/multiprocessing-on-python-3-jupyter>`__.
+
+Using multiple schedules and data collectors
+-------------
+
+One may occasionally want multiple schedules with distinct sets of agents and collect data on each of those schedules independently while allowing the agents to interact in the same model. This would be particularly useful if the agents in your model do not all have the same attributes and you want to collect data on them in groups. You can create a simple scheduler class to mix the agents in each schedule and data collect on each schedule independently.
+
+.. code:: python
+
+    from random import shuffle
+
+    # simple schedule mixer
+    class ScheduleMixerRandom:
+        """
+        Be careful: ScheduleMixerRandom as written here would not capture changes to
+        the agent lists in the sub-schedules. The user would need to be mindful if
+        this functionality were to be included.
+        """
+
+        def __init__(self, schedules):
+            self._agents = {}
+            for s in schedules:
+                for a in s.agents:
+                    self._agents[
+                        a.unique_id
+                    ] = a  # each agent must have a unique id across models
+
+        def step(self):
+            # do stuff with self._agents
+            ...
+
+    class MyModel(Model):
+        def __init__(self, ...):
+            super().__init__()
+            self.schedule_one = BaseScheduler()
+            self.schedule_two = BaseScheduler()
+
+            # each DataCollector instance is associated with a specific schedule (and therefore its agents)
+            self.datacollector_one = DataCollector(
+                agent_reporters={"value_one": "value_one"}, schedule=self.schedule_one
+            )
+            self.datacollector_two = DataCollector(
+                agent_reporters={"value_two": "value_two"}, schedule=self.schedule_two
+            )
+
+            self.schedule_mixer = ScheduleMixerRandom(
+                schedules=[self.schedule_one, self.schedule_two]
+            )
+
+        def step():
+            self.schedule_mixer.step()  # calls step() on each agent in both schedules, mixed together
+            self.datacollector_one.collect(
+                self
+            )  # collects data on the agents in schedule_one
+            self.datacollector_two.collect(
+                self
+            )  # collects data on the agents in schedule_two
+
+
+To get the data at the end of the model, simply call:
+
+.. code:: python
+
+  schedule_one_agent_data = model.datacollector_one.get_agent_vars_dataframe()
+  schedule_two_agent_data = model.datacollector_two.get_agent_vars_dataframe()

mesa/datacollection.py

@@ -37,7 +37,6 @@
 """
 from functools import partial
 import itertools
-from operator import attrgetter
 import pandas as pd
 import types
 
@@ -53,7 +52,9 @@ class DataCollector:
 
     """
 
-    def __init__(self, model_reporters=None, agent_reporters=None, tables=None):
+    def __init__(
+        self, model_reporters=None, agent_reporters=None, tables=None, schedule=None
+    ):
         """Instantiate a DataCollector with lists of model and agent reporters.
         Both model_reporters and agent_reporters accept a dictionary mapping a
         variable name to either an attribute name, or a method.
@@ -76,6 +77,7 @@ def __init__(self, model_reporters=None, agent_reporters=None, tables=None):
             model_reporters: Dictionary of reporter names and attributes/funcs
             agent_reporters: Dictionary of reporter names and attributes/funcs.
             tables: Dictionary of table names to lists of column names.
+            schedule: A scheduler from the mesa.time module. If not supplied, this defaults to `model.schedule`.
 
         Notes:
             If you want to pickle your model you must not use lambda functions.
@@ -100,6 +102,8 @@ class attributes of model
         self._agent_records = {}
         self.tables = {}
 
+        self.schedule = schedule
+
         if model_reporters is not None:
             for name, reporter in model_reporters.items():
                 self._new_model_reporter(name, reporter)
@@ -151,28 +155,28 @@ 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, schedule):
         """Record agents data in a mapping of functions and agents."""
         rep_funcs = self.agent_reporters.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]
-            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
+        def get_reports(agent):
+            _prefix = (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, schedule.agents)
         return agent_records
 
     def _reporter_decorator(self, reporter):
         return reporter()
 
     def collect(self, model):
         """Collect all the data for the given model object."""
+        if self.schedule is None:
+            schedule = model.schedule
+        else:
+            schedule = self.schedule
+
         if self.model_reporters:
 
             for var, reporter in self.model_reporters.items():
@@ -189,8 +193,8 @@ def collect(self, model):
                     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)
+            agent_records = self._record_agents(schedule)
+            self._agent_records[schedule.steps] = list(agent_records)
 
     def add_table_row(self, table_name, row, ignore_missing=False):
         """Add a row dictionary to a specific table.