Repository: projectmesa/mesa
Pull request #2283 (2024-09-06 10:47:41)

Add AgentSet.shuffle_do()
This PR introduces a new `shuffle_do` method to the `AgentSet` class, optimizing the process of shuffling agents and applying a method to them in a single operation.

### Motive
The current approach of shuffling agents and then applying a method (`shuffle().do()`) is inefficient, especially for large agent sets. This new method aims to significantly improve performance by combining these operations, reducing memory allocations and iterations.

### Implementation
- Added a new `shuffle_do` method to the `AgentSet` class in `mesa/agent.py`.
- The method takes a `method` parameter (either a string or callable) and additional args/kwargs.
- It shuffles the agents in-place.
- The specified method is then applied to each agent in the shuffled order.
- Added corresponding unit tests in `tests/test_agent.py`.
- Updated the benchmarks (BoltzmannWealth and Wolf-sheep).

### Usage Examples
```python
# Using a string method name
model.agents.shuffle_do("step")

# Using a lambda function
model.agents.shuffle_do(lambda agent: agent.move())
```

Before/After Performance Comparison:
| Configuration                | shuffle().do() | shuffle(inplace=True).do() | shuffle_do() |
|------------------------------|----------------|----------------------------|--------------|
| 10,000 agents, 1,000 steps   | 8.27 s         | 3.65 s                     | 3.06 s       |
| 100,000 agents, 100 steps    | 13.71 s        | 6.31 s                     | 3.71 s       |
| 1,000,000 agents, 10 steps   | 18.36 s        | 9.44 s                     | 5.75 s       |

As shown, `shuffle_do()` provides significant performance improvements, especially for larger agent sets.

### Additional Notes
- This change is backwards compatible and doesn't affect existing code using `shuffle().do()`.
- The performance gain is more pronounced for larger agent sets and more frequent shuffling operations.
- Examples and docs can be updated (other PR).

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]       | Run time [95% CI]        |
|-----------------|--------|--------------------------|--------------------------|
| BoltzmannWealth | small  | 🔵 -2.1% [-2.6%, -1.6%]  | 🔵 -0.5% [-0.6%, -0.4%]  |
| BoltzmannWealth | large  | 🔵 -9.1% [-23.6%, -1.3%] | 🔵 -2.7% [-7.5%, +3.2%]  |
| Schelling       | small  | 🔵 -2.0% [-2.3%, -1.6%]  | 🔵 -3.2% [-3.8%, -2.8%]  |
| Schelling       | large  | 🔵 -2.2% [-3.1%, -1.1%]  | 🟢 -9.1% [-12.0%, -6.2%] |
| WolfSheep       | small  | 🔵 -0.7% [-2.0%, +0.6%]  | 🔵 -1.3% [-1.7%, -0.8%]  |
| WolfSheep       | large  | 🔵 -1.9% [-2.8%, -1.0%]  | 🔵 -4.3% [-6.9%, -1.7%]  |
| BoidFlockers    | small  | 🔵 +2.2% [+1.8%, +2.6%]  | 🔵 +0.2% [-0.4%, +0.8%]  |
| BoidFlockers    | large  | 🔵 +2.6% [+2.1%, +3.2%]  | 🔵 -0.4% [-1.0%, +0.3%]  |

EwoutH: <details>
<summary>Small benchmark:</summary>

```Python
from mesa import Agent, Model

class MyAgent(Agent):
    def __init__(self, model):
        super().__init__(model)

    def step(self):
        pass

class MyModel1(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle().do("step")


class MyModel2(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle(inplace=True).do("step")


class MyModel3(Model):
    def __init__(self):
        super().__init__()
        # Create 10000
        for _ in range(10000):
            MyAgent(self)

    def step(self):
        self.agents.shuffle_do("step")


# Benchmark the three models
from time import time

models = [MyModel1(), MyModel2(), MyModel3()]
times = []
print("Running models...")
for model in models:
    print(f"Running {model}")
    start = time()
    for _ in range(1000):
        model.step()
    times.append(time() - start)

print(f"Results: {times}")
```
</details>

| Configuration                          | shuffle().do() | shuffle(inplace=True).do() | shuffle_do() |
|----------------------------------|----------------|----------------------------|--------------|
| **10,000 agents, 1,000 steps**   | 8.27 s         | 3.65 s                     | 3.06 s       |
| **100,000 agents, 100 steps**    | 13.71 s        | 6.31 s                     | 3.71 s       |
| **1,000,000 agents, 10 steps**   | 18.36 s        | 9.44 s                     | 5.75 s       |

EwoutH: A bit more radical approach: #2284
EwoutH: I was thinking of adding a `shuffle` keyword argument to some functions, starting with `do()`. So you can do:

```Python
model.agents.do("step", shuffle=True)
```

I think this is explicit and allows us to also apply this to other performance critical functions.

What does everybody think about this as an API? It should have the same performance as `shuffle_do`. 
quaquel: I am indifferent between the two. So pick which one you prefer and get this done. The performance benefits are clear and shuffling followed by do is probably the most common sequence of operations done with an agentset. This last point might be a reason to have it as a separate method rather than via a keyword argument. 

Patch:
diff --git a/benchmarks/BoltzmannWealth/boltzmann_wealth.py b/benchmarks/BoltzmannWealth/boltzmann_wealth.py
index 041445f1..93d4da14 100644
--- a/benchmarks/BoltzmannWealth/boltzmann_wealth.py
+++ b/benchmarks/BoltzmannWealth/boltzmann_wealth.py
@@ -60,7 +60,7 @@ class BoltzmannWealth(mesa.Model):
 
     def step(self):
         """Run the model for a single step."""
-        self.agents.shuffle().do("step")
+        self.agents.shuffle_do("step")
         # collect data
         self.datacollector.collect(self)
 
diff --git a/benchmarks/WolfSheep/wolf_sheep.py b/benchmarks/WolfSheep/wolf_sheep.py
index 0999fc77..16c01e00 100644
--- a/benchmarks/WolfSheep/wolf_sheep.py
+++ b/benchmarks/WolfSheep/wolf_sheep.py
@@ -227,8 +227,8 @@ class WolfSheep(Model):
 
     def step(self):
         """Run one step of the model."""
-        self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
-        self.agents_by_type[Wolf].shuffle(inplace=True).do("step")
+        self.agents_by_type[Sheep].shuffle_do("step")
+        self.agents_by_type[Wolf].shuffle_do("step")
 
 
 if __name__ == "__main__":
diff --git a/mesa/agent.py b/mesa/agent.py
index 2d098c35..ec98efaa 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -293,6 +293,23 @@ class AgentSet(MutableSet, Sequence):
 
         return self
 
+    def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
+        """Shuffle the agents in the AgentSet and then invoke a method or function on each agent.
+
+        It's a fast, optimized version of calling shuffle() followed by do().
+        """
+        agents = list(self._agents.keys())
+        self.random.shuffle(agents)
+
+        if isinstance(method, str):
+            for agent in agents:
+                getattr(agent, method)(*args, **kwargs)
+        else:
+            for agent in agents:
+                method(agent, *args, **kwargs)
+
+        return self
+
     def map(self, method: str | Callable, *args, **kwargs) -> list[Any]:
         """Invoke a method or function on each agent in the AgentSet and return the results.
 
diff --git a/mesa/datacollection.py b/mesa/datacollection.py
index bf50be2a..2fb4d241 100644
--- a/mesa/datacollection.py
+++ b/mesa/datacollection.py
@@ -1,18 +1,19 @@
 """Mesa Data Collection Module.
 
 DataCollector is meant to provide a simple, standard way to collect data
-generated by a Mesa model. It collects three types of data: model-level data,
-agent-level data, and tables.
+generated by a Mesa model. It collects four types of data: model-level data,
+agent-level data, agent-type-level data, and tables.
 
-A DataCollector is instantiated with two dictionaries of reporter names and
-associated variable names or functions for each, one for model-level data and
-one for agent-level data; a third dictionary provides table names and columns.
-Variable names are converted into functions which retrieve attributes of that
-name.
+A DataCollector is instantiated with three dictionaries of reporter names and
+associated variable names or functions for each, one for model-level data,
+one for agent-level data, and one for agent-type-level data; a fourth dictionary
+provides table names and columns. Variable names are converted into functions
+which retrieve attributes of that name.
 
 When the collect() method is called, each model-level function is called, with
 the model as the argument, and the results associated with the relevant
-variable. Then the agent-level functions are called on each agent.
+variable. Then the agent-level functions are called on each agent, and the
+agent-type-level functions are called on each agent of the specified type.
 
 Additionally, other objects can write directly to tables by passing in an
 appropriate dictionary object for a table row.
@@ -21,19 +22,18 @@ The DataCollector then stores the data it collects in dictionaries:
     * model_vars maps each reporter to a list of its values
     * tables maps each table to a dictionary, with each column as a key with a
       list as its value.
-    * _agent_records maps each model step to a list of each agents id
+    * _agent_records maps each model step to a list of each agent's id
       and its values.
+    * _agenttype_records maps each model step to a dictionary of agent types,
+      each containing a list of each agent's id and its values.
 
 Finally, DataCollector can create a pandas DataFrame from each collection.
-
-The default DataCollector here makes several assumptions:
-    * The model has an agent list called agents
-    * For collecting agent-level variables, agents must have a unique_id
 """
 
 import contextlib
 import itertools
 import types
+import warnings
 from copy import deepcopy
 from functools import partial
 
@@ -44,24 +44,25 @@ with contextlib.suppress(ImportError):
 class DataCollector:
     """Class for collecting data generated by a Mesa model.
 
-    A DataCollector is instantiated with dictionaries of names of model- and
-    agent-level variables to collect, associated with attribute names or
-    functions which actually collect them. When the collect(...) method is
-    called, it collects these attributes and executes these functions one by
-    one and stores the results.
+    A DataCollector is instantiated with dictionaries of names of model-,
+    agent-, and agent-type-level variables to collect, associated with
+    attribute names or functions which actually collect them. When the
+    collect(...) method is called, it collects these attributes and executes
+    these functions one by one and stores the results.
     """
 
     def __init__(
         self,
         model_reporters=None,
         agent_reporters=None,
+        agenttype_reporters=None,
         tables=None,
     ):
-        """Instantiate a DataCollector with lists of model and agent reporters.
+        """Instantiate a DataCollector with lists of model, agent, and agent-type reporters.
 
-        Both model_reporters and agent_reporters accept a dictionary mapping a
-        variable name to either an attribute name, a function, a method of a class/instance,
-        or a function with parameters placed in a list.
+        Both model_reporters, agent_reporters, and agenttype_reporters accept a
+        dictionary mapping a 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:
@@ -85,6 +86,10 @@ class DataCollector:
         4. Functions with parameters placed in a list:
            {"Agent_Function": [function, [param_1, param_2]]}
 
+        Agenttype reporters take a dictionary mapping agent types to dictionaries
+        of reporter names and attributes/funcs/methods, similar to agent_reporters:
+           {Wolf: {"energy": lambda a: a.energy}}
+
         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
@@ -94,6 +99,8 @@ class DataCollector:
         Args:
             model_reporters: Dictionary of reporter names and attributes/funcs/methods.
             agent_reporters: Dictionary of reporter names and attributes/funcs/methods.
+            agenttype_reporters: Dictionary of agent types to dictionaries of
+                                 reporter names and attributes/funcs/methods.
             tables: Dictionary of table names to lists of column names.
 
         Notes:
@@ -103,9 +110,11 @@ class DataCollector:
         """
         self.model_reporters = {}
         self.agent_reporters = {}
+        self.agenttype_reporters = {}
 
         self.model_vars = {}
         self._agent_records = {}
+        self._agenttype_records = {}
         self.tables = {}
 
         if model_reporters is not None:
@@ -116,6 +125,11 @@ class DataCollector:
             for name, reporter in agent_reporters.items():
                 self._new_agent_reporter(name, reporter)
 
+        if agenttype_reporters is not None:
+            for agent_type, reporters in agenttype_reporters.items():
+                for name, reporter in reporters.items():
+                    self._new_agenttype_reporter(agent_type, name, reporter)
+
         if tables is not None:
             for name, columns in tables.items():
                 self._new_table(name, columns)
@@ -163,6 +177,38 @@ class DataCollector:
 
         self.agent_reporters[name] = reporter
 
+    def _new_agenttype_reporter(self, agent_type, name, reporter):
+        """Add a new agent-type-level reporter to collect.
+
+        Args:
+            agent_type: The type of agent to collect data for.
+            name: Name of the agent-type-level variable to collect.
+            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.
+        """
+        if agent_type not in self.agenttype_reporters:
+            self.agenttype_reporters[agent_type] = {}
+
+        # Use the same logic as _new_agent_reporter
+        if isinstance(reporter, str):
+            attribute_name = reporter
+
+            def attr_reporter(agent):
+                return getattr(agent, attribute_name, None)
+
+            reporter = attr_reporter
+
+        elif isinstance(reporter, list):
+            func, params = reporter[0], reporter[1]
+
+            def func_with_params(agent):
+                return func(agent, *params)
+
+            reporter = func_with_params
+
+        self.agenttype_reporters[agent_type][name] = reporter
+
     def _new_table(self, table_name, table_columns):
         """Add a new table that objects can write to.
 
@@ -190,6 +236,34 @@ class DataCollector:
         )
         return agent_records
 
+    def _record_agenttype(self, model, agent_type):
+        """Record agent-type data in a mapping of functions and agents."""
+        rep_funcs = self.agenttype_reporters[agent_type].values()
+
+        def get_reports(agent):
+            _prefix = (agent.model.steps, agent.unique_id)
+            reports = tuple(rep(agent) for rep in rep_funcs)
+            return _prefix + reports
+
+        agent_types = model.agent_types
+        if agent_type in agent_types:
+            agents = model.agents_by_type[agent_type]
+        else:
+            from mesa import Agent
+
+            if issubclass(agent_type, Agent):
+                agents = [
+                    agent for agent in model.agents if isinstance(agent, agent_type)
+                ]
+            else:
+                # Raise error if agent_type is not in model.agent_types
+                raise ValueError(
+                    f"Agent type {agent_type} is not recognized as an Agent type in the model or Agent subclass. Use an Agent (sub)class, like {agent_types}."
+                )
+
+        agenttype_records = map(get_reports, agents)
+        return agenttype_records
+
     def collect(self, model):
         """Collect all the data for the given model object."""
         if self.model_reporters:
@@ -208,7 +282,6 @@ class DataCollector:
                 elif isinstance(reporter, list):
                     self.model_vars[var].append(deepcopy(reporter[0](*reporter[1])))
                 # Assume it's a callable otherwise (e.g., method)
-                # TODO: Check if method of a class explicitly
                 else:
                     self.model_vars[var].append(deepcopy(reporter()))
 
@@ -216,6 +289,14 @@ class DataCollector:
             agent_records = self._record_agents(model)
             self._agent_records[model.steps] = list(agent_records)
 
+        if self.agenttype_reporters:
+            self._agenttype_records[model.steps] = {}
+            for agent_type in self.agenttype_reporters:
+                agenttype_records = self._record_agenttype(model, agent_type)
+                self._agenttype_records[model.steps][agent_type] = list(
+                    agenttype_records
+                )
+
     def add_table_row(self, table_name, row, ignore_missing=False):
         """Add a row dictionary to a specific table.
 
@@ -272,6 +353,38 @@ class DataCollector:
         )
         return df
 
+    def get_agenttype_vars_dataframe(self, agent_type):
+        """Create a pandas DataFrame from the agent-type variables for a specific agent type.
+
+        The DataFrame has one column for each variable, with two additional
+        columns for tick and agent_id.
+
+        Args:
+            agent_type: The type of agent to get the data for.
+        """
+        # Check if self.agenttype_reporters dictionary is empty for this agent type, if so return empty DataFrame
+        if agent_type not in self.agenttype_reporters:
+            warnings.warn(
+                f"No agent-type reporters have been defined for {agent_type} in the DataCollector, returning empty DataFrame.",
+                UserWarning,
+                stacklevel=2,
+            )
+            return pd.DataFrame()
+
+        all_records = itertools.chain.from_iterable(
+            records[agent_type]
+            for records in self._agenttype_records.values()
+            if agent_type in records
+        )
+        rep_names = list(self.agenttype_reporters[agent_type])
+
+        df = pd.DataFrame.from_records(
+            data=all_records,
+            columns=["Step", "AgentID", *rep_names],
+            index=["Step", "AgentID"],
+        )
+        return df
+
     def get_table_dataframe(self, table_name):
         """Create a pandas DataFrame from a particular table.
 
diff --git a/mesa/experimental/devs/examples/epstein_civil_violence.py b/mesa/experimental/devs/examples/epstein_civil_violence.py
index ce6b835e..6f32e061 100644
--- a/mesa/experimental/devs/examples/epstein_civil_violence.py
+++ b/mesa/experimental/devs/examples/epstein_civil_violence.py
@@ -293,7 +293,7 @@ class EpsteinCivilViolence(Model):
 
     def step(self):
         """Run one step of the model."""
-        self.active_agents.shuffle(inplace=True).do("step")
+        self.active_agents.shuffle_do("step")
 
 
 if __name__ == "__main__":
diff --git a/mesa/experimental/devs/examples/wolf_sheep.py b/mesa/experimental/devs/examples/wolf_sheep.py
index b76a35bf..8d7d16d6 100644
--- a/mesa/experimental/devs/examples/wolf_sheep.py
+++ b/mesa/experimental/devs/examples/wolf_sheep.py
@@ -230,8 +230,8 @@ class WolfSheep(mesa.Model):
 
     def step(self):
         """Perform one step of the model."""
-        self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
-        self.agents_by_type[Wolf].shuffle(inplace=True).do("step")
+        self.agents_by_type[Sheep].shuffle_do("step")
+        self.agents_by_type[Wolf].shuffle_do("step")
 
 
 if __name__ == "__main__":
diff --git a/mesa/model.py b/mesa/model.py
index 1302c60e..a3227c10 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -198,6 +198,7 @@ class Model:
         self,
         model_reporters=None,
         agent_reporters=None,
+        agenttype_reporters=None,
         tables=None,
     ) -> None:
         """Initialize the data collector for the model.
@@ -205,6 +206,7 @@ class Model:
         Args:
             model_reporters: model reporters to collect
             agent_reporters: agent reporters to collect
+            agenttype_reporters: agent type reporters to collect
             tables: tables to collect
 
         """
@@ -219,6 +221,7 @@ class Model:
         self.datacollector = DataCollector(
             model_reporters=model_reporters,
             agent_reporters=agent_reporters,
+            agenttype_reporters=agenttype_reporters,
             tables=tables,
         )
         # Collect data for the first time during initialization.


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index f43a80c8..a0837de7 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -437,6 +437,42 @@ def test_agentset_map_callable():
     assert all(i == entry for i, entry in zip(results, range(1, 11)))
 
 
+def test_agentset_shuffle_do():
+    """Test AgentSet.shuffle_do method."""
+    model = Model()
+
+    class TestAgentShuffleDo(Agent):
+        def __init__(self, model):
+            super().__init__(model)
+            self.called = False
+
+        def test_method(self):
+            self.called = True
+
+    agents = [TestAgentShuffleDo(model) for _ in range(100)]
+    agentset = AgentSet(agents, model)
+
+    # Test shuffle_do with a string method name
+    agentset.shuffle_do("test_method")
+    assert all(agent.called for agent in agents)
+
+    # Reset the called flag
+    for agent in agents:
+        agent.called = False
+
+    # Test shuffle_do with a callable
+    agentset.shuffle_do(lambda agent: setattr(agent, "called", True))
+    assert all(agent.called for agent in agents)
+
+    # Verify that the order is indeed shuffled
+    original_order = list(agentset)
+    shuffled_order = []
+    agentset.shuffle_do(lambda agent: shuffled_order.append(agent))
+    assert (
+        original_order != shuffled_order
+    ), "The order should be different after shuffle_do"
+
+
 def test_agentset_get_attribute():
     """Test AgentSet.get for attributes."""
     model = Model()
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index 88792554..b2760fc1 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -3,6 +3,7 @@
 import unittest
 
 from mesa import Agent, Model
+from mesa.datacollection import DataCollector
 from mesa.time import BaseScheduler
 
 
@@ -28,6 +29,30 @@ class MockAgent(Agent):
         self.model.datacollector.add_table_row("Final_Values", row)
 
 
+class MockAgentA(MockAgent):
+    """Agent subclass A for testing agent-type-specific reporters."""
+
+    def __init__(self, model, val=0):  # noqa: D107
+        super().__init__(model, val)
+        self.type_a_val = val * 2
+
+    def step(self):  # noqa: D102
+        super().step()
+        self.type_a_val = self.val * 2
+
+
+class MockAgentB(MockAgent):
+    """Agent subclass B for testing agent-type-specific reporters."""
+
+    def __init__(self, model, val=0):  # noqa: D107
+        super().__init__(model, val)
+        self.type_b_val = val * 3
+
+    def step(self):  # noqa: D102
+        super().step()
+        self.type_b_val = self.val * 3
+
+
 def agent_function_with_params(agent, multiplier, offset):  # noqa: D103
     return (agent.val * multiplier) + offset
 
@@ -74,6 +99,33 @@ class MockModel(Model):
         self.datacollector.collect(self)
 
 
+class MockModelWithAgentTypes(Model):
+    """Model for testing agent-type-specific reporters."""
+
+    def __init__(self):  # noqa: D107
+        super().__init__()
+        self.model_val = 100
+
+        for i in range(10):
+            if i % 2 == 0:
+                MockAgentA(self, val=i)
+            else:
+                MockAgentB(self, val=i)
+
+        self.datacollector = DataCollector(
+            model_reporters={"total_agents": lambda m: len(m.agents)},
+            agent_reporters={"value": lambda a: a.val},
+            agenttype_reporters={
+                MockAgentA: {"type_a_val": lambda a: a.type_a_val},
+                MockAgentB: {"type_b_val": lambda a: a.type_b_val},
+            },
+        )
+
+    def step(self):  # noqa: D102
+        self.agents.do("step")
+        self.datacollector.collect(self)
+
+
 class TestDataCollector(unittest.TestCase):
     """Tests for DataCollector."""
 
@@ -206,5 +258,130 @@ class TestDataCollectorInitialization(unittest.TestCase):
         )
 
 
+class TestDataCollectorWithAgentTypes(unittest.TestCase):
+    """Tests for DataCollector with agent-type-specific reporters."""
+
+    def setUp(self):
+        """Create the model and run it a set number of steps."""
+        self.model = MockModelWithAgentTypes()
+        for _ in range(5):
+            self.model.step()
+
+    def test_agenttype_vars(self):
+        """Test agent-type-specific variable collection."""
+        data_collector = self.model.datacollector
+
+        # Test MockAgentA data
+        agent_a_data = data_collector.get_agenttype_vars_dataframe(MockAgentA)
+        self.assertIn("type_a_val", agent_a_data.columns)
+        self.assertEqual(len(agent_a_data), 25)  # 5 agents * 5 steps
+        for (step, agent_id), value in agent_a_data["type_a_val"].items():
+            expected_value = (agent_id - 1) * 2 + step * 2
+            self.assertEqual(value, expected_value)
+
+        # Test MockAgentB data
+        agent_b_data = data_collector.get_agenttype_vars_dataframe(MockAgentB)
+        self.assertIn("type_b_val", agent_b_data.columns)
+        self.assertEqual(len(agent_b_data), 25)  # 5 agents * 5 steps
+        for (step, agent_id), value in agent_b_data["type_b_val"].items():
+            expected_value = (agent_id - 1) * 3 + step * 3
+            self.assertEqual(value, expected_value)
+
+    def test_agenttype_and_agent_vars(self):
+        """Test that agent-type-specific and general agent variables are collected correctly."""
+        data_collector = self.model.datacollector
+
+        agent_vars = data_collector.get_agent_vars_dataframe()
+        agent_a_vars = data_collector.get_agenttype_vars_dataframe(MockAgentA)
+        agent_b_vars = data_collector.get_agenttype_vars_dataframe(MockAgentB)
+
+        # Check that general agent variables are present for all agents
+        self.assertIn("value", agent_vars.columns)
+
+        # Check that agent-type-specific variables are only present in their respective dataframes
+        self.assertIn("type_a_val", agent_a_vars.columns)
+        self.assertNotIn("type_a_val", agent_b_vars.columns)
+        self.assertIn("type_b_val", agent_b_vars.columns)
+        self.assertNotIn("type_b_val", agent_a_vars.columns)
+
+    def test_nonexistent_agenttype(self):
+        """Test that requesting data for a non-existent agent type raises a warning."""
+        data_collector = self.model.datacollector
+
+        class NonExistentAgent(Agent):
+            pass
+
+        with self.assertWarns(UserWarning):
+            non_existent_data = data_collector.get_agenttype_vars_dataframe(
+                NonExistentAgent
+            )
+            self.assertTrue(non_existent_data.empty)
+
+    def test_agenttype_reporter_string_attribute(self):
+        """Test agent-type-specific reporter with string attribute."""
+        model = MockModelWithAgentTypes()
+        model.datacollector._new_agenttype_reporter(MockAgentA, "string_attr", "val")
+        model.step()
+
+        agent_a_data = model.datacollector.get_agenttype_vars_dataframe(MockAgentA)
+        self.assertIn("string_attr", agent_a_data.columns)
+        for (_step, agent_id), value in agent_a_data["string_attr"].items():
+            expected_value = agent_id
+            self.assertEqual(value, expected_value)
+
+    def test_agenttype_reporter_function_with_params(self):
+        """Test agent-type-specific reporter with function and parameters."""
+
+        def test_func(agent, multiplier):
+            return agent.val * multiplier
+
+        model = MockModelWithAgentTypes()
+        model.datacollector._new_agenttype_reporter(
+            MockAgentB, "func_param", [test_func, [2]]
+        )
+        model.step()
+
+        agent_b_data = model.datacollector.get_agenttype_vars_dataframe(MockAgentB)
+        self.assertIn("func_param", agent_b_data.columns)
+        for (_step, agent_id), value in agent_b_data["func_param"].items():
+            expected_value = agent_id * 2
+            self.assertEqual(value, expected_value)
+
+    def test_agenttype_reporter_multiple_types(self):
+        """Test adding reporters for multiple agent types."""
+        model = MockModelWithAgentTypes()
+        model.datacollector._new_agenttype_reporter(
+            MockAgentA, "type_a_val", lambda a: a.type_a_val
+        )
+        model.datacollector._new_agenttype_reporter(
+            MockAgentB, "type_b_val", lambda a: a.type_b_val
+        )
+        model.step()
+
+        agent_a_data = model.datacollector.get_agenttype_vars_dataframe(MockAgentA)
+        agent_b_data = model.datacollector.get_agenttype_vars_dataframe(MockAgentB)
+
+        self.assertIn("type_a_val", agent_a_data.columns)
+        self.assertIn("type_b_val", agent_b_data.columns)
+        self.assertNotIn("type_b_val", agent_a_data.columns)
+        self.assertNotIn("type_a_val", agent_b_data.columns)
+
+    def test_agenttype_superclass_reporter(self):
+        """Test adding a reporter for a superclass of an agent type."""
+        model = MockModelWithAgentTypes()
+        model.datacollector._new_agenttype_reporter(MockAgent, "val", lambda a: a.val)
+        model.datacollector._new_agenttype_reporter(Agent, "val", lambda a: a.val)
+        for _ in range(3):
+            model.step()
+
+        super_data = model.datacollector.get_agenttype_vars_dataframe(MockAgent)
+        agent_data = model.datacollector.get_agenttype_vars_dataframe(Agent)
+        self.assertIn("val", super_data.columns)
+        self.assertIn("val", agent_data.columns)
+        self.assertEqual(len(super_data), 30)  # 10 agents * 3 steps
+        self.assertEqual(len(agent_data), 30)
+        self.assertTrue(super_data.equals(agent_data))
+
+
 if __name__ == "__main__":
     unittest.main()