"""Modelly Data Collection Module.
DataCollector is meant to provide a simple, standard way to collect data
generated by a Modelly model. It collects four types of data: model-level data,
agent-level data, agent-type-level data, and tables.
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, 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.
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 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,
and a pyarrow Table or a Parquet file (``to_arrow``, ``to_parquet``).
"""
import contextlib
from modelly.core.agent import Agent
from modelly.data.frames import DataFrameExportMixin
from modelly.data.io import ArrowExportMixin
from modelly.data.reporters import (
check_list_reporter,
collect_agent_rows,
compile_model_reporter,
make_agent_reporter,
snapshot,
validate_model_reporter,
)
from modelly.data.tables import TableMixin
with contextlib.suppress(ImportError):
import pandas as pd
[docs]
class DataCollector(TableMixin, DataFrameExportMixin, ArrowExportMixin):
"""Class for collecting data generated by a Modelly model.
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, agent, and agent-type reporters.
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, a partial function, or a function with
parameters placed in a list.
Model reporters can take five types of arguments::
{"agent_count": lambda m: len(m.agents)} # lambda
{"agent_count": self.get_agent_count} # bound method
{"agent_count": Model.get_agent_count} # method of a class
{"model_attribute": "model_attribute"} # attribute name
{"agent_count": functools.partial(count_agents, multiplier=2)}
{"Model_Function": [function, [param_1, param_2]]} # function + params
Agent reporters can similarly take::
{"energy": "energy"} # attribute name
{"energy": lambda a: a.energy} # lambda
{"agent_action": self.do_action} # bound method
{"agent_action": Agent.do_action} # method of a class
{"energy": functools.partial(get_energy, scale=2)}
{"Agent_Function": [function, [param_1, param_2]]} # function + params
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 like::
{"Lifespan": ["unique_id", "age"]}
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:
- 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.
- agenttype_reporters collect only agents of exactly the given type.
Subclasses are not included, so a reporter keyed on a superclass does
not pick up its subclasses. Add a separate reporter per concrete type
if you need subclass data.
"""
self._init_storage()
if model_reporters is not None:
for name, reporter in model_reporters.items():
self._new_model_reporter(name, reporter)
if agent_reporters is not None:
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)
def _init_storage(self):
"""Create the empty reporter registries and record stores."""
self.model_reporters = {}
self.agent_reporters = {}
self.agenttype_reporters = {}
self.model_vars = {}
# name -> (reporter, callable(model)) compiled once at registration
self._model_getters = {}
self._collection_steps = []
self._agent_records = {}
self._agenttype_records = {}
self.tables = {}
# (steps already in the frame, frame): get_agent_vars_dataframe only
# converts the steps collected since the last call.
self._agent_frame_cache: tuple[int, pd.DataFrame] | None = None
# add the signal of the validation of model reporter
self._validated = False
@staticmethod
def _check_list_reporter(name, reporter):
"""Validate that a list-style reporter has the form [callable, [params]]."""
check_list_reporter(name, reporter)
def _validate_model_reporter(self, name, reporter, model):
"""Validate model reporter and handle validation results appropriately.
Args:
name: Name of the reporter
reporter: Reporter definition (lambda/method/attribute/function list)
model: Model instance
Raises:
ValueError: If reporter is None or has invalid format
AttributeError: If model attribute doesn't exist
TypeError: If reporter type is not supported
RuntimeError: If reporter execution fails
"""
self._validated = True
validate_model_reporter(name, reporter, model)
def _new_model_reporter(self, name, reporter):
"""Add a new model-level reporter to collect.
Args:
name: Name of the model-level variable to collect.
reporter: Can be one of four types:
1. Attribute name (str): "attribute_name"
2. Lambda function: lambda m: len(m.agents)
3. Method: model.get_count or Model.get_count
4. List of [function, [parameters]]
"""
if isinstance(reporter, list):
check_list_reporter(name, reporter)
self.model_reporters[name] = reporter
self.model_vars[name] = []
self._model_getters[name] = (reporter, compile_model_reporter(name, reporter))
def _new_agent_reporter(self, name, reporter):
"""Add a new agent-level reporter to collect.
Args:
name: Name of the agent-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.
"""
self.agent_reporters[name] = make_agent_reporter(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.
"""
reporters = self.agenttype_reporters.setdefault(agent_type, {})
reporters[name] = make_agent_reporter(name, reporter)
def _record_agents(self, model):
"""Record agents data in a mapping of functions and agents."""
return collect_agent_rows(model.agents, self.agent_reporters, model.time)
def _record_agenttype(self, model, agent_type):
"""Record agent-type data in a mapping of functions and agents."""
if not (isinstance(agent_type, type) and issubclass(agent_type, Agent)):
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 {model.agent_types}."
)
# Collect only agents whose type is exactly agent_type. Subclasses are not
# included: a reporter keyed on a superclass does not pick up its subclasses.
# This is consistent regardless of whether agent_type has direct instances.
agents = model.agents_by_type.get(agent_type, [])
return collect_agent_rows(
agents, self.agenttype_reporters[agent_type], model.time
)
[docs]
def collect(self, model):
"""Collect all the data for the given model object."""
time = model.time
if self.model_reporters:
self._collect_model_vars(model, time)
if self.agent_reporters:
self._agent_records[time] = self._record_agents(model)
if self.agenttype_reporters:
self._agenttype_records[time] = {
agent_type: self._record_agenttype(model, agent_type)
for agent_type in self.agenttype_reporters
}
def _collect_model_vars(self, model, time):
"""Append one value per model reporter for the current step."""
self._collection_steps.append(time)
if not self._validated:
for name, reporter in self.model_reporters.items():
self._validate_model_reporter(name, reporter, model)
model_vars = self.model_vars
getters = self._model_getters
for var, reporter in self.model_reporters.items():
compiled = getters.get(var)
if compiled is None or compiled[0] is not reporter:
# a reporter was replaced after registration: compile the new one
compiled = getters[var] = (
reporter,
compile_model_reporter(var, reporter),
)
model_vars[var].append(snapshot(compiled[1](model)))
