Repository: projectmesa/mesa
Pull request #2300 (2024-09-19 15:06:01)

datacollector: Allow collecting data from Agent (sub)classes
Enhanced Mesa's DataCollector to allow collecting data from Agent (sub)classes, providing more flexible and granular data collection capabilities.

## Motive
To enable more comprehensive data collection in multi-agent simulations, allowing researchers to track attributes and behaviors specific to different agent types, including custom Agent subclasses.

## Implementation
- Modified `DataCollector` class to accept `agenttype_reporters` parameter
- Added `_new_agenttype_reporter` method for handling agent-type-specific reporters
- Updated `collect` method to handle agent-type-specific data collection
- Added `get_agenttype_vars_dataframe` method for retrieving agent-type-specific data
- Updated `Model` class to support `agenttype_reporters` in `initialize_data_collector`
- Added support for collecting data from all Agent subclasses, not just predefined agent types
- Updated docstrings and module-level documentation
- Added comprehensive unit tests for the new functionality

## Usage Examples
```python
class MyModel(Model):
    def __init__(self):
        super().__init__()
        self.datacollector = DataCollector(
            agent_reporters={"life_span": "life_span"},
            # The new agenttype_reporters argument
            agenttype_reporters={
                Wolf: {"sheep_eaten": "sheep_eaten"},
                Sheep: {"wool": "wool_amount"},
                Animal: {"energy": "energy"}  # Collects from all animals
            }
        )

# Retrieve data for a specific agent type
wolf_data = model.datacollector.get_agenttype_vars_dataframe(Wolf)
# Retrieve data for all Animal subclasses, which are in this case all Wolf and Sheep
animal_data = model.datacollector.get_agenttype_vars_dataframe(Animal)
```

## Additional Notes
- Backward compatible with existing DataCollector usage
- Supports collecting data from custom Agent subclasses and superclasses

Part of https://github.com/projectmesa/mesa/issues/348.

github-actions[bot]: Performance benchmarks:

| Model           | Size   | Init time [95% CI]      | Run time [95% CI]       |
|-----------------|--------|-------------------------|-------------------------|
| BoltzmannWealth | small  | 🔵 +1.3% [+0.1%, +2.5%] | 🔵 +0.2% [-0.1%, +0.5%] |
| BoltzmannWealth | large  | 🔵 -0.3% [-0.9%, +0.2%] | 🔵 -1.4% [-3.6%, +1.8%] |
| Schelling       | small  | 🔵 +0.1% [-0.3%, +0.4%] | 🔵 -0.4% [-0.6%, -0.2%] |
| Schelling       | large  | 🔵 +0.1% [-0.6%, +0.8%] | 🔵 -2.1% [-3.4%, -0.8%] |
| WolfSheep       | small  | 🔵 +0.9% [+0.6%, +1.1%] | 🔵 +0.0% [-0.2%, +0.3%] |
| WolfSheep       | large  | 🔵 -3.0% [-4.1%, -1.7%] | 🟢 -6.4% [-7.6%, -5.3%] |
| BoidFlockers    | small  | 🔵 +1.0% [+0.3%, +1.5%] | 🔵 -0.2% [-0.9%, +0.5%] |
| BoidFlockers    | large  | 🔵 +0.8% [+0.3%, +1.2%] | 🔵 +0.5% [-0.3%, +1.2%] |

rht: What about the case for e.g. in Epstein civil violence, you want to collect based on citizens that are jailed, which is not based on type?
EwoutH: This PR doesn't fix everything, we clearly need a new DataCollector. But it's a minimally invasive, fully backwards compatible extension to the current DataCollector that everybody knows. It's even backportable to Mesa 2.x.

In Epstein civil violence, you can log if agents are jailed or not and filter afterwards. Not ideal, but workable.

Of course I'm open for other suggestions for backwards compatible alternatives.
quaquel: There are a few lines currently not covered by the tests, but otherwise, this seems fine. It solves a long-standing open issue in a backward-compatible way. It makes sense to add it while we continue work on building the foundations for a better datacollector. 
EwoutH: Added some additional tests.

I consciously named it `agenttype_reporters`, to leave room for future other arguments like `agentset_ reporters` or `custom_agent_reporters`.

This functionality might even be able to be implemented more elegantly using more existing code.

But I think API is solid, and there are distinct use cases for it. One huge advantage of collecting by type, it that you can ensure all agents have those attributes and methods you’re collecting. It also elegantly using the pre-existing `model.agent_by_types` duct, which is a very fast way to get these agent groups (since select can be an expensive operation).
EwoutH: In 86b184d and 3dbb5b6 I updated this feature to allow collecting from any Agent (sub)class, so also a class like `Animals`, which is not present in `model.agents_by_type` but is a valid Agent (sub)class. It even allows inputting `Agent` itself, but that's ofcourse redundant with the `agent_reporters`.

```python
class MyModel(Model):
    def __init__(self):
        super().__init__()
        self.datacollector = DataCollector(
            agent_reporters={"life_span": "life_span"},
            # The new agenttype_reporters argument
            agenttype_reporters={
                Wolf: {"sheep_eaten": "sheep_eaten"},
                Sheep: {"wool": "wool_amount"},
                Animal: {"energy": "energy"}  # Collects from all animals
            }
        )

I would love some final reviews on this (also from @Corvince).

Patch:
diff --git a/.github/workflows/build_lint.yml b/.github/workflows/build_lint.yml
index d407b178..5e59c2f0 100644
--- a/.github/workflows/build_lint.yml
+++ b/.github/workflows/build_lint.yml
@@ -53,12 +53,8 @@ jobs:
         cache: 'pip'
     - name: Install uv
       run: pip install uv
-    - name: Install Mesa with uv pip
+    - name: Install Mesa and dependencies
       run: uv pip install --system .[dev]
-      if: matrix.python-version != '3.13'
-    - name: Install Mesa with pip
-      run: pip install .[dev]
-      if: matrix.python-version == '3.13'
     - name: Test with pytest
       run: pytest --durations=10 --cov=mesa tests/ --cov-report=xml
     - if: matrix.os == 'ubuntu'
diff --git a/docs/migration_guide.md b/docs/migration_guide.md
index 9097ea2e..c8c06860 100644
--- a/docs/migration_guide.md
+++ b/docs/migration_guide.md
@@ -29,10 +29,49 @@ The `mesa.flat` namespace is removed. Use the full namespace for your imports.
 - Ref: [PR #2091](https://github.com/projectmesa/mesa/pull/2091)
 
 
+### Mandatory Model initialization with `super().__init__()`
+In Mesa 3.0, it is now mandatory to call `super().__init__()` when initializing your model class. This ensures that all necessary Mesa model variables are correctly set up and agents are properly added to the model.
+
+Make sure all your model classes explicitly call `super().__init__()` in their `__init__` method:
+
+```python
+class MyModel(mesa.Model):
+    def __init__(self, *args, **kwargs):
+        super().__init__()  # This is now required!
+        # Your model initialization code here
+```
+
+This change ensures that all Mesa models are properly initialized, which is crucial for:
+- Correctly adding agents to the model
+- Setting up other essential Mesa model variables
+- Maintaining consistency across all models
+
+If you forget to call `super().__init__()`, you'll now see this error:
+
+```
+RuntimeError: The Mesa Model class was not initialized. You must explicitly initialize the Model by calling super().__init__() on initialization.
+```
+
+- Ref: [PR #2218](https://github.com/projectmesa/mesa/pull/2218), [PR #1928](https://github.com/projectmesa/mesa/pull/1928), Mesa-examples [PR #83](https://github.com/projectmesa/mesa-examples/pull/83)
+
+
 ### Automatic assignment of `unique_id` to Agents
-<!-- TODO -->
+In Mesa 3.0, `unique_id` for agents is now automatically assigned, simplifying agent creation and ensuring unique IDs across all agents in a model.
+
+1. Remove `unique_id` from agent initialization:
+   ```python
+   # Old
+   agent = MyAgent(self.next_id(), self, ...)
 
-- Ref: [PR #2226](https://github.com/projectmesa/mesa/pull/2226)
+   # New
+   agent = MyAgent(self, ...)
+   ```
+2. `Model.next_id()` is deprecated and will always return 0. Remove any calls to this method.
+3. `unique_id` is now unique relative to a Model instance and starts from 1.
+4. If you previously used custom `unique_id` values, you'll need to store that information in a separate attribute.
+5. Deprecation warning: Initializing an agent with two arguments (`unique_id` and `model`) will raise a warning. The `unique_id` argument will be ignored.
+
+- Ref: [PR #2226](https://github.com/projectmesa/mesa/pull/2226), [PR #2260](https://github.com/projectmesa/mesa/pull/2260), Mesa-examples [PR #194](https://github.com/projectmesa/mesa-examples/pull/194), [Issue #2213](https://github.com/projectmesa/mesa/issues/2213)
 
 
 ### AgentSet and `Model.agents`
@@ -64,4 +103,67 @@ You can access it by `Model.steps`, and it's internally in the datacollector, ba
 
 
 ### Visualisation
-<!-- TODO -->
+
+Mesa has adopted a new API for our frontend. If you already migrated to the experimental new SolaraViz you can still use
+the import from mesa.experimental. Otherwise here is a list of things you need to change.
+
+#### Model Initialization
+
+Previously SolaraViz was initialized by providing a `model_cls` and a `model_params`. This has changed to expect a model instance `model`. You can still provide (user-settable) `model_params`, but only if users should be able to change them. It is now also possible to pass in a "reactive model" by first calling `model = solara.reactive(model)`. This is useful for notebook environments. It allows you to pass the model to the SolaraViz Module, but continue to use the model. For example calling `model.value.step()` (notice the extra .value) will automatically update the plots. This currently only automatically works for the step method, you can force visualization updates by calling `model.value.force_update()`.
+
+#### Default space visualization
+
+Previously we included a default space drawer that you could configure with an `agent_portrayal` function. You now have to explicitly create a space drawer with the `agent_portrayal` function
+
+```python
+# old
+from mesa.experimental import SolaraViz
+
+SolaraViz(model_cls, model_params, agent_portrayal=agent_portrayal)
+
+# new
+from mesa.visualization import SolaraViz, make_space_matplotlib
+
+SolaraViz(model, components=[make_space_matplotlib(agent_portrayal)])
+```
+
+#### Plotting "measures"
+
+"Measure" plots also need to be made explicit here. Previously, measure could either be 1) A function that receives a model and returns a solara component or 2) A string or list of string of variables that are collected by the datacollector and are to be plotted as a line plot. 1) still works, but you can pass that function to "components" directly. 2) needs to explicitly call the `make_plot_measure()`function.
+
+```python
+# old
+from mesa.experimental import SolaraViz
+
+def make_plot(model):
+    ...
+
+SolaraViz(model_cls, model_params, measures=[make_plot, "foo", ["bar", "baz"]])
+
+# new
+from mesa.visualization import SolaraViz, make_plot_measure
+
+SolaraViz(model, components=[make_plot, make_plot_measure("foo"), make_plot_measure("bar", "baz")])
+```
+
+#### Plotting text
+
+To plot model-dependent text the experimental SolaraViz provided a `make_text` function that wraps another functions that receives the model and turns its string return value into a solara text component. Again, this other function can now be passed directly to the new SolaraViz components array. It is okay if your function just returns a string.
+
+```python
+# old
+from mesa.experimental import SolaraViz, make_text
+
+def show_steps(model):
+    return f"Steps: {model.steps}"
+
+SolaraViz(model_cls, model_params, measures=make_text(show_steps))
+
+# new
+from mesa.visualisation import SolaraViz
+
+def show_steps(model):
+    return f"Steps: {model.steps}"
+
+SolaraViz(model, components=[show_steps])
+```
diff --git a/mesa/agent.py b/mesa/agent.py
index 5badc825..2d098c35 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -15,7 +15,7 @@ import operator
 import warnings
 import weakref
 from collections import defaultdict
-from collections.abc import Callable, Iterable, Iterator, MutableSet, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, MutableSet, Sequence
 from random import Random
 
 # mypy
@@ -611,6 +611,29 @@ class GroupBy:
 
         return self
 
+    def count(self) -> dict[Any, int]:
+        """Return the count of agents in each group.
+
+        Returns:
+            dict: A dictionary mapping group names to the number of agents in each group.
+        """
+        return {k: len(v) for k, v in self.groups.items()}
+
+    def agg(self, attr_name: str, func: Callable) -> dict[Hashable, Any]:
+        """Aggregate the values of a specific attribute across each group using the provided function.
+
+        Args:
+            attr_name (str): The name of the attribute to aggregate.
+            func (Callable): The function to apply (e.g., sum, min, max, mean).
+
+        Returns:
+            dict[Hashable, Any]: A dictionary mapping group names to the result of applying the aggregation function.
+        """
+        return {
+            group_name: func([getattr(agent, attr_name) for agent in group])
+            for group_name, group in self.groups.items()
+        }
+
     def __iter__(self):  # noqa: D105
         return iter(self.groups.items())
 
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/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.
diff --git a/mesa/visualization/solara_viz.py b/mesa/visualization/solara_viz.py
index 129a9486..ed495732 100644
--- a/mesa/visualization/solara_viz.py
+++ b/mesa/visualization/solara_viz.py
@@ -23,8 +23,8 @@ See the Visualization Tutorial and example models for more details.
 
 from __future__ import annotations
 
+import asyncio
 import copy
-import time
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Literal
 
@@ -235,12 +235,12 @@ def ModelController(model: solara.Reactive[Model], play_interval=100):
 
     solara.use_effect(save_initial_model, [model.value])
 
-    def step():
+    async def step():
         while playing.value:
-            time.sleep(play_interval / 1000)
+            await asyncio.sleep(play_interval / 1000)
             do_step()
 
-    solara.use_thread(step, [playing.value])
+    solara.lab.use_task(step, dependencies=[playing.value], prefer_threaded=False)
 
     def do_step():
         """Advance the model by one step."""


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 769be4ec..f43a80c8 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -524,6 +524,7 @@ def test_agentset_groupby():
         def __init__(self, model):
             super().__init__(model)
             self.even = self.unique_id % 2 == 0
+            self.value = self.unique_id * 10
 
         def get_unique_identifier(self):
             return self.unique_id
@@ -560,6 +561,37 @@ def test_agentset_groupby():
     another_ref_to_groups = groups.do(lambda x: x.do("step"))
     assert groups == another_ref_to_groups
 
+    # New tests for count() method
+    groups = agentset.groupby("even")
+    count_result = groups.count()
+    assert count_result == {True: 5, False: 5}
+
+    # New tests for agg() method
+    groups = agentset.groupby("even")
+    sum_result = groups.agg("value", sum)
+    assert sum_result[True] == sum(agent.value for agent in agents if agent.even)
+    assert sum_result[False] == sum(agent.value for agent in agents if not agent.even)
+
+    max_result = groups.agg("value", max)
+    assert max_result[True] == max(agent.value for agent in agents if agent.even)
+    assert max_result[False] == max(agent.value for agent in agents if not agent.even)
+
+    min_result = groups.agg("value", min)
+    assert min_result[True] == min(agent.value for agent in agents if agent.even)
+    assert min_result[False] == min(agent.value for agent in agents if not agent.even)
+
+    # Test with a custom aggregation function
+    def custom_agg(values):
+        return sum(values) / len(values) if values else 0
+
+    custom_result = groups.agg("value", custom_agg)
+    assert custom_result[True] == custom_agg(
+        [agent.value for agent in agents if agent.even]
+    )
+    assert custom_result[False] == custom_agg(
+        [agent.value for agent in agents if not agent.even]
+    )
+
 
 def test_oldstyle_agent_instantiation():
     """Old behavior of Agent creation with unique_id and model as positional arguments.
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()