Repository: projectmesa/mesa

Issue:
Feature Request: Agent DataCollection Can't Handle Different Attributes in ActivationByType 
**Describe the Feature**
Receive attribute error when running an agent reporter through datacollector when using RandomActivationByType 

**Desired behavior**
I can collect against different attributes based on agent type

**To Reproduce**
Create a model with different agent types and attributes and try to collect against that agent attribute

**Additional context**
Feature would need to update at line 170 of datacollection.py `agent_records = map(get_reports, model.schedule.agents)`



Discussion:
@tpike3 I am unable to recreate the bug -- it looks like sugarscape_cg uses `RandomActivationByType`. Can you add it to that example & recreate it?


If you use this colab you will get the error https://github.com/edgeofchaos42/ComplexityExplorer/blob/main/Session_18_Data_Collector_agent.ipynb Attached is the landscape you will need to load into your google drive. or update the `sugar_distribution = ` line

Per the dev session this morning I should be able to add a custom function to filter out non-traders or at least return None, but so far unsuccessful on that code. 


[sugar-map.txt](https://github.com/projectmesa/mesa/files/9639657/sugar-map.txt)

I think a temporary workaround is to put the trade partners in the model reporter instead.
```
lambda m: {k: list(a.trade_partners) for k, a in m.schedule.agents_by_type[Trader].items()}
```

I used `list()` to copy the `a.trade_partners`.
This is indeed a (in my opinion big) limitation. I wrote a proof of concept implementation for it in https://github.com/projectmesa/mesa/pull/1142, and used it in a personal project. If anyone ones to built on it, feel free to use it in any way!
```
agent_reporters = {Cat: {"Weight": "weight", "Size": "size"}, Dog: {"Weight": "weight"}}
```

While this looks more concise for `RandomActivationByType`, it is a pattern that only works for this scheduler. We can't just modify `DataCollector` to expect more different types of schedulers. `DataCollector` will become more bloated. And this is not much more concise than the workaround I specified in my earlier comment. The lack of the feature is not a blocker in @tpike3's situation.

Either:
- Create a `DataCollectorRABT` (shorthand for `DataCollectorRandomActivationByType`) class that handles types for agent collector
- Tell people to put their type-specific data in the model collector, like in https://github.com/projectmesa/mesa/issues/1419#issuecomment-1257158578. Put the guide in in best-practices.rst

We should look at how other ABM libraries choose their trade-off here.
Could we create a new switch for the DataCollector, `multiple_agent_types`, for example, which is `False` by default but can be set to `True` if you want to collect data from multiple agent types? That way it's still contained to a single class, which can reduce the maintenance effort.
The maintenance effort is about the same with `DataCollectorRandomActivationByType`, since it will inherit from `DataCollector`, and only `_record_agents` needs to be replaced to handle agents by type. Adding an extra switch and if clause inside `DataCollector` doesn't scale, because what if there is a scheduler with 2-level type hierarchy (`Dict[str, Dict[str, Dict[int, Agent]]]`, see [Multi-level Mesa](https://github.com/tpike3/multilevel_mesa))? The switch and the if's keep growing.
That's a good point.

Would you be interested in creating an (draft) implementation?
My plate is rather full right now.
Thanks for the honest reply, let's discuss at the next development meeting how we can move this project forward.
What if you would restructure the [`DataCollector`](https://github.com/projectmesa/mesa/blob/af23b8bc854ec85235642a549ad1924bd19aa4a3/mesa/datacollection.py#L45) in such way that it's split in an agent-data and model-data part?

Maybe a `DataCollector` superclass with a `ModelDataCollector` and a `AgentDataCollector` subclasses? In that case the `ModelDataCollector` always expects a single model as input, and a `AgentDataCollector` a single agent type. If you want to collect data for more that one agent type, you simply call the `AgentDataCollector` multiple times.

Many of the code is now duplicate such as `_new_model_reporter` and `_new_agent_reporter`, so why not split those? And everything that's shared can sit in the `DataCollector` superclass.
I'm also running into this issue with respect to the agent attributes. After reading through the code, I wonder if it would be possible to just pass a specific `schedule` to the `DataCollector` constructor instead of assuming ([hard-coding](https://github.com/projectmesa/mesa/blob/659b00a0c547864f3a4359ac8f8e50359703d631/mesa/datacollection.py#L168)) that there exists an object in the model called `schedule`. That would seem to limit the bloat of `DataCollector` and remove one of the implicit assumptions in the class.

So something like this in the `Model`:
```python
self.datacollector_my_agents = mesa.DataCollector(
    schedule = self.schedule_my_agents, # arbitrary schedule with an .agents attribute that returns a list of agents
    agent_reporters = {...},
    model_reporters = {...}
)

self.datacollector_your_agents = mesa.DataCollector(
    schedule = self.schedule_your_agents, # arbitrary schedule with an .agents attribute that returns a list of agents
    agent_reporters = {...},
    model_reporters = {...}
)
```

Then to get the data:
```python
my_agent_data = model.datacollector_my_agents.get_agent_vars_dataframe()
your_agent_data = model.datacollector_your_agents.get_agent_vars_dataframe()
```
> Maybe a DataCollector superclass with a ModelDataCollector and a AgentDataCollector subclasses?
> ... 
> If you want to collect data for more that one agent type, you simply call the AgentDataCollector multiple times.

This definitely could provide ingredients for people who want to build their own data collector class. But I don't see it to be cleaner than just inheriting from the current `DataCollector` class and modifying the method you want to change as needed.

I think for now, for trading Sugarscape use case, we should implement the solution within the examples folder. Then wait for a while for it to become a part of the library.
@jacob-thrackle your solution could work. Though we also want to make sure the >80% of the use case stays simple and convenient, where the scheduler is assigned to the model attribute `schedule` anyway. So that users don't have to specify the scheduler object every time.
Absolutely. I started hacking around with it last night but the following section makes it difficult because of the string type for `model.schedule.steps` in `prefix` (line 158):
https://github.com/projectmesa/mesa/blob/373637ed28566f0c35832bfd5b51d44eae8739ab/mesa/datacollection.py#L156-L166

If I'm understanding the code correctly, the condition at line 157 is only `True` if all the `rep_funcs` are `partial` which is only true if the `reporter` is a `str`:
https://github.com/projectmesa/mesa/blob/373637ed28566f0c35832bfd5b51d44eae8739ab/mesa/datacollection.py#L137-L141

So having a reporter be a lambda or some other arbitrary function would make this not true and evaluate the `else` statement at line 161. I assume that it is done this way so that `get_reports = attrgetter(*prefix + attributes)` is really fast if the reporters are just attributes of the agents but this gain is nullified once you use anything other than a string in the reporter dictionary in the `DataCollector` itself. Since the `DataCollector.collect` method is typically only called once per model-step, and unless there are a huge number of agents in which one is only interested in getting attributes, it may be worthwhile to consider removing the condition and expression in lines 157-160 and just using the function at 163. 

Using the function has the added advantage of calling `self._getattr()` which has a default value of `None`, meaning that attributes which do not exist return `None` instead of throwing an `AttributeError`. The code in lines 158-160 never call `self._getattr()` and just use the `attribute_name` which is passed to `attrgetter()`.

This is all to say that I was able to make my suggestion work by simply removing lines 157-161, only using the function `get_reports` at line 163 with some minor adjustments, including a condition for backwards compatibility where the scheduler is assigned to the model attribute `schedule`:
```python
if self.schedule:
    schedule = self.schedule # new behavior
else:
    schedule = model.schedule # old behavior, self.schedule defaults to None in the constructor
```
This conditional has to be placed, unfortunately, in both `_record_agents` and `collect`, although one could limit it to one call and just pass in the schedule to `_record_agents` which is called from inside `collect`. Thoughts?


That L157-161 seems to be written by @Corvince in https://github.com/projectmesa/mesa/commit/aaa4c545fbaaeb7408c74148e96510421ea243fa. The original commit contains this comment for the `attrgetter`
```python
# Fast path if all reporters are attribute collecters
```
but had since been removed in the subsequent iterations.

But the original PR that did the specialization of the code when all the reporters are string-based, is #576. #575 contains some benchmark results
Thanks for digging that up @rht 

And yes, @jacob-thrackle you understood exactly right how it works. 
Basically, attrgetter gives the best possible performance, but only works if we have only string collectors.
This is a pure performance optimization. But, as you figured out, does not allow a default value.

Is is probably better to give up a tiny bit of performance by using getattr instead of attrgetter to have default values 
@Corvince Cool, I opened a PR from my personal account (this is @jacob-thrackle). Let me know when you guys get around to reviewing it and what changes, if any, you'd like me to make. I forgot to link it in the commit but this also fixes #976.

I made another version that kept the optimization but it only happened if two things were true: that `self.schedule == None` and if all the reporters in the dictionary are strings. I'm honestly not sure how often people use that (a lot?) but I can throw it back in if it's wanted. I removed it since the behavior would have been different if one had passed in `model.schedule` versus letting `self.schedule` be None. Felt inconsistent.

Resolving patch:
diff --git a/mesa/datacollection.py b/mesa/datacollection.py
--- a/mesa/datacollection.py
+++ b/mesa/datacollection.py
@@ -51,7 +51,13 @@ class DataCollector:
     one and stores the results.
     """
 
-    def __init__(self, model_reporters=None, agent_reporters=None, tables=None):
+    def __init__(
+        self,
+        model_reporters=None,
+        agent_reporters=None,
+        tables=None,
+        exclude_none_values=False,
+    ):
         """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.
@@ -74,6 +80,8 @@ 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.
+            exclude_none_values: Boolean of whether to drop records which values
+            are None, in the final result.
 
         Notes:
             If you want to pickle your model you must not use lambda functions.
@@ -97,6 +105,7 @@ class attributes of a model
         self.model_vars = {}
         self._agent_records = {}
         self.tables = {}
+        self.exclude_none_values = exclude_none_values
 
         if model_reporters is not None:
             for name, reporter in model_reporters.items():
@@ -151,7 +160,23 @@ def _new_table(self, table_name, table_columns):
     def _record_agents(self, model):
         """Record agents data in a mapping of functions and agents."""
         rep_funcs = self.agent_reporters.values()
+        if self.exclude_none_values:
+            # Drop records which values are None.
+
+            def get_reports(agent):
+                _prefix = (agent.model.schedule.steps, agent.unique_id)
+                reports = (rep(agent) for rep in rep_funcs)
+                reports_without_none = tuple(r for r in reports if r is not None)
+                if len(reports_without_none) == 0:
+                    return None
+                return _prefix + reports_without_none
+
+            agent_records = (get_reports(agent) for agent in model.schedule.agents)
+            agent_records_without_none = (r for r in agent_records if r is not None)
+            return agent_records_without_none
+
         if all(hasattr(rep, "attribute_name") for rep in rep_funcs):
+            # This branch is for performance optimization purpose.
             prefix = ["model.schedule.steps", "unique_id"]
             attributes = [func.attribute_name for func in rep_funcs]
             get_reports = attrgetter(*prefix + attributes)
diff --git a/mesa/model.py b/mesa/model.py
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -66,7 +66,11 @@ def reset_randomizer(self, seed: int | None = None) -> None:
         self._seed = seed
 
     def initialize_data_collector(
-        self, model_reporters=None, agent_reporters=None, tables=None
+        self,
+        model_reporters=None,
+        agent_reporters=None,
+        tables=None,
+        exclude_none_values=False,
     ) -> None:
         if not hasattr(self, "schedule") or self.schedule is None:
             raise RuntimeError(
@@ -80,6 +84,7 @@ def initialize_data_collector(
             model_reporters=model_reporters,
             agent_reporters=agent_reporters,
             tables=tables,
+            exclude_none_values=exclude_none_values,
         )
         # Collect data for the first time during initialization.
         self.datacollector.collect(self)


Test patch:
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -32,6 +32,14 @@ def write_final_values(self):
         self.model.datacollector.add_table_row("Final_Values", row)
 
 
+class DifferentMockAgent(MockAgent):
+    # We define a different MockAgent to test for attributes that are present
+    # only in 1 type of agent, but not the other.
+    def __init__(self, unique_id, model, val=0):
+        super().__init__(unique_id, model, val=val)
+        self.val3 = val + 42
+
+
 class MockModel(Model):
     """
     Minimalistic model for testing purposes.
@@ -39,13 +47,20 @@ class MockModel(Model):
 
     schedule = BaseScheduler(None)
 
-    def __init__(self):
+    def __init__(self, test_exclude_none_values=False):
         self.schedule = BaseScheduler(self)
         self.model_val = 100
 
-        for i in range(10):
-            a = MockAgent(i, self, val=i)
-            self.schedule.add(a)
+        self.n = 10
+        for i in range(self.n):
+            self.schedule.add(MockAgent(i, self, val=i))
+            if test_exclude_none_values:
+                self.schedule.add(DifferentMockAgent(self.n + i, self, val=i))
+        if test_exclude_none_values:
+            # Only DifferentMockAgent has val3.
+            agent_reporters = {"value": lambda a: a.val, "value3": "val3"}
+        else:
+            agent_reporters = {"value": lambda a: a.val, "value2": "val2"}
         self.initialize_data_collector(
             {
                 "total_agents": lambda m: m.schedule.get_agent_count(),
@@ -54,8 +69,9 @@ def __init__(self):
                 "model_calc_comp": [self.test_model_calc_comp, [3, 4]],
                 "model_calc_fail": [self.test_model_calc_comp, [12, 0]],
             },
-            {"value": lambda a: a.val, "value2": "val2"},
+            agent_reporters,
             {"Final_Values": ["agent_id", "final_value"]},
+            exclude_none_values=test_exclude_none_values,
         )
 
     def test_model_calc_comp(self, input1, input2):
@@ -195,5 +211,44 @@ def test_initialize_before_agents_added_to_scheduler(self):
         )
 
 
+class TestDataCollectorExcludeNone(unittest.TestCase):
+    def setUp(self):
+        """
+        Create the model and run it a set number of steps.
+        """
+        self.model = MockModel(test_exclude_none_values=True)
+        for i in range(7):
+            if i == 4:
+                self.model.schedule.remove(self.model.schedule._agents[3])
+            self.model.step()
+
+    def test_agent_records(self):
+        """
+        Test agent-level variable collection.
+        """
+        data_collector = self.model.datacollector
+        agent_table = data_collector.get_agent_vars_dataframe()
+
+        assert len(data_collector._agent_records) == 8
+        for step, records in data_collector._agent_records.items():
+            if step < 5:
+                assert len(records) == 20
+            else:
+                assert len(records) == 19
+
+            for values in records:
+                agent_id = values[1]
+                if agent_id < self.model.n:
+                    assert len(values) == 3
+                else:
+                    # Agents with agent_id >= self.model.n are
+                    # DifferentMockAgent, which additionally contains val3.
+                    assert len(values) == 4
+
+        assert "value" in list(agent_table.columns)
+        assert "value2" not in list(agent_table.columns)
+        assert "value3" in list(agent_table.columns)
+
+
 if __name__ == "__main__":
     unittest.main()