{ "cells": [ { "metadata": {}, "cell_type": "markdown", "source": [ "# Working with AgentSets\n", "### The Boltzmann Wealth Model" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "**Important:**\n", "- If you are just exploring Mesa and want the fastest way to execute the code we recommend executing this tutorial online in a Colab notebook. [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mesa/mesa/blob/main/docs/tutorials/1_agentset.ipynb) or if you do not have a Google account you can use [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/mesa/mesa/main?labpath=docs%2Ftutorials%2F1_agentset.ipynb) (This can take 30 seconds to 5 minutes to load)\n", "- If you are running locally, please ensure you have the latest Mesa version installed.\n", "## Tutorial Description\n", "This tutorial builds on the Boltzmann Wealth Model from the [First Model tutorial](0_first_model.ipynb).\n", "In the first tutorial you created agents, put them in a model, and had them exchange money.\n", "Now we'll explore **AgentSet** — Mesa's core tool for querying, filtering, grouping, and\n", "inspecting collections of agents.\n", "By the end of this tutorial, you will know how to:\n", "- Retrieve attribute values from agents\n", "- Filter agents with `select`\n", "- Compute aggregate statistics with `agg`\n", "- Group agents by attributes with `groupby`\n", "- Combine these tools to answer questions about your model's state\n", "The *next* tutorial covers how to use AgentSet for **activating** agents (calling their\n", "methods). We separate these concerns because querying agents and activating agents are\n", "conceptually different — you'll often query first and activate a subset." ] }, { "metadata": {}, "cell_type": "markdown", "source": "### IN COLAB? - Run the next cell" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": "# %pip install --quiet mesa[rec]" }, { "metadata": {}, "cell_type": "markdown", "source": "### Import Dependencies" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "import numpy as np\n", "import pandas as pd\n", "import seaborn as sns\n", "\n", "import mesa" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Setup: The Wealth Model with Ethnicities\n", "We'll use a slightly enriched version of the Boltzmann Wealth Model. Each agent has a\n", "`wealth` (starting at 1) and an `ethnicity` (randomly assigned from \"Green\", \"Blue\",\n", "or \"Mixed\"). This gives us meaningful attributes to query, filter, and group by." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "class MoneyAgent(mesa.Agent):\n", " \"\"\"An agent with fixed initial wealth and an ethnicity.\"\"\"\n", "\n", " def __init__(self, model, ethnicity):\n", " super().__init__(model)\n", " self.wealth = 1\n", " self.ethnicity = ethnicity\n", "\n", " def exchange(self):\n", " if self.wealth > 0:\n", " other_agent = self.random.choice(self.model.agents)\n", " other_agent.wealth += 1\n", " self.wealth -= 1\n", "\n", "\n", "class MoneyModel(mesa.Model):\n", " \"\"\"A model with some number of agents.\"\"\"\n", "\n", " def __init__(self, n=100):\n", " super().__init__()\n", " ethnicities = [\"Green\", \"Blue\", \"Mixed\"]\n", " MoneyAgent.create_agents(\n", " model=self,\n", " n=n,\n", " ethnicity=self.random.choices(ethnicities, k=n),\n", " )\n", "\n", " def step(self):\n", " self.agents.shuffle_do(\"exchange\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": "Let's create a model and run it for a while so agents have different wealth levels." }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "model = MoneyModel(100)\n", "model.run_for(50)" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## What is an AgentSet?\n", "Every Mesa model automatically tracks all its agents in `model.agents`. This is an\n", "**AgentSet** — an ordered collection of agents that provides powerful methods for\n", "querying, filtering, and manipulating groups of agents.\n", "You never need to create an AgentSet yourself for basic usage. Mesa creates and maintains\n", "`model.agents` automatically whenever agents are added to or removed from the model.\n", "Let's look at some basics:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# How many agents are in the model?\n", "print(f\"Total agents: {len(model.agents)}\")\n", "\n", "# Iterate over agents (just the first 5 for brevity)\n", "for agent in model.agents.select(at_most=5):\n", " print(\n", " f\" Agent {agent.unique_id}: wealth={agent.wealth}, ethnicity={agent.ethnicity}\"\n", " )" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Retrieving Attribute Values with `get`\n", "The `get` method retrieves attribute values from every agent in the set, returning them\n", "as a list. This is useful whenever you want to inspect or analyze a particular attribute\n", "across all agents." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Get all wealth values\n", "all_wealth = model.agents.get(\"wealth\")\n", "print(f\"First 10 wealth values: {all_wealth[:10]}\")\n", "print(f\"Total wealth in economy: {sum(all_wealth)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "You can also retrieve multiple attributes at once by passing a list of attribute names.\n", "This returns a list of lists — one inner list per agent." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Get both wealth and ethnicity for each agent\n", "wealth_and_ethnicity = model.agents.get([\"wealth\", \"ethnicity\"])\n", "print(\"First 5 agents (wealth, ethnicity):\")\n", "for values in wealth_and_ethnicity[:5]:\n", " print(f\" {values}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Handling missing attributes\n", "If some agents might not have a particular attribute, you can use the `handle_missing`\n", "parameter. By default, `get` raises an `AttributeError` for missing attributes. Setting\n", "`handle_missing=\"default\"` returns a default value instead." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# This would raise AttributeError if any agent lacks 'wealth':\n", "# model.agents.get(\"nonexistent_attr\")\n", "\n", "# Safe alternative — returns None for missing attributes:\n", "values = model.agents.get(\"wealth\", handle_missing=\"default\", default_value=0)\n", "print(f\"Retrieved {len(values)} values safely\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Filtering Agents with `select`\n", "The `select` method filters agents based on criteria, returning a new AgentSet containing\n", "only the agents that match. This is one of the most frequently used AgentSet operations.\n", "### Basic filtering with a function\n", "Pass a function (often a lambda) that takes an agent and returns `True` or `False`:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Select only wealthy agents (wealth >= 3)\n", "rich_agents = model.agents.select(lambda a: a.wealth >= 3)\n", "print(f\"Rich agents (wealth >= 3): {len(rich_agents)}\")\n", "\n", "# Select agents with no money\n", "broke_agents = model.agents.select(lambda a: a.wealth == 0)\n", "print(f\"Broke agents (wealth == 0): {len(broke_agents)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Filtering by agent type\n", "If your model has multiple agent classes, you can filter by type using the `agent_type`\n", "parameter. This is faster than using a lambda with `isinstance`." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# In this model we only have one type, but the syntax would be:\n", "money_agents = model.agents.select(agent_type=MoneyAgent)\n", "print(f\"MoneyAgents: {len(money_agents)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Limiting results with `at_most`\n", "The `at_most` parameter limits how many agents are returned. This is useful when you\n", "only need a few agents and want to avoid processing the entire set.\n", "- Pass an **integer** to get at most that many agents\n", "- Pass a **float between 0 and 1** to get at most that fraction of agents\n", "**Important:** `at_most` returns the *first* matching agents, not a random sample.\n", "If you want a random subset, call `shuffle()` first (covered in the activation tutorial)." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Get at most 5 rich agents\n", "some_rich = model.agents.select(lambda a: a.wealth >= 2, at_most=5)\n", "print(f\"Up to 5 rich agents: {len(some_rich)}\")\n", "\n", "# Get roughly 10% of agents\n", "ten_percent = model.agents.select(at_most=0.1)\n", "print(f\"~10% of agents: {len(ten_percent)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Combining criteria\n", "You can combine `filter_func`, `agent_type`, and `at_most` in a single call. All\n", "criteria are applied together (logical AND):" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# At most 10 MoneyAgents with wealth > 0\n", "subset = model.agents.select(\n", " filter_func=lambda a: a.wealth > 0,\n", " agent_type=MoneyAgent,\n", " at_most=10,\n", ")\n", "print(f\"Subset size: {len(subset)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Chaining selects\n", "Since `select` returns an AgentSet, you can chain multiple calls. Each successive\n", "`select` narrows the set further:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# First get Green agents, then filter for wealthy ones\n", "wealthy_green = model.agents.select(lambda a: a.ethnicity == \"Green\").select(\n", " lambda a: a.wealth >= 3\n", ")\n", "print(f\"Wealthy Green agents: {len(wealthy_green)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Computing Aggregates with `agg`\n", "The `agg` method computes aggregate statistics over an attribute for all agents in\n", "the set. Pass the attribute name and a function (like `min`, `max`, `sum`, or `np.mean`)." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Average wealth across all agents\n", "avg_wealth = model.agents.agg(\"wealth\", np.mean)\n", "print(f\"Average wealth: {avg_wealth:.2f}\")\n", "\n", "# Min and max wealth\n", "min_wealth = model.agents.agg(\"wealth\", min)\n", "max_wealth = model.agents.agg(\"wealth\", max)\n", "print(f\"Wealth range: {min_wealth} to {max_wealth}\")\n", "\n", "# Total wealth (should equal the number of agents, since money is conserved)\n", "total = model.agents.agg(\"wealth\", sum)\n", "print(f\"Total wealth: {total}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Multiple aggregations at once\n", "You can pass a list of functions to compute multiple statistics in a single call:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "min_w, max_w, avg_w = model.agents.agg(\"wealth\", [min, max, np.mean])\n", "print(f\"Min: {min_w}, Max: {max_w}, Mean: {avg_w:.2f}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Aggregating subsets\n", "Since `select` returns an AgentSet, you can chain `select` and `agg` to compute\n", "statistics for specific subgroups:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Average wealth of Green agents only\n", "green_avg = model.agents.select(lambda a: a.ethnicity == \"Green\").agg(\"wealth\", np.mean)\n", "blue_avg = model.agents.select(lambda a: a.ethnicity == \"Blue\").agg(\"wealth\", np.mean)\n", "mixed_avg = model.agents.select(lambda a: a.ethnicity == \"Mixed\").agg(\"wealth\", np.mean)\n", "\n", "print(\n", " f\"Average wealth — Green: {green_avg:.2f}, Blue: {blue_avg:.2f}, Mixed: {mixed_avg:.2f}\"\n", ")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "This pattern of select-then-aggregate is common, but when you want to do it for *all*\n", "groups at once, `groupby` is more elegant." ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Grouping Agents with `groupby`\n", "The `groupby` method splits agents into groups based on an attribute (or a callable),\n", "returning a `GroupBy` object. This is conceptually similar to pandas' `groupby` and\n", "is ideal when you want to analyze or act on agents by category." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Group agents by ethnicity\n", "grouped = model.agents.groupby(\"ethnicity\")\n", "\n", "# See how many agents are in each group\n", "print(\"Agents per ethnicity:\", grouped.count())" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Iterating over groups\n", "A `GroupBy` object is iterable. Each iteration yields a `(group_name, agent_set)` tuple:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "for ethnicity, group in grouped:\n", " avg = group.agg(\"wealth\", np.mean)\n", " print(f\" {ethnicity}: {len(group)} agents, avg wealth = {avg:.2f}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Aggregating across groups\n", "The `agg` method on `GroupBy` computes an aggregate for each group in one call:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Mean wealth by ethnicity\n", "mean_by_group = grouped.agg(\"wealth\", np.mean)\n", "print(\"Mean wealth by ethnicity:\", mean_by_group)" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Grouping by a function\n", "Instead of an attribute name, you can pass a callable that computes the group key\n", "for each agent. This is useful for creating custom groupings:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Group agents into wealth brackets\n", "def wealth_bracket(agent):\n", " if agent.wealth == 0:\n", " return \"broke\"\n", " elif agent.wealth <= 2:\n", " return \"modest\"\n", " else:\n", " return \"wealthy\"\n", "\n", "\n", "brackets = model.agents.groupby(wealth_bracket)\n", "print(\"Agents per wealth bracket:\", brackets.count())" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Setting Attributes with `set`\n", "The `set` method assigns a value to an attribute for all agents in the set. This is\n", "useful for bulk updates — for example, applying a policy change to a group of agents." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Give all broke agents a subsidy of 1\n", "broke = model.agents.select(lambda a: a.wealth == 0)\n", "print(f\"Broke agents before subsidy: {len(broke)}\")\n", "\n", "broke.set(\"wealth\", 1)\n", "\n", "# Verify\n", "still_broke = model.agents.select(lambda a: a.wealth == 0)\n", "print(f\"Broke agents after subsidy: {len(still_broke)}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "**Note:** `set` modifies agents in place and returns the AgentSet, so you can chain it:\n", "```python\n", "model.agents.select(lambda a: a.wealth > 10).set(\"taxed\", True)\n", "```" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Sorting Agents with `sort`\n", "The `sort` method orders agents by an attribute or a custom key function. By default,\n", "it returns a new sorted AgentSet (use `inplace=True` to sort in place)." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Sort by wealth (descending by default)\n", "richest_first = model.agents.sort(\"wealth\")\n", "top_5 = richest_first.select(at_most=5)\n", "print(\"Top 5 wealthiest agents:\")\n", "for agent in top_5:\n", " print(f\" Agent {agent.unique_id}: wealth={agent.wealth}\")\n", "\n", "# Sort ascending\n", "poorest_first = model.agents.sort(\"wealth\", ascending=True)\n", "bottom_5 = poorest_first.select(at_most=5)\n", "print(\"\\nBottom 5:\")\n", "for agent in bottom_5:\n", " print(f\" Agent {agent.unique_id}: wealth={agent.wealth}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Converting to a List\n", "If you need standard list operations like indexing or slicing, use the `to_list()` method\n", "to convert the AgentSet to a plain Python list:" ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "agent_list = model.agents.to_list()\n", "print(f\"First agent: {agent_list[0].unique_id}\")\n", "print(f\"Last agent: {agent_list[-1].unique_id}\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Putting It Together: Analyzing the Model\n", "Let's combine what we've learned to produce a summary analysis of the model state." ] }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "print(\"=== Model Summary After 50 Steps ===\\n\")\n", "\n", "# Overall statistics\n", "min_w, max_w, avg_w, total_w = model.agents.agg(\"wealth\", [min, max, np.mean, sum])\n", "print(f\"Agents: {len(model.agents)}\")\n", "print(\n", " f\"Total wealth: {total_w} (conserved: {'yes' if total_w == len(model.agents) else 'no, subsidy applied'})\"\n", ")\n", "print(f\"Wealth range: {min_w} to {max_w}, mean: {avg_w:.2f}\\n\")\n", "\n", "# By ethnicity\n", "print(\"By ethnicity:\")\n", "for ethnicity, group in model.agents.groupby(\"ethnicity\"):\n", " count = len(group)\n", " avg = group.agg(\"wealth\", np.mean)\n", " broke = len(group.select(lambda a: a.wealth == 0))\n", " print(\n", " f\" {ethnicity:6s}: {count:3d} agents, avg wealth = {avg:.2f}, broke = {broke}\"\n", " )\n", "\n", "# Wealth distribution\n", "print(\"\\nWealth brackets:\")\n", "for bracket, group in model.agents.groupby(wealth_bracket):\n", " print(f\" {bracket:8s}: {len(group)} agents\")" ] }, { "metadata": {}, "cell_type": "markdown", "source": "## Visualizing the Results" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Collect data for plotting\n", "data = []\n", "for agent in model.agents:\n", " data.append({\"wealth\": agent.wealth, \"ethnicity\": agent.ethnicity})\n", "df = pd.DataFrame(data)\n", "\n", "palette = {\"Green\": \"green\", \"Blue\": \"blue\", \"Mixed\": \"purple\"}\n", "g = sns.histplot(data=df, x=\"wealth\", hue=\"ethnicity\", discrete=True, palette=palette)\n", "g.set(\n", " title=\"Wealth distribution by ethnicity\", xlabel=\"Wealth\", ylabel=\"Number of agents\"\n", ")" ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Summary\n", "In this tutorial you learned the core AgentSet **query** methods:\n", "| Method | Purpose |\n", "|---|---|\n", "| `get(attr)` | Retrieve attribute values from all agents |\n", "| `select(func)` | Filter agents by criteria |\n", "| `agg(attr, func)` | Compute aggregate statistics |\n", "| `groupby(attr)` | Group agents by attribute or function |\n", "| `set(attr, value)` | Bulk-assign attribute values |\n", "| `sort(key)` | Order agents by attribute |\n", "| `to_list()` | Convert to a plain Python list |\n", "These methods are about *inspecting* and *organizing* agents. In the next tutorial,\n", "we'll cover how to **activate** agents — making them actually *do* things — using\n", "`do`, `shuffle_do`, and `map`." ] }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Next Steps\n", "Check out the [Agent Activation tutorial](2_agent_activation)\n", "to learn how to make your agents act, in different orders and patterns.\n" ] } ], "metadata": { "anaconda-cloud": {}, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.5" }, "widgets": { "state": {}, "version": "1.1.2" } }, "nbformat": 4, "nbformat_minor": 4 }