Repository: projectmesa/mesa
Pull request #1838 (2023-10-18 14:46:15)

DataCollector: Allow agent reporters to take class methods and functions with parameter lists
Modify the `DataCollector` class to allow agent reporters to take methods of a class/instance and functions with parameters placed in a list (like model reporters), by extending the `_new_agent_reporter` method.

This implementation starts by checking if the reporter is an attribute string. If so, it creates a function to retrieve the attribute from an agent. Next, it checks if the reporter is a list. If it is, this indicates that we have a function with parameters, so it wraps that function to pass those parameters when called.

For any other type (like lambdas or methods), we assume they're directly suitable to be used as reporters.

Now, with this modification, agent reporters in the `DataCollector` class can take:

1. Attribute strings
2. Function objects (like lambdas)
3. Methods of a class/instance
4. Functions with parameters placed in a list

This approach ensures backward compatibility because the existing checks for attribute strings and function objects remain unchanged. The added functionality only extends the capabilities of the class without altering the existing behavior.

Closes: #1837

**Todo**
 - [x] Update and fix tests
 - [x] Add more docstring

Can be done in separate PRs:
 - Update examples
 - Update tutorial (optional)

EwoutH: Code is ready to review. Will still work a bit on documentation.

> Is this a syntax sugar for an agent reporter of previously specified as lambda agent: func(agent, *params), but now with [func, params]?

Yes, that's completely right, the modification is essentially syntactic sugar, offering an alternative way to specify agent reporters that involve functions with additional parameters.

The big advantage of this however, is that it's now completely consistent with specifying model reporters. This way, if you can use the syntax for a model reporter, you can use it for an agent reporter, and visa versa.

A disadvantage could of course be the increased complexity of the code base. But, then I would argue, we should also report one of those syntax types from the model reporter. Because than that would be consistent, but I also don't think that's a desired option due to breaking backwards compatibility.
EwoutH: Updated the docs, this PR is ready for review.

I might update the tutorial and some examples in separate PRs.
rht: I personally think that passing `lambda agent: func(agent, *params)` doesn't require you to read extra documentation and is most likely what people would have come up with naturally.

I think it has nothing to do with the pickling issue in `multiprocessing` as I initially guessed in https://github.com/projectmesa/mesa/issues/1837#issuecomment-1769869159, because it will create a local anonymous function that is unpickleable anyway.

@jackiekazil @tpike3 should we merge this PR?
rht: Content: LGTM. Concept: I prefer passing `lambda agent: func(agent, *params)`, but others might think otherwise.
jackiekazil: I am good with this. @rht thank you for doing the review. 
LGTM. 

Patch:
diff --git a/mesa/datacollection.py b/mesa/datacollection.py
index 7f3dc847..8bddfa23 100644
--- a/mesa/datacollection.py
+++ b/mesa/datacollection.py
@@ -56,44 +56,49 @@ class DataCollector:
         agent_reporters=None,
         tables=None,
     ):
-        """Instantiate a DataCollector with lists of model and agent reporters.
+        """
+        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.
-        For example, if there was only one model-level reporter for number of
-        agents, it might look like:
-            {"agent_count": lambda m: m.schedule.get_agent_count() }
-        If there was only one agent-level reporter (e.g. the agent's energy),
-        it might look like this:
-            {"energy": "energy"}
-        or like this:
-            {"energy": lambda a: a.energy}
+        variable name to either an attribute name, a function, a method of a class/instance,
+        or a function with parameters placed in a list.
+
+        Model reporters can take four types of arguments:
+        1. Lambda function:
+           {"agent_count": lambda m: m.schedule.get_agent_count()}
+        2. Method of a class/instance:
+           {"agent_count": self.get_agent_count} # self here is a class instance
+           {"agent_count": Model.get_agent_count} # Model here is a class
+        3. Class attributes of a model:
+           {"model_attribute": "model_attribute"}
+        4. Functions with parameters that have been placed in a list:
+           {"Model_Function": [function, [param_1, param_2]]}
+
+        Agent reporters can similarly take:
+        1. Attribute name (string) referring to agent's attribute:
+           {"energy": "energy"}
+        2. Lambda function:
+           {"energy": lambda a: a.energy}
+        3. Method of an agent class/instance:
+           {"agent_action": self.do_action} # self here is an agent class instance
+           {"agent_action": Agent.do_action} # Agent here is a class
+        4. Functions with parameters placed in a list:
+           {"Agent_Function": [function, [param_1, param_2]]}
 
         The tables arg accepts a dictionary mapping names of tables to lists of
         columns. For example, if we want to allow agents to write their age
         when they are destroyed (to keep track of lifespans), it might look
         like:
-            {"Lifespan": ["unique_id", "age"]}
+           {"Lifespan": ["unique_id", "age"]}
 
         Args:
-            model_reporters: Dictionary of reporter names and attributes/funcs
-            agent_reporters: Dictionary of reporter names and attributes/funcs.
+            model_reporters: Dictionary of reporter names and attributes/funcs/methods.
+            agent_reporters: Dictionary of reporter names and attributes/funcs/methods.
             tables: Dictionary of table names to lists of column names.
 
         Notes:
-            If you want to pickle your model you must not use lambda functions.
-            If your model includes a large number of agents, you should *only*
-            use attribute names for the agent reporter, it will be much faster.
-
-            Model reporters can take four types of arguments:
-            lambda like above:
-            {"agent_count": lambda m: m.schedule.get_agent_count() }
-            method of a class/instance:
-            {"agent_count": self.get_agent_count} # self here is a class instance
-            {"agent_count": Model.get_agent_count} # Model here is a class
-            class attributes of a model
-            {"model_attribute": "model_attribute"}
-            functions with parameters that have placed in a list
-            {"Model_Function":[function, [param_1, param_2]]}
+            - If you want to pickle your model you must not use lambda functions.
+            - If your model includes a large number of agents, it is recommended to
+              use attribute names for the agent reporter, as it will be faster.
         """
         self.model_reporters = {}
         self.agent_reporters = {}
@@ -130,16 +135,31 @@ class DataCollector:
 
         Args:
             name: Name of the agent-level variable to collect.
-            reporter: Attribute string, or function object that returns the
-                      variable when given a model instance.
+            reporter: Attribute string, function object, method of a class/instance, or
+                      function with parameters placed in a list that returns the
+                      variable when given an agent instance.
         """
+        # Check if the reporter is an attribute string
         if isinstance(reporter, str):
             attribute_name = reporter
 
-            def reporter(agent):
+            def attr_reporter(agent):
                 return getattr(agent, attribute_name, None)
 
-            reporter.attribute_name = attribute_name
+            reporter = attr_reporter
+
+        # Check if the reporter is a function with arguments placed in a list
+        elif isinstance(reporter, list):
+            func, params = reporter[0], reporter[1]
+
+            def func_with_params(agent):
+                return func(agent, *params)
+
+            reporter = func_with_params
+
+        # For other types (like lambda functions, method of a class/instance),
+        # it's already suitable to be used as a reporter directly.
+
         self.agent_reporters[name] = reporter
 
     def _new_table(self, table_name, table_columns):


Test patch:
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index 8cdee440..a45d8891 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -24,6 +24,9 @@ class MockAgent(Agent):
         self.val += 1
         self.val2 += 1
 
+    def double_val(self):
+        return self.val * 2
+
     def write_final_values(self):
         """
         Write the final value to the appropriate table.
@@ -32,6 +35,10 @@ class MockAgent(Agent):
         self.model.datacollector.add_table_row("Final_Values", row)
 
 
+def agent_function_with_params(agent, multiplier, offset):
+    return (agent.val * multiplier) + offset
+
+
 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.
@@ -54,17 +61,21 @@ class MockModel(Model):
         self.n = 10
         for i in range(self.n):
             self.schedule.add(MockAgent(i, self, val=i))
-        agent_reporters = {"value": lambda a: a.val, "value2": "val2"}
         self.initialize_data_collector(
-            {
+            model_reporters={
                 "total_agents": lambda m: m.schedule.get_agent_count(),
                 "model_value": "model_val",
                 "model_calc": self.schedule.get_agent_count,
                 "model_calc_comp": [self.test_model_calc_comp, [3, 4]],
                 "model_calc_fail": [self.test_model_calc_comp, [12, 0]],
             },
-            agent_reporters,
-            {"Final_Values": ["agent_id", "final_value"]},
+            agent_reporters={
+                "value": lambda a: a.val,
+                "value2": "val2",
+                "double_value": MockAgent.double_val,
+                "value_with_params": [agent_function_with_params, [2, 3]],
+            },
+            tables={"Final_Values": ["agent_id", "final_value"]},
         )
 
     def test_model_calc_comp(self, input1, input2):
@@ -132,6 +143,19 @@ class TestDataCollector(unittest.TestCase):
         data_collector = self.model.datacollector
         agent_table = data_collector.get_agent_vars_dataframe()
 
+        assert "double_value" in list(agent_table.columns)
+        assert "value_with_params" in list(agent_table.columns)
+
+        # Check the double_value column
+        for (step, agent_id), value in agent_table["double_value"].items():
+            expected_value = (step + agent_id) * 2
+            self.assertEqual(value, expected_value)
+
+        # Check the value_with_params column
+        for (step, agent_id), value in agent_table["value_with_params"].items():
+            expected_value = ((step + agent_id) * 2) + 3
+            self.assertEqual(value, expected_value)
+
         assert len(data_collector._agent_records) == 8
         for step, records in data_collector._agent_records.items():
             if step < 5:
@@ -140,7 +164,7 @@ class TestDataCollector(unittest.TestCase):
                 assert len(records) == 9
 
             for values in records:
-                assert len(values) == 4
+                assert len(values) == 6
 
         assert "value" in list(agent_table.columns)
         assert "value2" in list(agent_table.columns)
@@ -175,7 +199,7 @@ class TestDataCollector(unittest.TestCase):
         agent_vars = data_collector.get_agent_vars_dataframe()
         table_df = data_collector.get_table_dataframe("Final_Values")
         assert model_vars.shape == (8, 5)
-        assert agent_vars.shape == (77, 2)
+        assert agent_vars.shape == (77, 4)
         assert table_df.shape == (9, 2)
 
         with self.assertRaises(Exception):