diff --git a/.github/workflows/docs-gh-pages.yml b/.github/workflows/docs-gh-pages.yml index 435af957..d60be123 100644 --- a/.github/workflows/docs-gh-pages.yml +++ b/.github/workflows/docs-gh-pages.yml @@ -1,40 +1,101 @@ -name: Build and Deploy Documentation +name: Docs โ€” Build & Preview on: push: - branches: - - main + branches: [ main ] # regular prod deploy + paths: + - 'mkdocs.yml' + - 'docs/**' + pull_request: # preview only when docs are touched + branches: [ '**' ] + paths: + - 'mkdocs.yml' + - 'docs/**' jobs: - build-and-deploy-docs: + build: runs-on: ubuntu-latest + outputs: + short_sha: ${{ steps.sha.outputs.short }} steps: - uses: actions/checkout@v4 - with: - fetch-depth: 0 # Fetch all history for .git-restore-mtime to work correctly - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.x' + with: { fetch-depth: 0 } + - uses: actions/setup-python@v5 + with: { python-version: '3.x' } + - uses: astral-sh/setup-uv@v6 - - name: Install uv via GitHub Action - uses: astral-sh/setup-uv@v6 - - - name: Install mesa-frames + docs dependencies + - name: Install mesa-frames + docs deps run: | uv pip install --system . uv pip install --group docs --system - - name: Build MkDocs site (general documentation) - run: mkdocs build --config-file mkdocs.yml --site-dir ./site + - name: Convert jupytext .py notebooks to .ipynb + run: | + set -euxo pipefail + # Convert any jupytext .py files to .ipynb without executing them. + # Enable nullglob so the pattern expands to empty when there are no matches + # and globstar so we recurse into subdirectories (e.g., user-guide/). + shopt -s nullglob globstar || true + files=(docs/general/**/*.py) + if [ ${#files[@]} -eq 0 ]; then + echo "No jupytext .py files found under docs/general" + else + for src in "${files[@]}"; do + [ -e "$src" ] || continue + dest="${src%.py}.ipynb" + echo "Converting $src -> $dest" + # jupytext will write the .ipynb alongside the source file + uv run jupytext --to notebook "$src" + done + fi + + - name: Build MkDocs site + run: uv run mkdocs build --config-file mkdocs.yml --site-dir ./site + + - name: Build Sphinx docs (API) + run: uv run sphinx-build -b html docs/api site/api - - name: Build Sphinx docs (API documentation) - run: sphinx-build -b html docs/api site/api + - name: Short SHA + id: sha + run: echo "short=$(git rev-parse --short HEAD)" >> "$GITHUB_OUTPUT" - - name: Deploy to GitHub Pages + - name: Upload site artifact + uses: actions/upload-artifact@v4 + with: + name: site + path: site + + deploy-main: + needs: build + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@v4 + with: { name: site, path: site } + - name: Deploy to GitHub Pages (main) uses: peaceiris/actions-gh-pages@v4 with: github_token: ${{ secrets.GITHUB_TOKEN }} + publish_branch: gh-pages publish_dir: ./site - force_orphan: true \ No newline at end of file + force_orphan: true + + deploy-preview: + needs: build + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@v4 + with: { name: site, path: site } + - name: Deploy preview under subfolder + uses: peaceiris/actions-gh-pages@v4 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_branch: gh-pages + publish_dir: ./site + destination_dir: preview/${{ github.head_ref || github.ref_name }}/${{ needs.build.outputs.short_sha }} + keep_files: true # keep previous previews + # DO NOT set force_orphan here + - name: Print preview URL + run: | + echo "Preview: https://${{ github.repository_owner }}.github.io/$(basename ${{ github.repository }})/preview/${{ github.head_ref || github.ref_name }}/${{ needs.build.outputs.short_sha }}/" diff --git a/.gitignore b/.gitignore index 4a189d56..45729158 100644 --- a/.gitignore +++ b/.gitignore @@ -154,3 +154,9 @@ cython_debug/ *.code-workspace llm_rules.md .python-version +docs/site +docs/api/_build +docs/general/user-guide/data_csv +docs/general/user-guide/data_parquet +docs/api/reference/**/mesa_frames.*.rst +examples/**/results \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 147b84d3..bb8b4148 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -58,28 +58,13 @@ Before you begin contributing, ensure that you have the necessary tools installe #### **Step 3: Install Dependencies** ๐Ÿ“ฆ -It is recommended to set up a virtual environment before installing dependencies. +We manage the development environment with [uv](https://docs.astral.sh/uv/): -- **Using UV**: +```sh +uv sync --all-extras +``` - ```sh - uv add --dev .[dev] - ``` - -- **Using Hatch**: - - ```sh - hatch env create dev - ``` - -- **Using Standard Python**: - - ```sh - python3 -m venv myenv - source myenv/bin/activate # macOS/Linux - myenv\Scripts\activate # Windows - pip install -e ".[dev]" - ``` +This creates `.venv/` and installs mesa-frames with the development extras. #### **Step 4: Make and Commit Changes** โœจ @@ -99,21 +84,19 @@ It is recommended to set up a virtual environment before installing dependencies - **Run pre-commit hooks** to enforce code quality standards: ```sh - pre-commit run + uv run pre-commit run -a ``` - **Run tests** to ensure your contribution does not break functionality: ```sh - pytest --cov + uv run pytest -q --cov=mesa_frames --cov-report=term-missing ``` - - If using UV: `uv run pytest --cov` - - **Optional: Enable runtime type checking** during development for enhanced type safety: ```sh - MESA_FRAMES_RUNTIME_TYPECHECKING=1 uv run pytest --cov + MESA_FRAMES_RUNTIME_TYPECHECKING=1 uv run pytest -q --cov=mesa_frames --cov-report=term-missing ``` !!! tip "Automatically Enabled" @@ -135,8 +118,7 @@ It is recommended to set up a virtual environment before installing dependencies - Preview your changes by running: ```sh - mkdocs serve - uv run mkdocs serve #If using uv + uv run mkdocs serve ``` - Open `http://127.0.0.1:8000` in your browser to verify documentation updates. diff --git a/README.md b/README.md index 6a16baad..9ff6205a 100644 --- a/README.md +++ b/README.md @@ -1,173 +1,135 @@ -# mesa-frames ๐Ÿš€ + +

+ Mesa logo +

-mesa-frames is an extension of the [mesa](https://github.com/projectmesa/mesa) framework, designed for complex simulations with thousands of agents. By storing agents in a DataFrame, mesa-frames significantly enhances the performance and scalability of mesa, while maintaining a similar syntax. mesa-frames allows for the use of [vectorized functions](https://stackoverflow.com/a/1422198) which significantly speeds up operations whenever simultaneous activation of agents is possible. +

mesa-frames

+ -## Why DataFrames? ๐Ÿ“Š +| | | +| ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| CI/CD | [![CI Checks](https://github.com/projectmesa/mesa-frames/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/projectmesa/mesa-frames/actions/workflows/build.yml) [![codecov](https://codecov.io/gh/projectmesa/mesa-frames/branch/main/graph/badge.svg)](https://app.codecov.io/gh/projectmesa/mesa-frames) | +| Package | [![PyPI - Version](https://img.shields.io/pypi/v/mesa-frames.svg?logo=pypi&label=PyPI&logoColor=gold)](https://pypi.org/project/mesa-frames/) [![PyPI - Downloads](https://img.shields.io/pypi/dm/mesa-frames.svg?color=blue&label=Downloads&logo=pypi&logoColor=gold)](https://pypi.org/project/mesa-frames/) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mesa-frames.svg?logo=python&label=Python&logoColor=gold)](https://pypi.org/project/mesa-frames/) | +| Meta | [![linting - Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://docs.astral.sh/ruff/) [![formatter - Ruff](https://img.shields.io/badge/formatter-Ruff-0f172a?logo=ruff&logoColor=white)](https://docs.astral.sh/ruff/formatter/) [![Hatch project](https://img.shields.io/badge/%F0%9F%A5%9A-Hatch-4051b5.svg)](https://github.com/pypa/hatch) [![Managed with uv](https://img.shields.io/badge/managed%20with-uv-5a4fcf?logo=uv&logoColor=white)](https://github.com/astral-sh/uv) | +| Chat | [![chat](https://img.shields.io/matrix/project-mesa:matrix.org?label=chat&logo=Matrix)](https://matrix.to/#/#project-mesa:matrix.org) | -DataFrames are optimized for simultaneous operations through [SIMD processing](https://en.wikipedia.org/wiki/Single_instruction,_multiple_data). At the moment, mesa-frames supports the use of Polars library. +--- -- [Polars](https://pola.rs/) is a new DataFrame library with a syntax similar to pandas but with several innovations, including a backend implemented in Rust, the Apache Arrow memory format, query optimization, and support for larger-than-memory DataFrames. +## Scale Mesa beyond its limits -The following is a performance graph showing execution time using mesa and mesa-frames for the [Boltzmann Wealth model](https://mesa.readthedocs.io/en/stable/tutorials/intro_tutorial.html). +Classic [Mesa](https://github.com/projectmesa/mesa) stores each agent as a Python object, which quickly becomes a bottleneck at scale. +**mesa-frames** reimagines agent storage using **Polars DataFrames**, so agents live in a columnar store rather than the Python heap. -![Performance Graph with Mesa](https://github.com/projectmesa/mesa-frames/blob/main/examples/boltzmann_wealth/boltzmann_with_mesa.png) +You keep the Mesa-style `Model` / `AgentSet` structure, but updates are vectorized and memory-efficient. -![Performance Graph without Mesa](https://github.com/projectmesa/mesa-frames/blob/main/examples/boltzmann_wealth/boltzmann_no_mesa.png) +### Why it matters -([You can check the script used to generate the graph here](https://github.com/projectmesa/mesa-frames/blob/main/examples/boltzmann_wealth/performance_plot.py), but if you want to additionally compare vs Mesa, you have to uncomment `mesa_implementation` and its label) +- โšก **10ร— faster** bulk updates on 10k+ agents ([see Benchmarks](#benchmarks)) +- ๐Ÿ“Š **Columnar execution** via [Polars](https://docs.pola.rs/): [SIMD](https://en.wikipedia.org/wiki/Single_instruction,_multiple_data) ops, multi-core support +- ๐Ÿ”„ **Declarative logic**: agent rules as transformations, not Python loops +- ๐Ÿš€ **Roadmap**: Lazy queries and GPU support for even faster models -## Installation +--- -### Install from PyPI +## Who is it for? -```bash -pip install mesa-frames -``` +- Researchers needing to scale to **tens or hundreds of thousands of agents** +- Users whose agent logic can be written as **vectorized, set-based operations** -### Install from Source +โŒ **Not a good fit if:** your model depends on strict per-agent sequencing, complex non-vectorizable methods, or fine-grained identity tracking. -To install the most updated version of mesa-frames, you can clone the repository and install the package in editable mode. +--- -#### Cloning the Repository +## Why DataFrames? -To get started with mesa-frames, first clone the repository from GitHub: +DataFrames enable SIMD and columnar operations that are far more efficient than Python loops. +mesa-frames currently uses **Polars** as its backend. -```bash -git clone https://github.com/projectmesa/mesa-frames.git -cd mesa_frames -``` +| Feature | mesa (classic) | mesa-frames | +| ---------------------- | -------------- | ----------- | +| Storage | Python objects | Polars DataFrame | +| Updates | Loops | Vectorized ops | +| Memory overhead | High | Low | +| Max agents (practical) | ~10^3 | ~10^6+ | -#### Installing in a Conda Environment +--- -If you want to install it into a new environment: +## Benchmarks -```bash -conda create -n myenv -``` +[![Reproduce Benchmarks](https://img.shields.io/badge/Reproduce%20Benchmarks-๐Ÿ“Š-orange?style=for-the-badge)](https://projectmesa.github.io/mesa-frames/general/benchmarks/) -If you want to install it into an existing environment: +mesa-frames delivers consistent speedups across both toy and canonical ABMs. +At 10k agents, it runs **~10ร— faster** than classic Mesa, and the gap grows with scale. -```bash -conda activate myenv -``` +![Benchmark: Boltzmann Wealth](examples/boltzmann_wealth/boltzmann_benchmark.png) -Then, to install mesa-frames itself: +![Benchmark: Sugarscape IG](examples/sugarscape/sugarscape_benchmark.png) -```bash -pip install -e . -``` +--- -#### Installing in a Python Virtual Environment +## Quick Start -If you want to install it into a new environment: +[![Explore the Tutorials](https://img.shields.io/badge/Explore%20the%20Tutorials-๐Ÿ“š-blue?style=for-the-badge)](https://projectmesa.github.io/mesa-frames/general/user-guide/) -```bash -python3 -m venv myenv -source myenv/bin/activate # On Windows, use `myenv\Scripts\activate` -``` - -If you want to install it into an existing environment: +1. **Install** ```bash -source myenv/bin/activate # On Windows, use `myenv\Scripts\activate` + pip install mesa-frames ``` -Then, to install mesa-frames itself: +Or for development: ```bash -pip install -e . +git clone https://github.com/projectmesa/mesa-frames.git +cd mesa-frames +uv sync --all-extras ``` -## Usage +1. **Create a model** -[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/2_introductory-tutorial.ipynb) + ```python + from mesa_frames import AgentSet, Model + import polars as pl -**Note:** mesa-frames is currently in its early stages of development. As such, the usage patterns and API are subject to change. Breaking changes may be introduced. Reports of feedback and issues are encouraged. + class MoneyAgents(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) -[You can find the API documentation here](https://projectmesa.github.io/mesa-frames/api). + def give_money(self): + self.select(self.wealth > 0) + other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) + self["active", "wealth"] -= 1 + new_wealth = other_agents.group_by("unique_id").len() + self[new_wealth, "wealth"] += new_wealth["len"] -### Creation of an Agent + def step(self): + self.do("give_money") -The agent implementation differs from base mesa. Agents are only defined at the AgentSet level. You can import `AgentSet`. As in mesa, you subclass and make sure to call `super().__init__(model)`. You can use the `add` method or the `+=` operator to add agents to the AgentSet. Most methods mirror the functionality of `mesa.AgentSet`. Additionally, `mesa-frames.AgentSet` implements many dunder methods such as `AgentSet[mask, attr]` to get and set items intuitively. All operations are by default inplace, but if you'd like to use functional programming, mesa-frames implements a fast copy method which aims to reduce memory usage, relying on reference-only and native copy methods. + class MoneyModelDF(Model): + def __init__(self, N: int): + super().__init__() + self.sets += MoneyAgents(N, self) -```python -from mesa-frames import AgentSet + def step(self): + self.sets.do("step") + ``` -class MoneyAgents(AgentSet): - def __init__(self, n: int, model: Model): - super().__init__(model) - # Adding the agents to the agent set - self += pl.DataFrame( - {"wealth": pl.ones(n, eager=True)} - ) +--- - def step(self) -> None: - # The give_money method is called - self.do("give_money") +## Roadmap - def give_money(self): - # Active agents are changed to wealthy agents - self.select(self.wealth > 0) +> Community contributions welcome โ€” see the [full roadmap](https://projectmesa.github.io/mesa-frames/general/roadmap) - # Receiving agents are sampled (only native expressions currently supported) - other_agents = self.df.sample( - n=len(self.active_agents), with_replacement=True - ) +- Transition to LazyFrames for optimization and GPU support +- Auto-vectorize existing Mesa models via decorator +- Increase possible Spaces +- Refine the API to align to Mesa - # Wealth of wealthy is decreased by 1 - self["active", "wealth"] -= 1 - - # Compute the income of the other agents (only native expressions currently supported) - new_wealth = other_agents.group_by("unique_id").len() - - # Add the income to the other agents - self[new_wealth, "wealth"] += new_wealth["len"] -``` - -### Creation of the Model - -Creation of the model is fairly similar to the process in mesa. You subclass `Model` and call `super().__init__()`. The `model.sets` attribute has the same interface as `mesa-frames.AgentSet`. You can use `+=` or `self.sets.add` with a `mesa-frames.AgentSet` (or a list of `AgentSet`) to add agents to the model. - -```python -from mesa-frames import Model - -class MoneyModelDF(Model): - def __init__(self, N: int, agents_cls): - super().__init__() - self.n_agents = N - self.sets += MoneyAgents(N, self) - - def step(self): - # Executes the step method for every agentset in self.sets - self.sets.do("step") - - def run_model(self, n): - for _ in range(n): - self.step() -``` - -## What's Next? ๐Ÿ”ฎ - -- Refine the API to make it more understandable for someone who is already familiar with the mesa package. The goal is to provide a seamless experience for users transitioning to or incorporating mesa-frames. -- Adding support for default mesa functions to ensure that the standard mesa functionality is preserved. -- Adding GPU functionality (cuDF and Dask-cuDF). -- Creating a decorator that will automatically vectorize an existing mesa model. This feature will allow users to easily tap into the performance enhancements that mesa-frames offers without significant code alterations. -- Creating a unique class for AgentSet, independent of the backend implementation. +--- ## License -Copyright 2024 Adam Amer, Project Mesa team and contributors - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. +Copyright ยฉ 2025 Adam Amer, Project Mesa team and contributors -For the full license text, see the [LICENSE](https://github.com/projectmesa/mesa-frames/blob/main/LICENSE) file in the GitHub repository. +Licensed under the [Apache License, Version 2.0](https://raw.githubusercontent.com/projectmesa/mesa-frames/refs/heads/main/LICENSE). diff --git a/benchmarks/cli.py b/benchmarks/cli.py new file mode 100644 index 00000000..3accba5d --- /dev/null +++ b/benchmarks/cli.py @@ -0,0 +1,234 @@ +"""Typer CLI for running mesa vs mesa-frames performance benchmarks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from time import perf_counter +from typing import Literal, Annotated, Protocol, Optional + +import matplotlib.pyplot as plt +import polars as pl +import seaborn as sns +import typer + +from examples.boltzmann_wealth import backend_frames as boltzmann_frames +from examples.boltzmann_wealth import backend_mesa as boltzmann_mesa +from examples.sugarscape_ig.backend_frames import model as sugarscape_frames +from examples.sugarscape_ig.backend_mesa import model as sugarscape_mesa + +app = typer.Typer(add_completion=False) + + +class RunnerP(Protocol): + def __call__(self, agents: int, steps: int, seed: int | None = None) -> None: ... + + +@dataclass(slots=True) +class Backend: + name: Literal["mesa", "frames"] + runner: RunnerP + + +@dataclass(slots=True) +class ModelConfig: + name: str + backends: list[Backend] + + +MODELS: dict[str, ModelConfig] = { + "boltzmann": ModelConfig( + name="boltzmann", + backends=[ + Backend(name="mesa", runner=boltzmann_mesa.simulate), + Backend(name="frames", runner=boltzmann_frames.simulate), + ], + ), + "sugarscape": ModelConfig( + name="sugarscape", + backends=[ + Backend( + name="mesa", + runner=sugarscape_mesa.simulate, + ), + Backend( + name="frames", + runner=sugarscape_frames.simulate, + ), + ], + ), +} + + +def _parse_agents(value: str) -> list[int]: + value = value.strip() + if ":" in value: + parts = value.split(":") + if len(parts) != 3: + raise typer.BadParameter("Ranges must use start:stop:step format") + try: + start, stop, step = (int(part) for part in parts) + except ValueError as exc: + raise typer.BadParameter("Range values must be integers") from exc + if step <= 0: + raise typer.BadParameter("Step must be positive") + if start <= 0 or stop <= 0: + raise typer.BadParameter("Range endpoints must be positive") + if start > stop: + raise typer.BadParameter("Range start must be <= stop") + counts = list(range(start, stop + step, step)) + if counts[-1] > stop: + counts.pop() + return counts + try: + agents = int(value) + except ValueError as exc: # pragma: no cover - defensive + raise typer.BadParameter("Agent count must be an integer") from exc + if agents <= 0: + raise typer.BadParameter("Agent count must be positive") + return [agents] + + +def _parse_models(value: str) -> list[str]: + """Parse models option into a list of model keys. + + Accepts: + - "all" -> returns all available model keys + - a single model name -> returns [name] + - a comma-separated list of model names -> returns list + + Validates that each selected model exists in MODELS. + """ + value = value.strip() + if value == "all": + return list(MODELS.keys()) + # support comma-separated lists + parts = [part.strip() for part in value.split(",") if part.strip()] + if not parts: + raise typer.BadParameter("Model selection must not be empty") + unknown = [p for p in parts if p not in MODELS] + if unknown: + raise typer.BadParameter(f"Unknown model selection: {', '.join(unknown)}") + # preserve order and uniqueness + seen = set() + result: list[str] = [] + for p in parts: + if p not in seen: + seen.add(p) + result.append(p) + return result + + +def _plot_performance( + df: pl.DataFrame, model_name: str, output_dir: Path, timestamp: str +) -> None: + if df.is_empty(): + return + for theme, style in {"light": "whitegrid", "dark": "darkgrid"}.items(): + sns.set_theme(style=style) + fig, ax = plt.subplots(figsize=(8, 5)) + sns.lineplot( + data=df.to_pandas(), + x="agents", + y="runtime_seconds", + hue="backend", + estimator="mean", + errorbar="sd", + marker="o", + ax=ax, + ) + ax.set_title(f"{model_name.title()} runtime vs agents") + ax.set_xlabel("Agents") + ax.set_ylabel("Runtime (seconds)") + fig.tight_layout() + filename = output_dir / f"{model_name}_runtime_{timestamp}_{theme}.png" + fig.savefig(filename, dpi=300) + plt.close(fig) + + +@app.command() +def run( + models: Annotated[ + str, + typer.Option( + help="Models to benchmark: boltzmann, sugarscape, or all", + callback=_parse_models, + ), + ] = "all", + agents: Annotated[ + list[int], + typer.Option( + help="Agent count or range (start:stop:step)", callback=_parse_agents + ), + ] = "1000:5000:1000", + steps: Annotated[ + int, + typer.Option( + min=0, + help="Number of steps per run.", + ), + ] = 100, + repeats: Annotated[int, typer.Option(help="Repeats per configuration.", min=1)] = 1, + seed: Annotated[int, typer.Option(help="Optional RNG seed.")] = 42, + save: Annotated[bool, typer.Option(help="Persist benchmark CSV results.")] = True, + plot: Annotated[bool, typer.Option(help="Render performance plots.")] = True, + results_dir: Annotated[ + Path, + typer.Option( + help="Directory for benchmark CSV results.", + ), + ] = Path(__file__).resolve().parent / "results", + plots_dir: Annotated[ + Path, + typer.Option( + help="Directory for benchmark plots.", + ), + ] = Path(__file__).resolve().parent / "plots", +) -> None: + """Run performance benchmarks for the models models.""" + rows: list[dict[str, object]] = [] + timestamp = datetime.now(datetime.UTC).strftime("%Y%m%d_%H%M%S") + for model in models: + config = MODELS[model] + typer.echo(f"Benchmarking {model} with agents {agents}") + for agents_count in agents: + for repeat_idx in range(repeats): + run_seed = seed + repeat_idx + for backend in config.backends: + start = perf_counter() + backend.runner(agents_count, steps, run_seed) + runtime = perf_counter() - start + rows.append( + { + "model": model, + "backend": backend.name, + "agents": agents_count, + "steps": steps, + "seed": run_seed, + "repeat_idx": repeat_idx, + "runtime_seconds": runtime, + "timestamp": timestamp, + } + ) + if not rows: + typer.echo("No benchmark data collected.") + return + df = pl.DataFrame(rows) + if save: + results_dir.mkdir(parents=True, exist_ok=True) + for model in models: + model_df = df.filter(pl.col("model") == model) + csv_path = results_dir / f"{model}_perf_{timestamp}.csv" + model_df.write_csv(csv_path) + typer.echo(f"Saved {model} results to {csv_path}") + if plot: + plots_dir.mkdir(parents=True, exist_ok=True) + for model in models: + model_df = df.filter(pl.col("model") == model) + _plot_performance(model_df, model, plots_dir, timestamp) + typer.echo(f"Saved {model} plots under {plots_dir}") + + +if __name__ == "__main__": + app() diff --git a/docs/api/Makefile b/docs/api/Makefile deleted file mode 100644 index d0c3cbf1..00000000 --- a/docs/api/Makefile +++ /dev/null @@ -1,20 +0,0 @@ -# Minimal makefile for Sphinx documentation -# - -# You can set these variables from the command line, and also -# from the environment for the first two. -SPHINXOPTS ?= -SPHINXBUILD ?= sphinx-build -SOURCEDIR = source -BUILDDIR = build - -# Put it first so that "make" without argument is like "make help". -help: - @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) - -.PHONY: help Makefile - -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). -%: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/examples/sugarscape_ig/__init__.py b/docs/api/_static/brand-core.css similarity index 100% rename from examples/sugarscape_ig/__init__.py rename to docs/api/_static/brand-core.css diff --git a/examples/sugarscape_ig/ss_mesa/__init__.py b/docs/api/_static/brand-pydata.css similarity index 100% rename from examples/sugarscape_ig/ss_mesa/__init__.py rename to docs/api/_static/brand-pydata.css diff --git a/docs/api/conf.py b/docs/api/conf.py index 43098ec2..745305dc 100644 --- a/docs/api/conf.py +++ b/docs/api/conf.py @@ -31,9 +31,26 @@ exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] # -- Options for HTML output ------------------------------------------------- +# Hide objects (classes/methods) from the page Table of Contents +toc_object_entries = False # NEW: stop adding class/method entries to the TOC + + html_theme = "pydata_sphinx_theme" html_static_path = ["_static"] html_show_sourcelink = False +html_logo = ( + "https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.png" +) +html_favicon = ( + "https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.ico" +) + +# Add custom branding CSS/JS (mesa_brand) to static files +html_css_files = [ + # Shared brand variables then theme adapter for pydata + "brand-core.css", + "brand-pydata.css", +] # -- Extension settings ------------------------------------------------------ # intersphinx mapping @@ -52,9 +69,18 @@ copybutton_prompt_is_regexp = True # -- Custom configurations --------------------------------------------------- +add_module_names = False autoclass_content = "class" autodoc_member_order = "bysource" -autodoc_default_options = {"special-members": True, "exclude-members": "__weakref__"} +autodoc_default_options = { + "members": True, + "inherited-members": True, + "undoc-members": True, + "member-order": "bysource", + "special-members": True, + "exclude-members": "__weakref__,__dict__,__module__,__annotations__,__firstlineno__,__static_attributes__,__abstractmethods__,__slots__", +} + # -- GitHub link and user guide settings ------------------------------------- github_root = "https://github.com/projectmesa/mesa-frames" @@ -64,7 +90,7 @@ "external_links": [ { "name": "User guide", - "url": f"{web_root}/user-guide/", + "url": f"{web_root}/user-guide/0_getting-started/", }, ], "icon_links": [ @@ -73,6 +99,12 @@ "url": github_root, "icon": "fa-brands fa-github", }, + { + "name": "Matrix", + "url": "https://matrix.to/#/#project-mesa:matrix.org", + "icon": "fa-solid fa-comments", + }, ], - "navbar_end": ["navbar-icon-links"], + "navbar_start": ["navbar-logo"], + "navbar_end": ["theme-switcher", "navbar-icon-links"], } diff --git a/docs/api/index.rst b/docs/api/index.rst index 936350d6..a7c2ab4c 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -1,34 +1,55 @@ mesa-frames API =============== -This page provides a high-level overview of all public mesa-frames objects, functions, and methods. All classes and functions exposed in the ``mesa_frames.*`` namespace are public. +.. toctree:: + :caption: Shortcuts + :maxdepth: 1 + :hidden: + + reference/agents/index + reference/model + reference/space/index + reference/datacollector + -.. grid:: - .. grid-item-card:: +Overview +-------- - .. toctree:: - :maxdepth: 2 +mesa-frames provides a DataFrame-first API for agent-based models. Instead of representing each agent as a distinct Python object, agents are stored in AgentSets (backed by DataFrames) and manipulated via vectorised operations. This leads to much lower memory overhead and faster bulk updates while keeping an object-oriented feel for model structure and lifecycle management. - reference/agents/index - .. grid-item-card:: +Mini usage flow +--------------- + +1. Create a Model and register AgentSets on ``model.sets``. +2. Populate AgentSets with agents (rows) and attributes (columns) via adding a DataFrame to the AgentSet. +3. Implement AgentSet methods that operate on DataFrames +4. Use ``model.sets.do("step")`` from the model loop to advance the simulation; datacollectors and reporters can sample model- and agent-level columns at each step. + +.. grid:: + :gutter: 2 - .. toctree:: - :maxdepth: 2 + .. grid-item-card:: Manage agent collections + :link: reference/agents/index + :link-type: doc - reference/model + Create and operate on ``AgentSets`` and ``AgentSetRegisties``: add/remove agents. - .. grid-item-card:: + .. grid-item-card:: Model orchestration + :link: reference/model + :link-type: doc - .. toctree:: - :maxdepth: 2 + ``Model`` API for registering sets, stepping the simulation, and integrating with datacollectors/reporters. - reference/space/index + .. grid-item-card:: Spatial support + :link: reference/space/index + :link-type: doc - .. grid-item-card:: + Placement and neighbourhood utilities for ``Grid`` and space - .. toctree:: - :maxdepth: 2 + .. grid-item-card:: Collect simulation data + :link: reference/datacollector + :link-type: doc - reference/datacollector \ No newline at end of file + Record model- and agent-level metrics over time with ``DataCollector``. Sample columns, run aggregations, and export cleaned frames for analysis. \ No newline at end of file diff --git a/docs/api/make.bat b/docs/api/make.bat deleted file mode 100644 index dc1312ab..00000000 --- a/docs/api/make.bat +++ /dev/null @@ -1,35 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -) -set SOURCEDIR=source -set BUILDDIR=build - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.https://www.sphinx-doc.org/ - exit /b 1 -) - -if "%1" == "" goto help - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/docs/api/reference/agents/index.rst b/docs/api/reference/agents/index.rst index a1c03126..cc5dfef0 100644 --- a/docs/api/reference/agents/index.rst +++ b/docs/api/reference/agents/index.rst @@ -3,15 +3,184 @@ Agents .. currentmodule:: mesa_frames +Quick intro +----------- -.. autoclass:: AgentSet - :members: - :inherited-members: - :autosummary: - :autosummary-nosignatures: - -.. autoclass:: AgentSetRegistry - :members: - :inherited-members: - :autosummary: - :autosummary-nosignatures: +- ``AgentSet`` stores agents as rows in a Polars-backed table and provides vectorised operations for high-performance updates. + +- ``AgentSetRegistry`` (available at ``model.sets``) is the container that holds all ``AgentSet`` instances for a model and provides convenience operations (add/remove sets, step all sets, rename). + +- Keep agent logic column-oriented and prefer Polars expressions for updates. + +Minimal example +--------------- + +.. code-block:: python + + from mesa_frames import Model, AgentSet + import polars as pl + + class MySet(AgentSet): + def __init__(self, model): + super().__init__(model) + self.add(pl.DataFrame({"age": [0, 5, 10]})) + + def step(self): + # vectorised update: increase age for all agents + self.df = self.df.with_columns((pl.col("age") + 1).alias("age")) + + class MyModel(Model): + def __init__(self): + super().__init__() + # register an AgentSet on the model's registry + self.sets += MySet(self) + + m = MyModel() + # step all registered sets (delegates to each AgentSet.step) + m.sets.do("step") + +API reference +--------------------------------- + +.. tab-set:: + + .. tab-item:: AgentSet + + .. tab-set:: + + .. tab-item:: Overview + + .. rubric:: Lifecycle / Core + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSet.__init__ + AgentSet.step + AgentSet.rename + AgentSet.copy + + .. rubric:: Accessors & Views + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSet.df + AgentSet.active_agents + AgentSet.inactive_agents + AgentSet.index + AgentSet.pos + AgentSet.name + AgentSet.get + AgentSet.contains + AgentSet.__len__ + AgentSet.__iter__ + AgentSet.__getitem__ + AgentSet.__contains__ + + .. rubric:: Mutators + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSet.add + AgentSet.remove + AgentSet.discard + AgentSet.set + AgentSet.select + AgentSet.shuffle + AgentSet.sort + AgentSet.do + + .. rubric:: Operators / Internal helpers + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSet.__add__ + AgentSet.__iadd__ + AgentSet.__sub__ + AgentSet.__isub__ + AgentSet.__repr__ + AgentSet.__reversed__ + + .. tab-item:: Full API + + .. autoclass:: AgentSet + :autosummary: + :autosummary-nosignatures: + + .. tab-item:: AgentSetRegistry + + .. tab-set:: + + .. tab-item:: Overview + + .. rubric:: Lifecycle / Core + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSetRegistry.__init__ + AgentSetRegistry.copy + AgentSetRegistry.rename + + .. rubric:: Accessors & Queries + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSetRegistry.get + AgentSetRegistry.contains + AgentSetRegistry.ids + AgentSetRegistry.keys + AgentSetRegistry.items + AgentSetRegistry.values + AgentSetRegistry.model + AgentSetRegistry.random + AgentSetRegistry.space + AgentSetRegistry.__len__ + AgentSetRegistry.__iter__ + AgentSetRegistry.__getitem__ + AgentSetRegistry.__contains__ + + .. rubric:: Mutators / Coordination + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSetRegistry.add + AgentSetRegistry.remove + AgentSetRegistry.discard + AgentSetRegistry.replace + AgentSetRegistry.shuffle + AgentSetRegistry.sort + AgentSetRegistry.do + AgentSetRegistry.__setitem__ + AgentSetRegistry.__add__ + AgentSetRegistry.__iadd__ + AgentSetRegistry.__sub__ + AgentSetRegistry.__isub__ + + .. rubric:: Representation + + .. autosummary:: + :nosignatures: + :toctree: + + AgentSetRegistry.__repr__ + AgentSetRegistry.__str__ + AgentSetRegistry.__reversed__ + + .. tab-item:: Full API + + .. autoclass:: AgentSetRegistry + :autosummary: + :autosummary-nosignatures: \ No newline at end of file diff --git a/docs/api/reference/datacollector.rst b/docs/api/reference/datacollector.rst index bdf38cfd..f1f2c68e 100644 --- a/docs/api/reference/datacollector.rst +++ b/docs/api/reference/datacollector.rst @@ -3,8 +3,66 @@ Data Collection .. currentmodule:: mesa_frames -.. autoclass:: DataCollector - :members: - :inherited-members: - :autosummary: - :autosummary-nosignatures: \ No newline at end of file +Quick intro +----------- + +``DataCollector`` samples model- and agent-level columns over time and returns cleaned DataFrames suitable for analysis. Typical patterns: + +- Provide ``model_reporters`` (callables producing scalars) and ``agent_reporters`` (column selectors or callables that operate on an AgentSet). +- Call ``collector.collect(model)`` inside the model step or use built-in integration if the model calls the collector automatically. + +Minimal example +--------------- + +.. code-block:: python + + from mesa_frames import DataCollector, Model, AgentSet + import polars as pl + + class P(AgentSet): + def __init__(self, model): + super().__init__(model) + self.add(pl.DataFrame({'x': [1,2]})) + + class M(Model): + def __init__(self): + super().__init__() + self.sets += P(self) + self.dc = DataCollector(model_reporters={'count': lambda m: len(m.sets['P'])}, + agent_reporters='x') + + m = M() + m.dc.collect() + +API reference +------------- + +.. tab-set:: + + .. tab-item:: Overview + + .. rubric:: Lifecycle / Core + + .. autosummary:: + :nosignatures: + :toctree: + + DataCollector.__init__ + DataCollector.collect + DataCollector.conditional_collect + DataCollector.flush + DataCollector.data + + .. rubric:: Reporting / Internals + + .. autosummary:: + :nosignatures: + :toctree: + + DataCollector.seed + + .. tab-item:: Full API + + .. autoclass:: DataCollector + :autosummary: + :autosummary-nosignatures: \ No newline at end of file diff --git a/docs/api/reference/model.rst b/docs/api/reference/model.rst index 099e601b..74b7e4e5 100644 --- a/docs/api/reference/model.rst +++ b/docs/api/reference/model.rst @@ -3,8 +3,67 @@ Model .. currentmodule:: mesa_frames -.. autoclass:: Model - :members: - :inherited-members: - :autosummary: - :autosummary-nosignatures: \ No newline at end of file +Quick intro +----------- + +``Model`` orchestrates the simulation lifecycle: creating and registering ``AgentSet``s, stepping the simulation, and integrating with ``DataCollector`` and spatial ``Grid``s. Typical usage: + +- Instantiate ``Model``, add ``AgentSet`` instances to ``model.sets``. +- Call ``model.sets.do('step')`` inside your model loop to trigger set-level updates. +- Use ``DataCollector`` to sample model- and agent-level columns each step. + +Minimal example +--------------- + +.. code-block:: python + + from mesa_frames import Model, AgentSet, DataCollector + import polars as pl + + class People(AgentSet): + def step(self): + self.add(pl.DataFrame({'wealth': [1, 2, 3]})) + + class MyModel(Model): + def __init__(self): + super().__init__() + self.sets += People(self) + self.dc = DataCollector(model_reporters={'avg_wealth': lambda m: m.sets['People'].df['wealth'].mean()}) + + m = MyModel() + m.step() + +API reference +------------- + +.. tab-set:: + + .. tab-item:: Overview + + .. rubric:: Lifecycle / Core + + .. autosummary:: + :nosignatures: + :toctree: + + Model.__init__ + Model.step + Model.run_model + Model.reset_randomizer + + .. rubric:: Accessors / Properties + + .. autosummary:: + :nosignatures: + :toctree: + + Model.steps + Model.sets + Model.space + Model.seed + + .. tab-item:: Full API + + .. autoclass:: Model + :autosummary: + :autosummary-nosignatures: \ No newline at end of file diff --git a/docs/api/reference/space/index.rst b/docs/api/reference/space/index.rst index 8741b6b6..c11b140d 100644 --- a/docs/api/reference/space/index.rst +++ b/docs/api/reference/space/index.rst @@ -4,8 +4,63 @@ This page provides a high-level overview of possible space objects for mesa-fram .. currentmodule:: mesa_frames -.. autoclass:: Grid - :members: - :inherited-members: - :autosummary: - :autosummary-nosignatures: \ No newline at end of file +Quick intro +----------- + + + +Currently we only support the ``Grid``. Typical usage: + +- Construct ``Grid(model, (width, height))`` and use ``place``/ ``move`` helpers to update agent positional columns. +- Use neighbourhood queries to produce masks or index lists and then apply vectorised updates to selected rows. + +Minimal example +--------------- + +.. code-block:: python + + from mesa_frames import Model, Grid, AgentSet + import polars as pl + + class P(AgentSet): + pass + + class M(Model): + def __init__(self): + super().__init__() + self.space = Grid(self, (10, 10)) + self.sets += P(self) + self.space.place_to_empty(self.sets) + + m = M() + m.space.move_to_available(m.sets) + + +API reference +------------- + +.. tab-set:: + + .. tab-item:: Overview + + .. rubric:: Lifecycle / Core + + .. autosummary:: + :nosignatures: + :toctree: + + Grid.__init__ + + .. rubric:: Sampling & Queries + + .. autosummary:: + :nosignatures: + :toctree: + + Grid.remaining_capacity + + .. tab-item:: Full API + + .. autoclass:: Grid + :autosummary: + :autosummary-nosignatures: \ No newline at end of file diff --git a/docs/general/index.md b/docs/general/index.md index 9859d2ee..cee3f109 100644 --- a/docs/general/index.md +++ b/docs/general/index.md @@ -1,89 +1 @@ -# Welcome to mesa-frames ๐Ÿš€ - -mesa-frames is an extension of the [mesa](https://github.com/projectmesa/mesa) framework, designed for complex simulations with thousands of agents. By storing agents in a DataFrame, mesa-frames significantly enhances the performance and scalability of mesa, while maintaining a similar syntax. - -You can get a model which is multiple orders of magnitude faster based on the number of agents - the more agents, the faster the relative performance. - -## Why DataFrames? ๐Ÿ“Š - -DataFrames are optimized for simultaneous operations through [SIMD processing](https://en.wikipedia.org/wiki/Single_instruction,_multiple_data). Currently, mesa-frames supports the library: - -- [Polars](https://pola.rs/): A new DataFrame library with a Rust backend, offering innovations like Apache Arrow memory format and support for larger-than-memory DataFrames. - -## Performance Boost ๐ŸŽ๏ธ - -Check out our performance graphs comparing mesa and mesa-frames for the [Boltzmann Wealth model](https://mesa.readthedocs.io/en/stable/tutorials/intro_tutorial.html): - -![Performance Graph with Mesa](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_with_mesa.png) - -![Performance Graph without Mesa](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_no_mesa.png) - -## Quick Start ๐Ÿš€ - -### Installation - -#### Installing from PyPI - -```bash -pip install mesa-frames -``` - -#### Installing from Source - -```bash -git clone https://github.com/projectmesa/mesa-frames.git -cd mesa_frames -pip install -e . -``` - -### Basic Usage - -Here's a quick example of how to create a model using mesa-frames: - -```python -from mesa_frames import AgentSet, Model -import polars as pl - -class MoneyAgents(AgentSet): - def __init__(self, n: int, model: Model): - super().__init__(model) - self += pl.DataFrame( - {"wealth": pl.ones(n, eager=True)} - ) - - def step(self) -> None: - self.do("give_money") - - def give_money(self): - # ... (implementation details) - -class MoneyModel(Model): - def __init__(self, N: int): - super().__init__() - self.sets += MoneyAgents(N, self) - - def step(self): - self.sets.do("step") - - def run_model(self, n): - for _ in range(n): - self.step() -``` - -## What's Next? ๐Ÿ”ฎ - -- API refinement for seamless transition from mesa -- Support for mesa functions -- Multiple other spaces: GeoGrid, ContinuousSpace, Network... -- Additional backends: Dask, cuDF (GPU), Dask-cuDF (GPU)... -- More examples: Schelling model, ... -- Automatic vectorization of existing mesa models -- Backend-agnostic AgentSet class - -## Get Involved! ๐Ÿค - -mesa-frames is in its early stages, and we welcome your feedback and contributions! Check out our [GitHub repository](https://github.com/projectmesa/mesa-frames) to get started. - -## License - -mesa-frames is available under the MIT License. See the [LICENSE](https://github.com/projectmesa/mesa-frames/blob/main/LICENSE) file for full details. +{% include-markdown "../../README.md" %} diff --git a/docs/general/user-guide/0_getting-started.md b/docs/general/user-guide/0_getting-started.md index 1edc1587..51ebe319 100644 --- a/docs/general/user-guide/0_getting-started.md +++ b/docs/general/user-guide/0_getting-started.md @@ -21,7 +21,7 @@ mesa-frames leverages the power of vectorized operations provided by DataFrame l - This approach is significantly faster than iterating over individual agents - Complex behaviors can be expressed in fewer lines of code -You should never use loops to iterate through your agents. Instead, use vectorized operations and implemented methods. If you need to loop, loop through vectorized operations (see the advanced tutorial SugarScape IG for more information). +Default to vectorized operations when expressing agent behaviour; that's where mesa-frames gains most of its speed-ups. If your agents must act sequentially (for example, to resolve conflicts or enforce ordering), fall back to loops or staged vectorized passesโ€”mesa-frames will behave more like base mesa in those situations. We'll unpack these trade-offs in the SugarScape advanced tutorial. It's important to note that in traditional `mesa` models, the order in which agents are activated can significantly impact the results of the model (see [Comer, 2014](http://mars.gmu.edu/bitstream/handle/1920/9070/Comer_gmu_0883E_10539.pdf)). `mesa-frames`, by default, doesn't have this issue as all agents are processed simultaneously. However, this comes with the trade-off of needing to carefully implement conflict resolution mechanisms when sequential processing is required. We'll discuss how to handle these situations later in this guide. @@ -42,7 +42,7 @@ Here's a comparison between mesa-frames and mesa: self.select(self.wealth > 0) # Receiving agents are sampled (only native expressions currently supported) - other_agents = self.sets.sample( + other_agents = self.model.sets.sample( n=len(self.active_agents), with_replacement=True ) @@ -92,7 +92,7 @@ If you're familiar with mesa, this guide will help you understand the key differ }) def step(self): givers = self.wealth > 0 - receivers = self.sets.sample(n=len(self.active_agents)) + receivers = self.model.sets.sample(n=len(self.active_agents)) self[givers, "wealth"] -= 1 new_wealth = receivers.groupby("unique_id").count() self[new_wealth["unique_id"], "wealth"] += new_wealth["count"] diff --git a/docs/general/user-guide/2_introductory-tutorial.ipynb b/docs/general/user-guide/2_introductory-tutorial.ipynb deleted file mode 100644 index 11391f9d..00000000 --- a/docs/general/user-guide/2_introductory-tutorial.ipynb +++ /dev/null @@ -1,449 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7ee055b2", - "metadata": {}, - "source": [ - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/2_introductory-tutorial.ipynb)" - ] - }, - { - "cell_type": "markdown", - "id": "8bd0381e", - "metadata": {}, - "source": [ - "## Installation (if running in Colab)\n", - "\n", - "Run the following cell to install `mesa-frames` if you are using Google Colab." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "df4d8623", - "metadata": {}, - "outputs": [], - "source": [ - "# !pip install git+https://github.com/projectmesa/mesa-frames mesa" - ] - }, - { - "cell_type": "markdown", - "id": "11515dfc", - "metadata": {}, - "source": [ - " # Introductory Tutorial: Boltzmann Wealth Model with mesa-frames ๐Ÿ’ฐ๐Ÿš€\n", - "\n", - "In this tutorial, we'll implement the Boltzmann Wealth Model using mesa-frames. This model simulates the distribution of wealth among agents, where agents randomly give money to each other.\n", - "\n", - "## Setting Up the Model ๐Ÿ—๏ธ\n", - "\n", - "First, let's import the necessary modules and set up our model class:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fc0ee981", - "metadata": {}, - "outputs": [ - { - "ename": "ImportError", - "evalue": "cannot import name 'Model' from partially initialized module 'mesa_frames' (most likely due to a circular import) (/home/adam/projects/mesa-frames/mesa_frames/__init__.py)", - "output_type": "error", - "traceback": [ - "\u001b[31m---------------------------------------------------------------------------\u001b[39m", - "\u001b[31mImportError\u001b[39m Traceback (most recent call last)", - "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Model, AgentSet, DataCollector\n\u001b[32m 4\u001b[39m \u001b[38;5;28;01mclass\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mMoneyModelDF\u001b[39;00m(Model):\n\u001b[32m 5\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, N: \u001b[38;5;28mint\u001b[39m, agents_cls):\n", - "\u001b[36mFile \u001b[39m\u001b[32m~/projects/mesa-frames/mesa_frames/__init__.py:65\u001b[39m\n\u001b[32m 63\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mconcrete\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01magentset\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m AgentSet\n\u001b[32m 64\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mconcrete\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01magentsetregistry\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m AgentSetRegistry\n\u001b[32m---> \u001b[39m\u001b[32m65\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mconcrete\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mdatacollector\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m DataCollector\n\u001b[32m 66\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mconcrete\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mmodel\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Model\n\u001b[32m 67\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mconcrete\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mspace\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Grid\n", - "\u001b[36mFile \u001b[39m\u001b[32m~/projects/mesa-frames/mesa_frames/concrete/datacollector.py:62\u001b[39m\n\u001b[32m 60\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtempfile\u001b[39;00m\n\u001b[32m 61\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mpsycopg2\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m62\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mabstract\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mdatacollector\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m AbstractDataCollector\n\u001b[32m 63\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtyping\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Any, Literal\n\u001b[32m 64\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mcollections\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mabc\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Callable\n", - "\u001b[36mFile \u001b[39m\u001b[32m~/projects/mesa-frames/mesa_frames/abstract/datacollector.py:50\u001b[39m\n\u001b[32m 48\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtyping\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Any, Literal\n\u001b[32m 49\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mcollections\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mabc\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Callable\n\u001b[32m---> \u001b[39m\u001b[32m50\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mmesa_frames\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Model\n\u001b[32m 51\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mpolars\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mpl\u001b[39;00m\n\u001b[32m 52\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mthreading\u001b[39;00m\n", - "\u001b[31mImportError\u001b[39m: cannot import name 'Model' from partially initialized module 'mesa_frames' (most likely due to a circular import) (/home/adam/projects/mesa-frames/mesa_frames/__init__.py)" - ] - } - ], - "source": [ - "from mesa_frames import Model, AgentSet, DataCollector\n", - "\n", - "\n", - "class MoneyModel(Model):\n", - " def __init__(self, N: int, agents_cls):\n", - " super().__init__()\n", - " self.n_agents = N\n", - " self.sets += agents_cls(N, self)\n", - " self.datacollector = DataCollector(\n", - " model=self,\n", - " model_reporters={\n", - " \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum()\n", - " },\n", - " agent_reporters={\"wealth\": \"wealth\"},\n", - " storage=\"csv\",\n", - " storage_uri=\"./data\",\n", - " trigger=lambda m: m.schedule.steps % 2 == 0,\n", - " )\n", - "\n", - " def step(self):\n", - " # Executes the step method for every agentset in self.sets\n", - " self.sets.do(\"step\")\n", - "\n", - " def run_model(self, n):\n", - " for _ in range(n):\n", - " self.step()\n", - " self.datacollector.conditional_collect\n", - " self.datacollector.flush()" - ] - }, - { - "cell_type": "markdown", - "id": "00e092c4", - "metadata": {}, - "source": [ - "## Implementing the AgentSet ๐Ÿ‘ฅ\n", - "\n", - "Now, let's implement our `MoneyAgents` using polars backends." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2bac0126", - "metadata": {}, - "outputs": [], - "source": [ - "import polars as pl\n", - "\n", - "\n", - "class MoneyAgents(AgentSet):\n", - " def __init__(self, n: int, model: Model):\n", - " super().__init__(model)\n", - " self += pl.DataFrame({\"wealth\": pl.ones(n, eager=True)})\n", - "\n", - " def step(self) -> None:\n", - " self.do(\"give_money\")\n", - "\n", - " def give_money(self):\n", - " self.select(self.wealth > 0)\n", - " other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True)\n", - " self[\"active\", \"wealth\"] -= 1\n", - " new_wealth = other_agents.group_by(\"unique_id\").len()\n", - " self[new_wealth[\"unique_id\"], \"wealth\"] += new_wealth[\"len\"]" - ] - }, - { - "cell_type": "markdown", - "id": "3b141016", - "metadata": {}, - "source": [ - "\n", - "## Running the Model โ–ถ๏ธ\n", - "\n", - "Now that we have our model and agent set defined, let's run a simulation:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "65da4e6f", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "shape: (9, 2)\n", - "โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”\n", - "โ”‚ statistic โ”† wealth โ”‚\n", - "โ”‚ --- โ”† --- โ”‚\n", - "โ”‚ str โ”† f64 โ”‚\n", - "โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก\n", - "โ”‚ count โ”† 1000.0 โ”‚\n", - "โ”‚ null_count โ”† 0.0 โ”‚\n", - "โ”‚ mean โ”† 1.0 โ”‚\n", - "โ”‚ std โ”† 1.134587 โ”‚\n", - "โ”‚ min โ”† 0.0 โ”‚\n", - "โ”‚ 25% โ”† 0.0 โ”‚\n", - "โ”‚ 50% โ”† 1.0 โ”‚\n", - "โ”‚ 75% โ”† 2.0 โ”‚\n", - "โ”‚ max โ”† 8.0 โ”‚\n", - "โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜\n" - ] - } - ], - "source": [ - "# Create and run the model\n", - "model = MoneyModel(1000, MoneyAgents)\n", - "model.run_model(100)\n", - "\n", - "wealth_dist = list(model.sets.df.values())[0]\n", - "\n", - "# Print the final wealth distribution\n", - "print(wealth_dist.select(pl.col(\"wealth\")).describe())" - ] - }, - { - "cell_type": "markdown", - "id": "812da73b", - "metadata": {}, - "source": [ - "\n", - "This output shows the statistical summary of the wealth distribution after 100 steps of the simulation with 1000 agents.\n", - "\n", - "## Performance Comparison ๐ŸŽ๏ธ๐Ÿ’จ\n", - "\n", - "One of the key advantages of mesa-frames is its performance with large numbers of agents. Let's compare the performance of mesa and polars:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fbdb540810924de8", - "metadata": {}, - "outputs": [], - "source": [ - "class MoneyAgentsConcise(AgentSet):\n", - " def __init__(self, n: int, model: Model):\n", - " super().__init__(model)\n", - " ## Adding the agents to the agent set\n", - " # 1. Changing the df attribute directly (not recommended, if other agents were added before, they will be lost)\n", - " \"\"\"self.df = pl.DataFrame(\n", - " {\"wealth\": pl.ones(n, eager=True)}\n", - " )\"\"\"\n", - " # 2. Adding the dataframe with add\n", - " \"\"\"self.add(\n", - " pl.DataFrame(\n", - " {\n", - " \"wealth\": pl.ones(n, eager=True),\n", - " }\n", - " )\n", - " )\"\"\"\n", - " # 3. Adding the dataframe with __iadd__\n", - " self += pl.DataFrame({\"wealth\": pl.ones(n, eager=True)})\n", - "\n", - " def step(self) -> None:\n", - " # The give_money method is called\n", - " # self.give_money()\n", - " self.do(\"give_money\")\n", - "\n", - " def give_money(self):\n", - " ## Active agents are changed to wealthy agents\n", - " # 1. Using the __getitem__ method\n", - " # self.select(self[\"wealth\"] > 0)\n", - " # 2. Using the fallback __getattr__ method\n", - " self.select(self.wealth > 0)\n", - "\n", - " # Receiving agents are sampled (only native expressions currently supported)\n", - " other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True)\n", - "\n", - " # Wealth of wealthy is decreased by 1\n", - " # 1. Using the __setitem__ method with self.active_agents mask\n", - " # self[self.active_agents, \"wealth\"] -= 1\n", - " # 2. Using the __setitem__ method with \"active\" mask\n", - " self[\"active\", \"wealth\"] -= 1\n", - "\n", - " # Compute the income of the other agents (only native expressions currently supported)\n", - " new_wealth = other_agents.group_by(\"unique_id\").len()\n", - "\n", - " # Add the income to the other agents\n", - " # 1. Using the set method\n", - " \"\"\"self.set(\n", - " attr_names=\"wealth\",\n", - " values=pl.col(\"wealth\") + new_wealth[\"len\"],\n", - " mask=new_wealth,\n", - " )\"\"\"\n", - "\n", - " # 2. Using the __setitem__ method\n", - " self[new_wealth, \"wealth\"] += new_wealth[\"len\"]\n", - "\n", - "\n", - "class MoneyAgentsNative(AgentSet):\n", - " def __init__(self, n: int, model: Model):\n", - " super().__init__(model)\n", - " self += pl.DataFrame({\"wealth\": pl.ones(n, eager=True)})\n", - "\n", - " def step(self) -> None:\n", - " self.do(\"give_money\")\n", - "\n", - " def give_money(self):\n", - " ## Active agents are changed to wealthy agents\n", - " self.select(pl.col(\"wealth\") > 0)\n", - "\n", - " other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True)\n", - "\n", - " # Wealth of wealthy is decreased by 1\n", - " self.df = self.df.with_columns(\n", - " wealth=pl.when(\n", - " pl.col(\"unique_id\").is_in(self.active_agents[\"unique_id\"].implode())\n", - " )\n", - " .then(pl.col(\"wealth\") - 1)\n", - " .otherwise(pl.col(\"wealth\"))\n", - " )\n", - "\n", - " new_wealth = other_agents.group_by(\"unique_id\").len()\n", - "\n", - " # Add the income to the other agents\n", - " self.df = (\n", - " self.df.join(new_wealth, on=\"unique_id\", how=\"left\")\n", - " .fill_null(0)\n", - " .with_columns(wealth=pl.col(\"wealth\") + pl.col(\"len\"))\n", - " .drop(\"len\")\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "496196d999f18634", - "metadata": {}, - "source": [ - "Add Mesa implementation of MoneyAgent and MoneyModel classes to test Mesa performance" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9dbe761af964af5b", - "metadata": {}, - "outputs": [], - "source": [ - "import mesa\n", - "\n", - "\n", - "class MesaMoneyAgent(mesa.Agent):\n", - " \"\"\"An agent with fixed initial wealth.\"\"\"\n", - "\n", - " def __init__(self, model):\n", - " # Pass the parameters to the parent class.\n", - " super().__init__(model)\n", - "\n", - " # Create the agent's variable and set the initial values.\n", - " self.wealth = 1\n", - "\n", - " def step(self):\n", - " # Verify agent has some wealth\n", - " if self.wealth > 0:\n", - " other_agent: MesaMoneyAgent = self.model.random.choice(self.model.agents)\n", - " if other_agent is not None:\n", - " other_agent.wealth += 1\n", - " self.wealth -= 1\n", - "\n", - "\n", - "class MesaMoneyModel(mesa.Model):\n", - " \"\"\"A model with some number of agents.\"\"\"\n", - "\n", - " def __init__(self, N: int):\n", - " super().__init__()\n", - " self.num_agents = N\n", - " for _ in range(N):\n", - " self.agents.add(MesaMoneyAgent(self))\n", - "\n", - " def step(self):\n", - " \"\"\"Advance the model by one step.\"\"\"\n", - " self.agents.shuffle_do(\"step\")\n", - "\n", - " def run_model(self, n_steps) -> None:\n", - " for _ in range(n_steps):\n", - " self.step()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2d864cd3", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Execution times:\n", - "---------------\n", - "mesa:\n", - " Number of agents: 100, Time: 0.03 seconds\n", - " Number of agents: 1001, Time: 1.45 seconds\n", - " Number of agents: 2000, Time: 5.40 seconds\n", - "---------------\n", - "---------------\n", - "mesa-frames (pl concise):\n", - " Number of agents: 100, Time: 1.60 seconds\n", - " Number of agents: 1001, Time: 2.68 seconds\n", - " Number of agents: 2000, Time: 3.04 seconds\n", - "---------------\n", - "---------------\n", - "mesa-frames (pl native):\n", - " Number of agents: 100, Time: 0.62 seconds\n", - " Number of agents: 1001, Time: 0.80 seconds\n", - " Number of agents: 2000, Time: 1.10 seconds\n", - "---------------\n" - ] - } - ], - "source": [ - "import time\n", - "\n", - "\n", - "def run_simulation(model: MesaMoneyModel | MoneyModel, n_steps: int):\n", - " start_time = time.time()\n", - " model.run_model(n_steps)\n", - " end_time = time.time()\n", - " return end_time - start_time\n", - "\n", - "\n", - "# Compare mesa and mesa-frames implementations\n", - "n_agents_list = [10**2, 10**3 + 1, 2 * 10**3]\n", - "n_steps = 100\n", - "print(\"Execution times:\")\n", - "for implementation in [\n", - " \"mesa\",\n", - " \"mesa-frames (pl concise)\",\n", - " \"mesa-frames (pl native)\",\n", - "]:\n", - " print(f\"---------------\\n{implementation}:\")\n", - " for n_agents in n_agents_list:\n", - " if implementation == \"mesa\":\n", - " ntime = run_simulation(MesaMoneyModel(n_agents), n_steps)\n", - " elif implementation == \"mesa-frames (pl concise)\":\n", - " ntime = run_simulation(MoneyModel(n_agents, MoneyAgentsConcise), n_steps)\n", - " elif implementation == \"mesa-frames (pl native)\":\n", - " ntime = run_simulation(MoneyModel(n_agents, MoneyAgentsNative), n_steps)\n", - "\n", - " print(f\" Number of agents: {n_agents}, Time: {ntime:.2f} seconds\")\n", - " print(\"---------------\")" - ] - }, - { - "cell_type": "markdown", - "id": "6dfc6d34", - "metadata": {}, - "source": [ - "\n", - "## Conclusion ๐ŸŽ‰\n", - "\n", - "- All mesa-frames implementations significantly outperform the original mesa implementation. ๐Ÿ†\n", - "- The native implementation for Polars shows better performance than their concise counterparts. ๐Ÿ’ช\n", - "- The Polars native implementation shows the most impressive speed-up, ranging from 10.86x to 17.60x faster than mesa! ๐Ÿš€๐Ÿš€๐Ÿš€\n", - "- The performance advantage of mesa-frames becomes more pronounced as the number of agents increases. ๐Ÿ“ˆ" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "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.3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/general/user-guide/2_introductory_tutorial.py b/docs/general/user-guide/2_introductory_tutorial.py new file mode 100644 index 00000000..8560034a --- /dev/null +++ b/docs/general/user-guide/2_introductory_tutorial.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +# %% [markdown] +"""[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/2_introductory-tutorial.ipynb)""" + +# %% [markdown] +"""## Installation (if running in Colab) + +Run the following cell to install `mesa-frames` if you are using Google Colab.""" + +# %% +# !pip install git+https://github.com/projectmesa/mesa-frames mesa + +# %% [markdown] +""" # Introductory Tutorial: Boltzmann Wealth Model with mesa-frames ๐Ÿ’ฐ๐Ÿš€ + +In this tutorial, we'll implement the Boltzmann Wealth Model using mesa-frames. This model simulates the distribution of wealth among agents, where agents randomly give money to each other. + +## Setting Up the Model ๐Ÿ—๏ธ + +First, let's import the necessary modules and set up our model class:""" + +# %% +from mesa_frames import Model, AgentSet, DataCollector + + +class MoneyModel(Model): + def __init__(self, N: int, agents_cls): + super().__init__() + self.n_agents = N + self.sets += agents_cls(N, self) + self.datacollector = DataCollector( + model=self, + model_reporters={ + "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum() + }, + agent_reporters={"wealth": "wealth"}, + storage="csv", + storage_uri="./data", + trigger=lambda m: m.schedule.steps % 2 == 0, + ) + + def step(self): + # Executes the step method for every agentset in self.sets + self.sets.do("step") + + def run_model(self, n): + for _ in range(n): + self.step() + self.datacollector.conditional_collect + self.datacollector.flush() + + +# %% [markdown] +"""## Implementing the AgentSet ๐Ÿ‘ฅ + +Now, let's implement our `MoneyAgents` using polars backends.""" + +# %% +import polars as pl + + +class MoneyAgents(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) + + def step(self) -> None: + self.do("give_money") + + def give_money(self): + self.select(self.wealth > 0) + other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) + self["active", "wealth"] -= 1 + new_wealth = other_agents.group_by("unique_id").len() + self[new_wealth["unique_id"], "wealth"] += new_wealth["len"] + + +# %% [markdown] +""" +## Running the Model โ–ถ๏ธ + +Now that we have our model and agent set defined, let's run a simulation:""" + +# %% +# Create and run the model +model = MoneyModel(1000, MoneyAgents) +model.run_model(100) + +wealth_dist = list(model.sets.df.values())[0] + +# Print the final wealth distribution +print(wealth_dist.select(pl.col("wealth")).describe()) + +# %% [markdown] +""" +This output shows the statistical summary of the wealth distribution after 100 steps of the simulation with 1000 agents. + +## Performance Comparison ๐ŸŽ๏ธ๐Ÿ’จ + +One of the key advantages of mesa-frames is its performance with large numbers of agents. Let's compare the performance of mesa and polars:""" + + +# %% +class MoneyAgentsConcise(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + ## Adding the agents to the agent set + # 1. Changing the df attribute directly (not recommended, if other agents were added before, they will be lost) + """self.df = pl.DataFrame( + {"wealth": pl.ones(n, eager=True)} + )""" + # 2. Adding the dataframe with add + """self.add( + pl.DataFrame( + { + "wealth": pl.ones(n, eager=True), + } + ) + )""" + # 3. Adding the dataframe with __iadd__ + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) + + def step(self) -> None: + # The give_money method is called + # self.give_money() + self.do("give_money") + + def give_money(self): + ## Active agents are changed to wealthy agents + # 1. Using the __getitem__ method + # self.select(self["wealth"] > 0) + # 2. Using the fallback __getattr__ method + self.select(self.wealth > 0) + + # Receiving agents are sampled (only native expressions currently supported) + other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) + + # Wealth of wealthy is decreased by 1 + # 1. Using the __setitem__ method with self.active_agents mask + # self[self.active_agents, "wealth"] -= 1 + # 2. Using the __setitem__ method with "active" mask + self["active", "wealth"] -= 1 + + # Compute the income of the other agents (only native expressions currently supported) + new_wealth = other_agents.group_by("unique_id").len() + + # Add the income to the other agents + # 1. Using the set method + """self.set( + attr_names="wealth", + values=pl.col("wealth") + new_wealth["len"], + mask=new_wealth, + )""" + + # 2. Using the __setitem__ method + self[new_wealth, "wealth"] += new_wealth["len"] + + +class MoneyAgentsNative(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) + + def step(self) -> None: + self.do("give_money") + + def give_money(self): + ## Active agents are changed to wealthy agents + self.select(pl.col("wealth") > 0) + + other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) + + # Wealth of wealthy is decreased by 1 + self.df = self.df.with_columns( + wealth=pl.when( + pl.col("unique_id").is_in(self.active_agents["unique_id"].implode()) + ) + .then(pl.col("wealth") - 1) + .otherwise(pl.col("wealth")) + ) + + new_wealth = other_agents.group_by("unique_id").len() + + # Add the income to the other agents + self.df = ( + self.df.join(new_wealth, on="unique_id", how="left") + .fill_null(0) + .with_columns(wealth=pl.col("wealth") + pl.col("len")) + .drop("len") + ) + + +# %% [markdown] +"""Add Mesa implementation of MoneyAgent and MoneyModel classes to test Mesa performance""" + +# %% +import mesa + + +class MesaMoneyAgent(mesa.Agent): + """An agent with fixed initial wealth.""" + + def __init__(self, model): + # Pass the parameters to the parent class. + super().__init__(model) + + # Create the agent's variable and set the initial values. + self.wealth = 1 + + def step(self): + # Verify agent has some wealth + if self.wealth > 0: + other_agent: MesaMoneyAgent = self.model.random.choice(self.model.agents) + if other_agent is not None: + other_agent.wealth += 1 + self.wealth -= 1 + + +class MesaMoneyModel(mesa.Model): + """A model with some number of agents.""" + + def __init__(self, N: int): + super().__init__() + self.num_agents = N + for _ in range(N): + self.agents.add(MesaMoneyAgent(self)) + + def step(self): + """Advance the model by one step.""" + self.agents.shuffle_do("step") + + def run_model(self, n_steps) -> None: + for _ in range(n_steps): + self.step() + + +# %% +import time + + +def run_simulation(model: MesaMoneyModel | MoneyModel, n_steps: int): + start_time = time.time() + model.run_model(n_steps) + end_time = time.time() + return end_time - start_time + + +# Compare mesa and mesa-frames implementations +n_agents_list = [10**2, 10**3 + 1, 2 * 10**3] +n_steps = 100 +print("Execution times:") +for implementation in [ + "mesa", + "mesa-frames (pl concise)", + "mesa-frames (pl native)", +]: + print(f"---------------\n{implementation}:") + for n_agents in n_agents_list: + if implementation == "mesa": + ntime = run_simulation(MesaMoneyModel(n_agents), n_steps) + elif implementation == "mesa-frames (pl concise)": + ntime = run_simulation(MoneyModel(n_agents, MoneyAgentsConcise), n_steps) + elif implementation == "mesa-frames (pl native)": + ntime = run_simulation(MoneyModel(n_agents, MoneyAgentsNative), n_steps) + + print(f" Number of agents: {n_agents}, Time: {ntime:.2f} seconds") + print("---------------") + +# %% [markdown] +""" +## Conclusion ๐ŸŽ‰ + +- All mesa-frames implementations significantly outperform the original mesa implementation. ๐Ÿ† +- The native implementation for Polars shows better performance than their concise counterparts. ๐Ÿ’ช +- The Polars native implementation shows the most impressive speed-up, ranging from 10.86x to 17.60x faster than mesa! ๐Ÿš€๐Ÿš€๐Ÿš€ +- The performance advantage of mesa-frames becomes more pronounced as the number of agents increases. ๐Ÿ“ˆ""" diff --git a/docs/general/user-guide/3_advanced-tutorial.md b/docs/general/user-guide/3_advanced-tutorial.md deleted file mode 100644 index 8a2eae55..00000000 --- a/docs/general/user-guide/3_advanced-tutorial.md +++ /dev/null @@ -1,4 +0,0 @@ -# Advanced Tutorial: SugarScape with Instantaneous Growback ๐Ÿฌ๐Ÿ”„ - -!!! warning "Work in Progress ๐Ÿšง" - This tutorial is coming soon! ๐Ÿ”œโœจ In the meantime, you can check out the code in the `examples/sugarscape-ig` directory of the mesa-frames repository. diff --git a/docs/general/user-guide/3_advanced_tutorial.py b/docs/general/user-guide/3_advanced_tutorial.py new file mode 100644 index 00000000..b1009734 --- /dev/null +++ b/docs/general/user-guide/3_advanced_tutorial.py @@ -0,0 +1,1644 @@ +from __future__ import annotations + +# %% [markdown] +""" +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/3_advanced_tutorial.ipynb) + +# Advanced Tutorial โ€” Rebuilding Sugarscape with mesa-frames + +We revisit the classic Sugarscape instant-growback model described in chapter 2 of [Growing Artificial Societies](https://direct.mit.edu/books/monograph/2503/Growing-Artificial-SocietiesSocial-Science-from) (Epstein & Axtell, +1996) and rebuild it step by step using `mesa-frames`. Along the way we highlight why the traditional definition is not ideal for high-performance with mesa-frames and how a simple relaxation can unlock vectorisation and lead to similar macro behaviour. + +## Sugarscape in Plain Terms + +We model a population of *ants* living on a rectangular grid rich in sugar. Each +cell can host at most one ant and holds a fixed amount of sugar. Every time step +unfolds as follows: + +* **Sense:** each ant looks outward along the four cardinal directions up to its + `vision` radius and spots open cells. +* **Move:** the ant chooses the cell with highest sugar (breaking ties by + distance and coordinates). The sugar on cells that are already occupied (including its own) is 0. +* **Eat & survive:** ants harvest the sugar on the cell they occupy. If their + sugar stock falls below their `metabolism` cost, they die. +* **Regrow:** sugar instantly regrows to its maximum level on empty cells. The + landscape is drawn from a uniform distribution, so resources are homogeneous + on average and the interesting dynamics come from agent heterogeneity and + congestion. + +The update schedule matters for micro-behaviour, so we study three variants: + +1. **Sequential loop (asynchronous):** This is the traditional definition. Ants move one at a time in random order. +This cannot be vectorised easily as the best move for an ant might depend on the moves of earlier ants (for example, if they target the same cell). +2. **Sequential with Numba:** matches the first variant but relies on a compiled + helper for speed. +3. **Parallel (synchronous):** all ants propose moves; conflicts are resolved at + random before applying the winners simultaneously (and the losers get to their second-best cell, etc). + +The first variant (pure Python loops) is a natural starting point, but it is **not** the mesa-frames philosophy. +The latter two are: we aim to **write rules declaratively** and let the dataframe engine worry about performance. +Our guiding principle is to **focus on modelling first and performance second**. Only when a rule is truly +inherently sequential do we fall back to a compiled kernel (Numba or JAX). + +Our goal is to show that, under instantaneous growback and uniform resources, +the model converges to the *same* macroscopic inequality pattern regardless of +whether agents act sequentially or in parallel and that As long as the random draws do +not push the system into extinction, the long-run Gini coefficient of wealth and +the wealthโ€“trait correlations line up within sampling error โ€” a classic example +of emergent macro regularities in agent-based models. +""" + +# %% [markdown] +# First, let's install and import the necessary packages. + +# %% [markdown] +# If you're running this tutorial on Google Colab or another fresh environment, +# uncomment the cell below to install the required dependencies. + +# %% +# !pip install git+https://github.com/projectmesa/mesa-frames polars numba numpy + +# %% [markdown] +"""## 1. Imports""" + +# %% +from time import perf_counter + +import numpy as np +import polars as pl +from numba import njit + +from mesa_frames import AgentSet, DataCollector, Grid, Model + +# %% [markdown] +"""## 2. Model definition + +In this section we define some helpers and the model class that wires +together the grid and the agents. The `agent_type` parameter stays flexible so +we can plug in different movement policies later, but the model now owns the +logic that generates the sugar field and the initial agent frame. Because both +helpers use `self.random`, instantiating each variant with the same seed keeps +the initial conditions identical across the sequential, Numba, and parallel +implementations. + +The space is a von Neumann grid (which means agents can only move up, down, left, +or right) with capacity 1, meaning each cell can host at most one agent. The sugar +field is stored as part of the cell data frame, with columns for current sugar +and maximum sugar (for regrowth). The model also sets up a data collector to +track aggregate statistics and agent traits over time. + +The `step` method advances the sugar field, triggers the agent set's step. + +We also define some useful functions to compute metrics like the Gini coefficient and correlations. +""" + + +# %% + +# Model-level reporters + + +def gini(model: Model) -> float: + """Compute the Gini coefficient of agent sugar holdings. + + The function reads the primary agent set from ``model.sets[0]`` and + computes the population Gini coefficient on the ``sugar`` column. The + implementation is robust to empty sets and zero-total sugar. + + Parameters + ---------- + model : Model + The simulation model that contains agent sets. The primary agent set + is expected to be at ``model.sets[0]`` and to expose a Polars DataFrame + under ``.df`` with a ``sugar`` column. + + Returns + ------- + float + Gini coefficient in the range [0, 1] if defined, ``0.0`` when the + total sugar is zero, and ``nan`` when the agent set is empty or too + small to measure. + """ + if len(model.sets) == 0: + return float("nan") + + primary_set = model.sets[0] + if len(primary_set) == 0: + return float("nan") + + sugar = primary_set.df["sugar"].to_numpy().astype(np.float64) + + if sugar.size == 0: + return float("nan") + sorted_vals = np.sort(sugar.astype(np.float64)) + n = sorted_vals.size + if n == 0: + return float("nan") + cumulative = np.cumsum(sorted_vals) + total = cumulative[-1] + if total == 0: + return 0.0 + index = np.arange(1, n + 1, dtype=np.float64) + return float((2.0 * np.dot(index, sorted_vals) / (n * total)) - (n + 1) / n) + + +def corr_sugar_metabolism(model: Model) -> float: + """Pearson correlation between agent sugar and metabolism. + + This reporter extracts the ``sugar`` and ``metabolism`` columns from the + primary agent set and returns their Pearson correlation coefficient. When + the agent set is empty or contains insufficient variation the function + returns ``nan``. + + Parameters + ---------- + model : Model + The simulation model that contains agent sets. The primary agent set + is expected to be at ``model.sets[0]`` and provide a Polars DataFrame + with ``sugar`` and ``metabolism`` columns. + + Returns + ------- + float + Pearson correlation coefficient between sugar and metabolism, or + ``nan`` when the correlation is undefined (empty set or constant + values). + """ + if len(model.sets) == 0: + return float("nan") + + primary_set = model.sets[0] + if len(primary_set) == 0: + return float("nan") + + agent_df = primary_set.df + sugar = agent_df["sugar"].to_numpy().astype(np.float64) + metabolism = agent_df["metabolism"].to_numpy().astype(np.float64) + return _safe_corr(sugar, metabolism) + + +def corr_sugar_vision(model: Model) -> float: + """Pearson correlation between agent sugar and vision. + + Extracts the ``sugar`` and ``vision`` columns from the primary agent set + and returns their Pearson correlation coefficient. If the reporter cannot + compute a meaningful correlation (for example, when the agent set is + empty or values are constant) it returns ``nan``. + + Parameters + ---------- + model : Model + The simulation model that contains agent sets. The primary agent set + is expected to be at ``model.sets[0]`` and provide a Polars DataFrame + with ``sugar`` and ``vision`` columns. + + Returns + ------- + float + Pearson correlation coefficient between sugar and vision, or ``nan`` + when the correlation is undefined. + """ + if len(model.sets) == 0: + return float("nan") + + primary_set = model.sets[0] + if len(primary_set) == 0: + return float("nan") + + agent_df = primary_set.df + sugar = agent_df["sugar"].to_numpy().astype(np.float64) + vision = agent_df["vision"].to_numpy().astype(np.float64) + return _safe_corr(sugar, vision) + + +def _safe_corr(x: np.ndarray, y: np.ndarray) -> float: + """Safely compute Pearson correlation between two 1-D arrays. + + This helper guards against degenerate inputs (too few observations or + constant arrays) which would make the Pearson correlation undefined or + numerically unstable. When a valid correlation can be computed the + function returns a Python float. + + Parameters + ---------- + x : np.ndarray + One-dimensional numeric array containing the first variable to + correlate. + y : np.ndarray + One-dimensional numeric array containing the second variable to + correlate. + + Returns + ------- + float + Pearson correlation coefficient as a Python float, or ``nan`` if the + correlation is undefined (fewer than 2 observations or constant + inputs). + """ + if x.size < 2 or y.size < 2: + return float("nan") + if np.allclose(x, x[0]) or np.allclose(y, y[0]): + return float("nan") + return float(np.corrcoef(x, y)[0, 1]) + + +class Sugarscape(Model): + """Minimal Sugarscape model used throughout the tutorial. + + This class wires together a grid that stores ``sugar`` per cell, an + agent set implementation (passed in as ``agent_type``), and a + data collector that records model- and agent-level statistics. + + The model's responsibilities are to: + - create the sugar landscape (cells with current and maximum sugar) + - create and place agents on the grid + - advance the sugar regrowth rule each step + - run the model for a fixed number of steps and collect data + + Parameters + ---------- + agent_type : type[AntsBase] + The :class:`AgentSet` subclass implementing the movement rules + (sequential, numba-accelerated, or parallel). + n_agents : int + Number of agents to create and place on the grid. + width : int + Grid width (number of columns). + height : int + Grid height (number of rows). + max_sugar : int, optional + Upper bound for the randomly initialised sugar values on the grid, + by default 4. + seed : int | None, optional + RNG seed to make runs reproducible across variants, by default None. + + Notes + ----- + The grid uses a von Neumann neighbourhood and capacity 1 (at most one + agent per cell). Both the sugar landscape and initial agent traits are + drawn from ``self.random`` so different movement variants can be + instantiated with identical initial conditions by passing the same seed. + """ + + def __init__( + self, + agent_type: type[AntsBase], + n_agents: int, + *, + width: int, + height: int, + max_sugar: int = 4, + seed: int | None = None, + ) -> None: + if n_agents > width * height: + raise ValueError( + "Cannot place more agents than grid cells when capacity is 1." + ) + super().__init__(seed) + + # 1. Let's create the sugar grid and set up the space + + sugar_grid_df = self._generate_sugar_grid(width, height, max_sugar) + self.space = Grid( + self, [width, height], neighborhood_type="von_neumann", capacity=1 + ) + self.space.set_cells(sugar_grid_df) + self._max_sugar = sugar_grid_df.select(["dim_0", "dim_1", "max_sugar"]) + + # 2. Now we create the agents and place them on the grid + + agent_frame = self._generate_agent_frame(n_agents) + main_set = agent_type(self, agent_frame) + self.sets += main_set + self.space.place_to_empty(self.sets) + + # 3. Finally we set up the data collector + self.datacollector = DataCollector( + model=self, + model_reporters={ + "mean_sugar": lambda m: 0.0 + if len(m.sets[0]) == 0 + else float(m.sets[0].df["sugar"].mean()), + "total_sugar": lambda m: float(m.sets[0].df["sugar"].sum()) + if len(m.sets[0]) + else 0.0, + "agents_alive": lambda m: float(len(m.sets[0])) if len(m.sets) else 0.0, + "gini": gini, + "corr_sugar_metabolism": corr_sugar_metabolism, + "corr_sugar_vision": corr_sugar_vision, + }, + agent_reporters={"traits": ["sugar", "metabolism", "vision"]}, + ) + self.datacollector.collect() + + def _generate_sugar_grid( + self, width: int, height: int, max_sugar: int + ) -> pl.DataFrame: + """Generate a random sugar grid. + + Parameters + ---------- + width : int + Grid width (number of columns). + height : int + Grid height (number of rows). + max_sugar : int + Maximum sugar value (inclusive) for each cell. + + Returns + ------- + pl.DataFrame + DataFrame with columns ``dim_0``, ``dim_1``, ``sugar`` (current + amount) and ``max_sugar`` (regrowth target). + """ + sugar_vals = self.random.integers( + 0, max_sugar + 1, size=(width, height), dtype=np.int64 + ) + dim_0 = pl.Series("dim_0", pl.arange(width, eager=True)).to_frame() + dim_1 = pl.Series("dim_1", pl.arange(height, eager=True)).to_frame() + return dim_0.join(dim_1, how="cross").with_columns( + sugar=sugar_vals.flatten(), max_sugar=sugar_vals.flatten() + ) + + def _generate_agent_frame(self, n_agents: int) -> pl.DataFrame: + """Create the initial agent frame populated with agent traits. + + Parameters + ---------- + n_agents : int + Number of agents to create. + + Returns + ------- + pl.DataFrame + DataFrame with columns ``sugar``, ``metabolism`` and ``vision`` + (integer values) for each agent. + """ + rng = self.random + return pl.DataFrame( + { + "sugar": rng.integers(6, 25, size=n_agents, dtype=np.int64), + "metabolism": rng.integers(2, 5, size=n_agents, dtype=np.int64), + "vision": rng.integers(1, 6, size=n_agents, dtype=np.int64), + } + ) + + def step(self) -> None: + """Advance the model by one step. + + Notes + ----- + The per-step ordering is important and this tutorial implements the + classic Sugarscape "instant growback": agents move and eat first, + and then empty cells are refilled immediately (move -> eat -> regrow + -> collect). + """ + if len(self.sets[0]) == 0: + self.running = False + return + self.sets[0].step() + self._advance_sugar_field() + self.datacollector.collect() + if len(self.sets[0]) == 0: + self.running = False + + def run(self, steps: int) -> None: + """Run the model for a fixed number of steps. + + Parameters + ---------- + steps : int + Maximum number of steps to run. The model may terminate earlier if + ``self.running`` is set to ``False`` (for example, when all agents + have died). + """ + for _ in range(steps): + if not self.running: + break + self.step() + + def _advance_sugar_field(self) -> None: + """Apply the instant-growback sugar regrowth rule. + + Empty cells (no agent present) are refilled to their ``max_sugar`` + value. Cells that are occupied are set to zero because agents harvest + the sugar when they eat. The method uses vectorised DataFrame joins + and writes to keep the operation efficient. + """ + empty_cells = self.space.empty_cells + if not empty_cells.is_empty(): + # Look up the maximum sugar for each empty cell and restore it. + refresh = empty_cells.join( + self._max_sugar, on=["dim_0", "dim_1"], how="left" + ) + self.space.set_cells(empty_cells, {"sugar": refresh["max_sugar"]}) + full_cells = self.space.full_cells + if not full_cells.is_empty(): + # Occupied cells have just been harvested; set their sugar to 0. + zeros = pl.Series(np.zeros(len(full_cells), dtype=np.int64)) + self.space.set_cells(full_cells, {"sugar": zeros}) + + +# %% [markdown] + +""" +## 3. Agent definition + +### 3.1 Base agent class + +Now let's define the agent class (the ant class). We start with a base class which implements the common logic for eating and starvation, while leaving the `move` method abstract. +The base class also provides helper methods for sensing visible cells and choosing the best cell based on sugar, distance, and coordinates. +This will allow us to define different movement policies (sequential, Numba-accelerated, and parallel) as subclasses that only need to implement the `move` method. +""" + +# %% + + +class AntsBase(AgentSet): + """Base agent set for the Sugarscape tutorial. + + This class implements the common behaviour shared by all agent + movement variants (sequential, numba-accelerated and parallel). + + Notes + ----- + - Agents are expected to have integer traits: ``sugar``, ``metabolism`` + and ``vision``. These are validated in :meth:`__init__`. + - Subclasses must implement :meth:`move` which changes agent positions + on the grid (via :meth:`mesa_frames.Grid` helpers). + """ + + def __init__(self, model: Model, agent_frame: pl.DataFrame) -> None: + """Initialise the agent set and validate required trait columns. + + Parameters + ---------- + model : Model + The parent model which provides RNG and space. + agent_frame : pl.DataFrame + A Polars DataFrame with at least the columns ``sugar``, + ``metabolism`` and ``vision`` for each agent. + + Raises + ------ + ValueError + If required trait columns are missing from ``agent_frame``. + """ + super().__init__(model) + required = {"sugar", "metabolism", "vision"} + missing = required.difference(agent_frame.columns) + if missing: + raise ValueError( + f"Initial agent frame must include columns {sorted(required)}; missing {sorted(missing)}." + ) + self.add(agent_frame.clone()) + + def step(self) -> None: + """Advance the agent set by one time step. + + The update order is important: agents are first shuffled to randomise + move order (this is important only for sequential variants), then they move, harvest sugar + from their occupied cells, and finally any agents whose sugar falls + to zero or below are removed. + """ + # Randomise ordering for movement decisions when required by the + # implementation (e.g. sequential update uses this shuffle). + self.shuffle(inplace=True) + # Movement policy implemented by subclasses. + self.move() + # Agents harvest sugar on their occupied cells. + self.eat() + # Remove agents that starved after eating. + self._remove_starved() + + def move(self) -> None: # pragma: no cover + """Abstract movement method. + + Subclasses must override this method to update agent positions on the + grid. Implementations should use :meth:`mesa_frames.Grid.move_agents` + or similar helpers provided by the space API. + """ + raise NotImplementedError + + def eat(self) -> None: + """Agents harvest sugar from the cells they currently occupy. + + Behaviour: + - Look up the set of occupied cells (cells that reference an agent + id). + - For each occupied cell, add the cell sugar to the agent's sugar + stock and subtract the agent's metabolism cost. + - After agents harvest, set the sugar on those cells to zero (they + were consumed). + """ + # Map of currently occupied agent ids on the grid. + occupied_ids = self.index + # `occupied_ids` is a Polars Series; calling `is_in` with a Series + # of the same datatype is ambiguous in newer Polars. Use `implode` + # to collapse the Series into a list-like value for membership checks. + occupied_cells = self.space.cells.filter( + pl.col("agent_id").is_in(occupied_ids.implode()) + ) + if occupied_cells.is_empty(): + return + # The agent ordering here uses the agent_id values stored in the + # occupied cells frame; indexing the agent set with that vector updates + # the matching agents' sugar values in one vectorised write. + agent_ids = occupied_cells["agent_id"] + self[agent_ids, "sugar"] = ( + self[agent_ids, "sugar"] + + occupied_cells["sugar"] + - self[agent_ids, "metabolism"] + ) + # After harvesting, occupied cells have zero sugar. + self.space.set_cells( + occupied_cells.select(["dim_0", "dim_1"]), + {"sugar": pl.Series(np.zeros(len(occupied_cells), dtype=np.int64))}, + ) + + def _remove_starved(self) -> None: + """Discard agents whose sugar stock has fallen to zero or below. + + This method performs a vectorised filter on the agent frame and + removes any matching rows from the set. + """ + starved = self.df.filter(pl.col("sugar") <= 0) + if not starved.is_empty(): + # ``discard`` accepts a DataFrame of agents to remove. + self.discard(starved) + + +# %% [markdown] + +"""### 3.2 Sequential movement + +We now implement the simplest movement policy: sequential (asynchronous). Each agent moves one at a time in the current ordering, choosing the best visible cell according to the rules. + +This implementation uses plain Python loops as the logic cannot be easily vectorised. As a result, it is slow for large populations and grids. We will later show how to speed it up with Numba. +""" + +# %% + + +class AntsSequential(AntsBase): + def _visible_cells( + self, origin: tuple[int, int], vision: int + ) -> list[tuple[int, int]]: + """List cells visible from an origin along the four cardinal axes. + + The visibility set includes the origin cell itself and cells at + Manhattan distances 1..vision along the four cardinal directions + (up, down, left, right), clipped to the grid bounds. + + Parameters + ---------- + origin : tuple[int, int] + The agent's current coordinate ``(x, y)``. + vision : int + Maximum Manhattan radius to consider along each axis. + + Returns + ------- + list[tuple[int, int]] + Ordered list of visible cells (origin first, then increasing + step distance along each axis). + """ + x0, y0 = origin + width, height = self.space.dimensions + cells: list[tuple[int, int]] = [origin] + # Look outward one step at a time in the four cardinal directions. + for step in range(1, vision + 1): + if x0 + step < width: + cells.append((x0 + step, y0)) + if x0 - step >= 0: + cells.append((x0 - step, y0)) + if y0 + step < height: + cells.append((x0, y0 + step)) + if y0 - step >= 0: + cells.append((x0, y0 - step)) + return cells + + def _choose_best_cell( + self, + origin: tuple[int, int], + vision: int, + sugar_map: dict[tuple[int, int], int], + blocked: set[tuple[int, int]] | None, + ) -> tuple[int, int]: + """Select the best visible cell according to the movement rules. + + Tie-break rules (in order): + 1. Prefer cells with strictly greater sugar. + 2. If equal sugar, prefer the cell with smaller distance from the + origin (measured with the Frobenius norm returned by + ``space.get_distances``). + 3. If still tied, prefer the cell with smaller coordinates (lexicographic + ordering of the ``(x, y)`` tuple). + + Parameters + ---------- + origin : tuple[int, int] + Agent's current coordinate. + vision : int + Maximum vision radius along cardinal axes. + sugar_map : dict[tuple[int, int], int] + Mapping from ``(x, y)`` to sugar amount. + blocked : set[tuple[int, int]] | None + Optional set of coordinates that should be considered occupied and + therefore skipped (except the origin which is always allowed). + + Returns + ------- + tuple[int, int] + Chosen target coordinate (may be the origin if no better cell is + available). + """ + best_cell = origin + best_sugar = sugar_map.get(origin, 0) + best_distance = 0 + ox, oy = origin + for candidate in self._visible_cells(origin, vision): + # Skip blocked cells (occupied by other agents) unless it's the + # agent's current cell which we always consider. + if blocked and candidate != origin and candidate in blocked: + continue + sugar_here = sugar_map.get(candidate, 0) + # Use step-based Manhattan distance (number of steps along cardinal + # axes) which is the same metric used by the Numba path. This avoids + # calling the heavier `space.get_distances` per candidate. + cx, cy = candidate + distance = abs(cx - ox) + abs(cy - oy) + better = False + # Primary criterion: strictly more sugar. + if sugar_here > best_sugar: + better = True + elif sugar_here == best_sugar: + # Secondary: closer distance. + if distance < best_distance: + better = True + # Tertiary: lexicographic tie-break on coordinates. + elif distance == best_distance and candidate < best_cell: + better = True + if better: + best_cell = candidate + best_sugar = sugar_here + best_distance = distance + return best_cell + + def _current_sugar_map(self) -> dict[tuple[int, int], int]: + """Return a mapping from grid coordinates to the current sugar value. + + Returns + ------- + dict[tuple[int, int], int] + Keys are ``(x, y)`` tuples and values are the integer sugar amount + on that cell (zero if missing/None). + """ + cells = self.space.cells.select(["dim_0", "dim_1", "sugar"]) + # Build a plain Python dict for fast lookups in the movement code. + return { + (int(x), int(y)): 0 if sugar is None else int(sugar) + for x, y, sugar in cells.iter_rows() + } + + def move(self) -> None: + sugar_map = self._current_sugar_map() + state = self.df.join(self.pos, on="unique_id", how="left") + positions = { + int(row["unique_id"]): (int(row["dim_0"]), int(row["dim_1"])) + for row in state.iter_rows(named=True) + } + taken: set[tuple[int, int]] = set(positions.values()) + + for row in state.iter_rows(named=True): + agent_id = int(row["unique_id"]) + vision = int(row["vision"]) + current = positions[agent_id] + taken.discard(current) + target = self._choose_best_cell(current, vision, sugar_map, taken) + taken.add(target) + positions[agent_id] = target + if target != current: + self.space.move_agents(agent_id, target) + + +# %% [markdown] +""" +### 3.3 Speeding Up the Loop with Numba + +As we will see later, the previous sequential implementation is slow for large populations and grids because it relies on plain Python loops. We can speed it up significantly by using Numba to compile the movement logic. + +Numba compiles numerical Python code to fast machine code at runtime. To use Numba, we need to rewrite the movement logic in a way that is compatible with Numba's restrictions (using tightly typed numpy arrays and accessing data indexes directly). +""" + + +# %% +@njit(cache=True) +def _numba_should_replace( + best_sugar: int, + best_distance: int, + best_x: int, + best_y: int, + candidate_sugar: int, + candidate_distance: int, + candidate_x: int, + candidate_y: int, +) -> bool: + """Numba helper: decide whether a candidate cell should replace the + current best cell according to the movement tie-break rules. + + This implements the same ordering used in :meth:`_choose_best_cell` but + in a tightly-typed, compiled form suitable for Numba loops. + + Parameters + ---------- + best_sugar : int + Sugar at the current best cell. + best_distance : int + Manhattan distance from the origin to the current best cell. + best_x : int + X coordinate of the current best cell. + best_y : int + Y coordinate of the current best cell. + candidate_sugar : int + Sugar at the candidate cell. + candidate_distance : int + Manhattan distance from the origin to the candidate cell. + candidate_x : int + X coordinate of the candidate cell. + candidate_y : int + Y coordinate of the candidate cell. + + Returns + ------- + bool + True if the candidate should replace the current best cell. + """ + # Primary criterion: prefer strictly greater sugar. + if candidate_sugar > best_sugar: + return True + # If sugar ties, prefer the closer cell. + if candidate_sugar == best_sugar: + if candidate_distance < best_distance: + return True + # If distance ties as well, compare coordinates lexicographically. + if candidate_distance == best_distance: + if candidate_x < best_x: + return True + if candidate_x == best_x and candidate_y < best_y: + return True + return False + + +@njit(cache=True) +def _numba_find_best_cell( + x0: int, + y0: int, + vision: int, + sugar_array: np.ndarray, + occupied: np.ndarray, +) -> tuple[int, int]: + width, height = sugar_array.shape + best_x = x0 + best_y = y0 + best_sugar = sugar_array[x0, y0] + best_distance = 0 + + # Examine visible cells along the four cardinal directions, increasing + # step by step. The 'occupied' array marks cells that are currently + # unavailable (True = occupied). The origin cell is allowed as the + # default; callers typically clear the origin before searching. + for step in range(1, vision + 1): + nx = x0 + step + if nx < width and not occupied[nx, y0]: + sugar_here = sugar_array[nx, y0] + if _numba_should_replace( + best_sugar, best_distance, best_x, best_y, sugar_here, step, nx, y0 + ): + best_x = nx + best_y = y0 + best_sugar = sugar_here + best_distance = step + + nx = x0 - step + if nx >= 0 and not occupied[nx, y0]: + sugar_here = sugar_array[nx, y0] + if _numba_should_replace( + best_sugar, best_distance, best_x, best_y, sugar_here, step, nx, y0 + ): + best_x = nx + best_y = y0 + best_sugar = sugar_here + best_distance = step + + ny = y0 + step + if ny < height and not occupied[x0, ny]: + sugar_here = sugar_array[x0, ny] + if _numba_should_replace( + best_sugar, best_distance, best_x, best_y, sugar_here, step, x0, ny + ): + best_x = x0 + best_y = ny + best_sugar = sugar_here + best_distance = step + + ny = y0 - step + if ny >= 0 and not occupied[x0, ny]: + sugar_here = sugar_array[x0, ny] + if _numba_should_replace( + best_sugar, best_distance, best_x, best_y, sugar_here, step, x0, ny + ): + best_x = x0 + best_y = ny + best_sugar = sugar_here + best_distance = step + + return best_x, best_y + + +@njit(cache=True) +def sequential_move_numba( + dim0: np.ndarray, + dim1: np.ndarray, + vision: np.ndarray, + sugar_array: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Numba-accelerated sequential movement helper. + + This function emulates the traditional asynchronous (sequential) update + where agents move one at a time in the current ordering. It accepts + numpy arrays describing agent positions and vision ranges, and a 2D + sugar array for lookup. + + Parameters + ---------- + dim0 : np.ndarray + 1D integer array of length n_agents containing the x coordinates + for each agent. + dim1 : np.ndarray + 1D integer array of length n_agents containing the y coordinates + for each agent. + vision : np.ndarray + 1D integer array of vision radii for each agent. + sugar_array : np.ndarray + 2D array shaped (width, height) containing per-cell sugar values. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + Updated arrays of x and y coordinates after sequential movement. + """ + n_agents = dim0.shape[0] + width, height = sugar_array.shape + # Copy inputs to avoid mutating caller arrays in-place. + new_dim0 = dim0.copy() + new_dim1 = dim1.copy() + # Occupancy grid: True when a cell is currently occupied by an agent. + occupied = np.zeros((width, height), dtype=np.bool_) + + # Mark initial occupancy. + for i in range(n_agents): + occupied[new_dim0[i], new_dim1[i]] = True + + # Process agents in order. For each agent we clear its current cell in + # the occupancy grid (so it can consider moving into it), search for the + # best unoccupied visible cell, and mark the chosen destination as + # occupied. This models agents moving one-by-one. + for i in range(n_agents): + x0 = new_dim0[i] + y0 = new_dim1[i] + # Free the agent's current cell so it is considered available during + # the search (agents may choose to stay, in which case we'll re-mark + # it below). + occupied[x0, y0] = False + best_x, best_y = _numba_find_best_cell( + x0, y0, int(vision[i]), sugar_array, occupied + ) + # Claim the chosen destination. + occupied[best_x, best_y] = True + new_dim0[i] = best_x + new_dim1[i] = best_y + + return new_dim0, new_dim1 + + +class AntsNumba(AntsBase): + def move(self) -> None: + state = self.df.join(self.pos, on="unique_id", how="left") + if state.is_empty(): + return + agent_ids = state["unique_id"] + dim0 = state["dim_0"].to_numpy().astype(np.int64) + dim1 = state["dim_1"].to_numpy().astype(np.int64) + vision = state["vision"].to_numpy().astype(np.int64) + + sugar_array = ( + self.space.cells.sort(["dim_0", "dim_1"]) + .with_columns(pl.col("sugar").fill_null(0))["sugar"] + .to_numpy() + .reshape(self.space.dimensions) + ) + + new_dim0, new_dim1 = sequential_move_numba(dim0, dim1, vision, sugar_array) + coords = pl.DataFrame({"dim_0": new_dim0.tolist(), "dim_1": new_dim1.tolist()}) + self.space.move_agents(agent_ids, coords) + + +# %% [markdown] +""" +### 3.5 Simultaneous Movement with Conflict Resolution (the Polars mesa-frames idiomatic way) + +The previous implementation is optimal speed-wise but it's a bit low-level. It requires maintaining an occupancy grid and imperative loops and it might become tricky to extend with more complex movement rules or models. +To stay in mesa-frames idiom, we can implement a parallel movement policy that uses Polars DataFrame operations to resolve conflicts when multiple agents target the same cell. +These conflicts are resolved in rounds: in each round, each agent proposes its current best candidate cell; winners per cell are chosen at random, and losers are promoted to their next-ranked choice. This continues until all agents have moved. +This implementation is a tad slower but still efficient and easier to read (for a Polars user). +""" + +# %% + + +class AntsParallel(AntsBase): + def move(self) -> None: + """Move agents in parallel by ranking visible cells and resolving conflicts. + + Declarative mental model: express *what* each agent wants (ranked candidates), + then use dataframe ops to *allocate* (joins, group_by with a lottery). + Performance is handled by Polars/LazyFrames; avoid premature micro-optimisations. + + Returns + ------- + None + Movement updates happen in-place on the underlying space. + """ + # Early exit if there are no agents. + if len(self.df) == 0: + return + + # current_pos columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_center โ”† dim_1_center โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + current_pos = self.pos.select( + [ + pl.col("unique_id").alias("agent_id"), + pl.col("dim_0").alias("dim_0_center"), + pl.col("dim_1").alias("dim_1_center"), + ] + ) + + neighborhood = self._build_neighborhood_frame(current_pos) + choices, origins, max_rank = self._rank_candidates(neighborhood, current_pos) + if choices.is_empty(): + return + + assigned = self._resolve_conflicts_in_rounds(choices, origins, max_rank) + if assigned.is_empty(): + return + + # move_df columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ unique_id โ”† dim_0 โ”† dim_1 โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + move_df = pl.DataFrame( + { + "unique_id": assigned["agent_id"], + "dim_0": assigned["dim_0_candidate"], + "dim_1": assigned["dim_1_candidate"], + } + ) + # `move_agents` accepts IdsLike and SpaceCoordinates (Polars Series/DataFrame), + # so pass Series/DataFrame directly rather than converting to Python lists. + self.space.move_agents(move_df["unique_id"], move_df.select(["dim_0", "dim_1"])) + + def _build_neighborhood_frame(self, current_pos: pl.DataFrame) -> pl.DataFrame: + """Assemble the sugar-weighted neighbourhood for each sensing agent. + + Parameters + ---------- + current_pos : pl.DataFrame + DataFrame with columns ``agent_id``, ``dim_0_center`` and + ``dim_1_center`` describing the current position of each agent. + + Returns + ------- + pl.DataFrame + DataFrame with columns ``agent_id``, ``radius``, ``dim_0_candidate``, + ``dim_1_candidate`` and ``sugar`` describing the visible cells for + each agent. + """ + # Build a neighbourhood frame: for each agent and visible cell we + # attach the cell sugar. The raw offsets contain the candidate + # cell coordinates and the center coordinates for the sensing agent. + # Raw neighborhood columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ dim_0 โ”† dim_1 โ”† radius โ”† dim_0_center โ”† dim_1_center โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ i64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + neighborhood_cells = self.space.get_neighborhood( + radius=self["vision"], agents=self, include_center=True + ) + + # sugar_cells columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ dim_0 โ”† dim_1 โ”† sugar โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”‚ + # โ”‚ i64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ก + + sugar_cells = self.space.cells.select(["dim_0", "dim_1", "sugar"]) + + neighborhood_cells = ( + neighborhood_cells.join(sugar_cells, on=["dim_0", "dim_1"], how="left") + .with_columns(pl.col("sugar").fill_null(0)) + .rename({"dim_0": "dim_0_candidate", "dim_1": "dim_1_candidate"}) + ) + + neighborhood_cells = neighborhood_cells.join( + current_pos, + left_on=["dim_0_center", "dim_1_center"], + right_on=["dim_0_center", "dim_1_center"], + how="left", + ) + + # Final neighborhood columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† radius โ”† dim_0_candidate โ”† dim_1_candidate โ”† sugar โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ก + neighborhood_cells = neighborhood_cells.drop( + ["dim_0_center", "dim_1_center"] + ).select(["agent_id", "radius", "dim_0_candidate", "dim_1_candidate", "sugar"]) + + return neighborhood_cells + + def _rank_candidates( + self, + neighborhood: pl.DataFrame, + current_pos: pl.DataFrame, + ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]: + """Rank candidate destination cells for each agent. + + Parameters + ---------- + neighborhood : pl.DataFrame + Output of :meth:`_build_neighborhood_frame` with columns + ``agent_id``, ``radius``, ``dim_0_candidate``, ``dim_1_candidate`` + and ``sugar``. + current_pos : pl.DataFrame + Frame with columns ``agent_id``, ``dim_0_center`` and + ``dim_1_center`` describing where each agent currently stands. + + Returns + ------- + choices : pl.DataFrame + Ranked candidates per agent with columns ``agent_id``, + ``dim_0_candidate``, ``dim_1_candidate``, ``sugar``, ``radius`` and + ``rank``. + origins : pl.DataFrame + Original coordinates per agent with columns ``agent_id``, + ``dim_0`` and ``dim_1``. + max_rank : pl.DataFrame + Maximum available rank per agent with columns ``agent_id`` and + ``max_rank``. + """ + # Create ranked choices per agent: sort by sugar (desc), radius + # (asc), then coordinates. Keep the first unique entry per cell. + # choices columns (after select): + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_candidate โ”† dim_1_candidate โ”† sugar โ”† radius โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ก + choices = ( + neighborhood.select( + [ + "agent_id", + "dim_0_candidate", + "dim_1_candidate", + "sugar", + "radius", + ] + ) + .with_columns(pl.col("radius")) + .sort( + ["agent_id", "sugar", "radius", "dim_0_candidate", "dim_1_candidate"], + descending=[False, True, False, False, False], + ) + .unique( + subset=["agent_id", "dim_0_candidate", "dim_1_candidate"], + keep="first", + maintain_order=True, + ) + .with_columns(pl.col("agent_id").cum_count().over("agent_id").alias("rank")) + ) + + # Precompute perโ€‘agent candidate rank once so conflict resolution can + # promote losers by incrementing a cheap `current_rank` counter, + # without re-sorting after each round. Alternative: drop taken cells + # and re-rank by sugar every round; simpler conceptually but requires + # repeated sorts and deduplication, which is heavier than filtering by + # `rank >= current_rank`. + + # Origins for fallback (if an agent exhausts candidates it stays put). + # origins columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0 โ”† dim_1 โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + origins = current_pos.select( + [ + "agent_id", + pl.col("dim_0_center").alias("dim_0"), + pl.col("dim_1_center").alias("dim_1"), + ] + ) + + # Track the maximum available rank per agent to clamp promotions. + # This bounds `current_rank`; once an agent reaches `max_rank` and + # cannot secure a cell, they fall back to origin cleanly instead of + # chasing nonexistent ranks. + # max_rank columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† max_rank โ”‚ + # โ”‚ --- โ”† --- โ”‚ + # โ”‚ u64 โ”† u32 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + max_rank = choices.group_by("agent_id").agg( + pl.col("rank").max().alias("max_rank") + ) + return choices, origins, max_rank + + def _resolve_conflicts_in_rounds( + self, + choices: pl.DataFrame, + origins: pl.DataFrame, + max_rank: pl.DataFrame, + ) -> pl.DataFrame: + """Resolve movement conflicts through iterative lottery rounds. + + Parameters + ---------- + choices : pl.DataFrame + Ranked candidate cells per agent with headers matching the + ``choices`` frame returned by :meth:`_rank_candidates`. + origins : pl.DataFrame + Agent origin coordinates with columns ``agent_id``, ``dim_0`` and + ``dim_1``. + max_rank : pl.DataFrame + Maximum rank offset per agent with columns ``agent_id`` and + ``max_rank``. + + Returns + ------- + pl.DataFrame + Allocated movements with columns ``agent_id``, ``dim_0_candidate`` + and ``dim_1_candidate``; each row records the destination assigned + to an agent. + """ + # Prepare unresolved agents and working tables. + agent_ids = choices["agent_id"].unique(maintain_order=True) + + # unresolved columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† current_rank โ”‚ + # โ”‚ --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + unresolved = pl.DataFrame( + { + "agent_id": agent_ids, + "current_rank": pl.Series(np.zeros(len(agent_ids), dtype=np.int64)), + } + ) + + # assigned columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_candidate โ”† dim_1_candidate โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + assigned = pl.DataFrame( + { + "agent_id": pl.Series( + name="agent_id", values=[], dtype=agent_ids.dtype + ), + "dim_0_candidate": pl.Series( + name="dim_0_candidate", values=[], dtype=pl.Int64 + ), + "dim_1_candidate": pl.Series( + name="dim_1_candidate", values=[], dtype=pl.Int64 + ), + } + ) + + # taken columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ dim_0_candidate โ”† dim_1_candidate โ”‚ + # โ”‚ --- โ”† --- โ”‚ + # โ”‚ i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + taken = pl.DataFrame( + { + "dim_0_candidate": pl.Series( + name="dim_0_candidate", values=[], dtype=pl.Int64 + ), + "dim_1_candidate": pl.Series( + name="dim_1_candidate", values=[], dtype=pl.Int64 + ), + } + ) + + # Resolve in rounds: each unresolved agent proposes its current-ranked + # candidate; winners per-cell are selected at random and losers are + # promoted to their next choice. + while unresolved.height > 0: + # Using precomputed `rank` lets us select candidates with + # `rank >= current_rank` and avoid re-ranking after each round. + # Alternative: remove taken cells and re-sort remaining candidates + # by sugar/distance per round (heavier due to repeated sort/dedupe). + # candidate_pool columns (after join with unresolved): + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_candidate โ”† dim_1_candidate โ”† sugar โ”† radius โ”† rank โ”† current_rank โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”† u32 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + candidate_pool = choices.join(unresolved, on="agent_id") + candidate_pool = candidate_pool.filter( + pl.col("rank") >= pl.col("current_rank") + ) + if not taken.is_empty(): + candidate_pool = candidate_pool.join( + taken, + on=["dim_0_candidate", "dim_1_candidate"], + how="anti", + ) + + if candidate_pool.is_empty(): + # No available candidates โ€” everyone falls back to origin. + # Note: this covers both agents with no visible cells left and + # the case where all remaining candidates are already taken. + # fallback columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0 โ”† dim_1 โ”† current_rank โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + fallback = unresolved.join(origins, on="agent_id", how="left") + assigned = pl.concat( + [ + assigned, + fallback.select( + [ + "agent_id", + pl.col("dim_0").alias("dim_0_candidate"), + pl.col("dim_1").alias("dim_1_candidate"), + ] + ), + ], + how="vertical", + ) + break + + # best_candidates columns (per agent first choice): + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_candidate โ”† dim_1_candidate โ”† sugar โ”† radius โ”† rank โ”† current_rank โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”† u32 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + best_candidates = ( + candidate_pool.sort(["agent_id", "rank"]) + .group_by("agent_id", maintain_order=True) + .first() + ) + + # Agents that had no candidate this round fall back to origin. + # missing columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† current_rank โ”‚ + # โ”‚ --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + missing = unresolved.join( + best_candidates.select("agent_id"), on="agent_id", how="anti" + ) + if not missing.is_empty(): + # fallback (missing) columns match fallback table above. + fallback = missing.join(origins, on="agent_id", how="left") + assigned = pl.concat( + [ + assigned, + fallback.select( + [ + "agent_id", + pl.col("dim_0").alias("dim_0_candidate"), + pl.col("dim_1").alias("dim_1_candidate"), + ] + ), + ], + how="vertical", + ) + taken = pl.concat( + [ + taken, + fallback.select( + [ + pl.col("dim_0").alias("dim_0_candidate"), + pl.col("dim_1").alias("dim_1_candidate"), + ] + ), + ], + how="vertical", + ) + unresolved = unresolved.join( + missing.select("agent_id"), on="agent_id", how="anti" + ) + best_candidates = best_candidates.join( + missing.select("agent_id"), on="agent_id", how="anti" + ) + if unresolved.is_empty() or best_candidates.is_empty(): + continue + + # Add a small random lottery to break ties deterministically for + # each candidate set. + lottery = pl.Series("lottery", self.random.random(best_candidates.height)) + best_candidates = best_candidates.with_columns(lottery) + + # winners columns: + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† dim_0_candidate โ”† dim_1_candidate โ”† sugar โ”† radius โ”† rank โ”† current_rank โ”‚ lottery โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”† i64 โ”† i64 โ”† i64 โ”† u32 โ”† i64 โ”† f64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + winners = ( + best_candidates.sort(["dim_0_candidate", "dim_1_candidate", "lottery"]) + .group_by(["dim_0_candidate", "dim_1_candidate"], maintain_order=True) + .first() + ) + + assigned = pl.concat( + [ + assigned, + winners.select( + [ + "agent_id", + pl.col("dim_0_candidate"), + pl.col("dim_1_candidate"), + ] + ), + ], + how="vertical", + ) + taken = pl.concat( + [ + taken, + winners.select(["dim_0_candidate", "dim_1_candidate"]), + ], + how="vertical", + ) + + winner_ids = winners.select("agent_id") + unresolved = unresolved.join(winner_ids, on="agent_id", how="anti") + if unresolved.is_empty(): + break + + # loser candidates columns mirror best_candidates (minus winners). + losers = best_candidates.join(winner_ids, on="agent_id", how="anti") + if losers.is_empty(): + continue + + # loser_updates columns (after select): + # โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + # โ”‚ agent_id โ”† next_rank โ”‚ + # โ”‚ --- โ”† --- โ”‚ + # โ”‚ u64 โ”† i64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + loser_updates = ( + losers.select( + "agent_id", + (pl.col("rank") + 1).cast(pl.Int64).alias("next_rank"), + ) + .join(max_rank, on="agent_id", how="left") + .with_columns( + pl.min_horizontal(pl.col("next_rank"), pl.col("max_rank")).alias( + "next_rank" + ) + ) + .select(["agent_id", "next_rank"]) + ) + + # Promote losers' current_rank (if any) and continue. + # unresolved (updated) retains columns agent_id/current_rank. + unresolved = ( + unresolved.join(loser_updates, on="agent_id", how="left") + .with_columns( + pl.when(pl.col("next_rank").is_not_null()) + .then(pl.col("next_rank")) + .otherwise(pl.col("current_rank")) + .alias("current_rank") + ) + .drop("next_rank") + ) + + return assigned + + +# %% [markdown] +""" +## 4. Run the Model Variants + +We iterate over each movement policy with a shared helper so all runs reuse the same seed. The tutorial runs all three variants (Python sequential, Numba sequential, and parallel) by default; edit the script if you want to skip the slow pure-Python baseline. + +""" + +# %% + +GRID_WIDTH = 40 +GRID_HEIGHT = 40 +NUM_AGENTS = 400 +MODEL_STEPS = 60 +MAX_SUGAR = 4 +SEED = 42 + + +def run_variant( + agent_cls: type[AntsBase], + *, + steps: int, + seed: int, +) -> tuple[Sugarscape, float]: + model = Sugarscape( + agent_type=agent_cls, + n_agents=NUM_AGENTS, + width=GRID_WIDTH, + height=GRID_HEIGHT, + max_sugar=MAX_SUGAR, + seed=seed, + ) + start = perf_counter() + model.run(steps) + return model, perf_counter() - start + + +variant_specs: dict[str, type[AntsBase]] = { + "Sequential (Python loop)": AntsSequential, + "Sequential (Numba)": AntsNumba, + "Parallel (Polars)": AntsParallel, +} + +models: dict[str, Sugarscape] = {} +frames: dict[str, pl.DataFrame] = {} +runtimes: dict[str, float] = {} + +for variant_name, agent_cls in variant_specs.items(): + model, runtime = run_variant(agent_cls, steps=MODEL_STEPS, seed=SEED) + models[variant_name] = model + frames[variant_name] = model.datacollector.data["model"] + runtimes[variant_name] = runtime + + print(f"{variant_name} aggregate trajectory (last 5 steps):") + print( + frames[variant_name] + .select(["step", "mean_sugar", "total_sugar", "agents_alive"]) + .tail(5) + ) + print(f"{variant_name} runtime: {runtime:.3f} s") + print() + +runtime_table = ( + pl.DataFrame( + [ + { + "update_rule": variant_name, + "runtime_seconds": runtimes.get(variant_name, float("nan")), + } + for variant_name in variant_specs.keys() + ] + ) + .with_columns(pl.col("runtime_seconds").round(4)) + .sort("runtime_seconds", descending=False, nulls_last=True) +) + +print("Runtime comparison (fastest first):") +print(runtime_table) + +# Access models/frames on demand; keep namespace minimal. +numba_model_frame = frames.get("Sequential (Numba)", pl.DataFrame()) +par_model_frame = frames.get("Parallel (Polars)", pl.DataFrame()) + + +# %% [markdown] +""" +## 5. Comparing the Update Rules + +Even though micro rules differ, aggregate trajectories remain qualitatively similar (sugar trends up while population gradually declines). +When we join the traces step-by-step, we see small but noticeable deviations introduced by synchronous conflict resolution (e.g., a few more retirements when conflicts cluster). +In our run (seed=42), the final-step Gini differs by โ‰ˆ0.005, and wealthโ€“trait correlations match within ~1e-3. +These gaps vary by seed and grid size, but they consistently stay modest, supporting the relaxed parallel update as a faithful macro-level approximation.""" + +# %% +comparison = numba_model_frame.select( + ["step", "mean_sugar", "total_sugar", "agents_alive"] +).join( + par_model_frame.select(["step", "mean_sugar", "total_sugar", "agents_alive"]), + on="step", + how="inner", + suffix="_parallel", +) +comparison = comparison.with_columns( + (pl.col("mean_sugar") - pl.col("mean_sugar_parallel")).abs().alias("mean_diff"), + (pl.col("total_sugar") - pl.col("total_sugar_parallel")).abs().alias("total_diff"), + (pl.col("agents_alive") - pl.col("agents_alive_parallel")) + .abs() + .alias("count_diff"), +) +print("Step-level absolute differences (first 10 steps):") +print(comparison.select(["step", "mean_diff", "total_diff", "count_diff"]).head(10)) + + +# Build the steadyโ€‘state metrics table from the DataCollector output rather than +# recomputing reporters directly on the model objects. The collector already +# stored the modelโ€‘level reporters (gini, correlations, etc.) every step. +def _last_row(df: pl.DataFrame) -> pl.DataFrame: + if df.is_empty(): + return df + # Ensure we take the final time step in case steps < MODEL_STEPS due to extinction. + return df.sort("step").tail(1) + + +numba_last = _last_row(frames.get("Sequential (Numba)", pl.DataFrame())) +parallel_last = _last_row(frames.get("Parallel (Polars)", pl.DataFrame())) + +metrics_pieces: list[pl.DataFrame] = [] +if not numba_last.is_empty(): + metrics_pieces.append( + numba_last.select( + [ + pl.lit("Sequential (Numba)").alias("update_rule"), + "gini", + "corr_sugar_metabolism", + "corr_sugar_vision", + pl.col("agents_alive"), + ] + ) + ) +if not parallel_last.is_empty(): + metrics_pieces.append( + parallel_last.select( + [ + pl.lit("Parallel (random tie-break)").alias("update_rule"), + "gini", + "corr_sugar_metabolism", + "corr_sugar_vision", + pl.col("agents_alive"), + ] + ) + ) + +metrics_table = ( + pl.concat(metrics_pieces, how="vertical") if metrics_pieces else pl.DataFrame() +) + +print("\nSteady-state inequality metrics:") +print( + metrics_table.select( + [ + "update_rule", + pl.col("gini").round(4), + pl.col("corr_sugar_metabolism").round(4), + pl.col("corr_sugar_vision").round(4), + pl.col("agents_alive"), + ] + ) +) + +if metrics_table.height >= 2: + numba_gini = metrics_table.filter(pl.col("update_rule") == "Sequential (Numba)")[ + "gini" + ][0] + par_gini = metrics_table.filter( + pl.col("update_rule") == "Parallel (random tie-break)" + )["gini"][0] + print(f"Absolute Gini gap (numba vs parallel): {abs(numba_gini - par_gini):.4f}") + +# %% [markdown] +""" +## 6. Takeaways and Next Steps + +Some final notes: +- mesa-frames should preferably be used when you have many agents and operations can be vectorized. +- If your model is not easily vectorizable, consider using Numba or reducing your microscopic rule to a vectorizable form. As we saw, the macroscopic behavior can remain consistent (and be more similar to real-world systems). + + +Currently, the Polars implementation spends most of the time in join operations. + +**Polars + LazyFrames roadmap** โ€“ future mesa-frames releases will expose + LazyFrame-powered sets and spaces (which can also use a GPU cuda accelerated backend which greatly accelerates joins), so the same Polars + code you wrote here will scale even further without touching Numba. +""" diff --git a/docs/general/user-guide/4_datacollector.ipynb b/docs/general/user-guide/4_datacollector.ipynb deleted file mode 100644 index 0809caa2..00000000 --- a/docs/general/user-guide/4_datacollector.ipynb +++ /dev/null @@ -1,501 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7fb27b941602401d91542211134fc71a", - "metadata": {}, - "source": [ - "# Data Collector Tutorial\n", - "\n", - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/4_datacollector.ipynb)\n", - "\n", - "This notebook walks you through using the concrete `DataCollector` in `mesa-frames` to collect model- and agent-level data and write it to different storage backends: **memory, CSV, Parquet, S3, and PostgreSQL**.\n", - "\n", - "It also shows how to use **conditional triggers** and how the **schema validation** behaves for PostgreSQL.\n" - ] - }, - { - "cell_type": "markdown", - "id": "acae54e37e7d407bbb7b55eff062a284", - "metadata": {}, - "source": [ - "## Installation (Colab or fresh env)\n", - "\n", - "Uncomment and run the next cell if you're in Colab or a clean environment.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "9a63283cbaf04dbcab1f6479b197f3a8", - "metadata": { - "editable": true - }, - "outputs": [], - "source": [ - "# !pip install git+https://github.com/projectmesa/mesa-frames mesa" - ] - }, - { - "cell_type": "markdown", - "id": "8dd0d8092fe74a7c96281538738b07e2", - "metadata": {}, - "source": [ - "## Minimal Example Model\n", - "\n", - "We create a tiny model using the `Model` and an `AgentSet`-style agent container. This is just to demonstrate collection APIs.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "72eea5119410473aa328ad9291626812", - "metadata": { - "editable": true - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{'model': shape: (5, 5)\n", - " โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”\n", - " โ”‚ step โ”† seed โ”† batch โ”† total_wealth โ”† n_agents โ”‚\n", - " โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚\n", - " โ”‚ i64 โ”† str โ”† i64 โ”† f64 โ”† i64 โ”‚\n", - " โ•žโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก\n", - " โ”‚ 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”† 1000.0 โ”† 1000 โ”‚\n", - " โ”‚ 4 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”† 1000.0 โ”† 1000 โ”‚\n", - " โ”‚ 6 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”† 1000.0 โ”† 1000 โ”‚\n", - " โ”‚ 8 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”† 1000.0 โ”† 1000 โ”‚\n", - " โ”‚ 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”† 1000.0 โ”† 1000 โ”‚\n", - " โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜,\n", - " 'agent': shape: (5_000, 4)\n", - " โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”\n", - " โ”‚ wealth_MoneyAgents โ”† step โ”† seed โ”† batch โ”‚\n", - " โ”‚ --- โ”† --- โ”† --- โ”† --- โ”‚\n", - " โ”‚ f64 โ”† i32 โ”† str โ”† i32 โ”‚\n", - " โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•ก\n", - " โ”‚ 3.0 โ”† 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 2.0 โ”† 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 1.0 โ”† 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 2 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ โ€ฆ โ”† โ€ฆ โ”† โ€ฆ โ”† โ€ฆ โ”‚\n", - " โ”‚ 0.0 โ”† 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ”‚ 0.0 โ”† 10 โ”† 332212815818606584686857770936โ€ฆ โ”† 0 โ”‚\n", - " โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜}" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from mesa_frames import Model, AgentSet, DataCollector\n", - "import polars as pl\n", - "\n", - "\n", - "class MoneyAgents(AgentSet):\n", - " def __init__(self, n: int, model: Model):\n", - " super().__init__(model)\n", - " # one column, one unit of wealth each\n", - " self += pl.DataFrame({\"wealth\": pl.ones(n, eager=True)})\n", - "\n", - " def step(self) -> None:\n", - " self.select(self.wealth > 0)\n", - " receivers = self.df.sample(n=len(self.active_agents), with_replacement=True)\n", - " self[\"active\", \"wealth\"] -= 1\n", - " income = receivers.group_by(\"unique_id\").len()\n", - " self[income[\"unique_id\"], \"wealth\"] += income[\"len\"]\n", - "\n", - "\n", - "class MoneyModel(Model):\n", - " def __init__(self, n: int):\n", - " super().__init__()\n", - " self.sets.add(MoneyAgents(n, self))\n", - " self.dc = DataCollector(\n", - " model=self,\n", - " model_reporters={\n", - " \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n", - " \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n", - " },\n", - " agent_reporters={\n", - " \"wealth\": \"wealth\", # pull existing column\n", - " },\n", - " storage=\"memory\", # we'll switch this per example\n", - " storage_uri=None,\n", - " trigger=lambda m: m.steps % 2\n", - " == 0, # collect every 2 steps via conditional_collect\n", - " reset_memory=True,\n", - " )\n", - "\n", - " def step(self):\n", - " self.sets.do(\"step\")\n", - "\n", - " def run(self, steps: int, conditional: bool = True):\n", - " for _ in range(steps):\n", - " self.step()\n", - " self.dc.conditional_collect() # or .collect if you want to collect every step regardless of trigger\n", - "\n", - "\n", - "model = MoneyModel(1000)\n", - "model.run(10)\n", - "model.dc.data # peek in-memory dataframes" - ] - }, - { - "cell_type": "markdown", - "id": "3d3ca41d", - "metadata": {}, - "source": [ - "## Saving the data for later use \n", - "\n", - "`DataCollector` supports multiple storage backends. \n", - "Files are saved with **step number** and **batch number** (e.g., `model_step10_batch2.csv`) so multiple collects at the same step donโ€™t overwrite. \n", - " \n", - "- **CSV:** `storage=\"csv\"` โ†’ writes `model_step{n}_batch{k}.csv`, easy to open anywhere. \n", - "- **Parquet:** `storage=\"parquet\"` โ†’ compressed, efficient for large datasets. \n", - "- **S3:** `storage=\"S3-csv\"`/`storage=\"S3-parquet\"` โ†’ saves CSV/Parquet directly to Amazon S3. \n", - "- **PostgreSQL:** `storage=\"postgresql\"` โ†’ inserts results into `model_data` and `agent_data` tables for querying. \n" - ] - }, - { - "cell_type": "markdown", - "id": "8edb47106e1a46a883d545849b8ab81b", - "metadata": {}, - "source": [ - "## Writing to Local CSV\n", - "\n", - "Switch the storage to `csv` and provide a folder path. Files are written as `model_step{n}.csv` and `agent_step{n}.csv`.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5f14f38c", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "[]" - ] - }, - "execution_count": 8, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import os\n", - "\n", - "os.makedirs(\"./data_csv\", exist_ok=True)\n", - "model_csv = MoneyModel(1000)\n", - "model_csv.dc = DataCollector(\n", - " model=model_csv,\n", - " model_reporters={\n", - " \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n", - " \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n", - " },\n", - " agent_reporters={\n", - " \"wealth\": \"wealth\",\n", - " },\n", - " storage=\"csv\", # saving as csv\n", - " storage_uri=\"./data_csv\",\n", - " trigger=lambda m: m._steps % 2 == 0,\n", - " reset_memory=True,\n", - ")\n", - "model_csv.run(10)\n", - "model_csv.dc.flush()\n", - "os.listdir(\"./data_csv\")" - ] - }, - { - "cell_type": "markdown", - "id": "10185d26023b46108eb7d9f57d49d2b3", - "metadata": {}, - "source": [ - "## Writing to Local Parquet\n", - "\n", - "Use `parquet` for columnar output.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8763a12b2bbd4a93a75aff182afb95dc", - "metadata": { - "editable": true - }, - "outputs": [ - { - "data": { - "text/plain": [ - "[]" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "os.makedirs(\"./data_parquet\", exist_ok=True)\n", - "model_parq = MoneyModel(1000)\n", - "model_parq.dc = DataCollector(\n", - " model=model_parq,\n", - " model_reporters={\n", - " \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n", - " \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n", - " },\n", - " agent_reporters={\n", - " \"wealth\": \"wealth\",\n", - " },\n", - " storage=\"parquet\", # save as parquet\n", - " storage_uri=\"data_parquet\",\n", - " trigger=lambda m: m._steps % 2 == 0,\n", - " reset_memory=True,\n", - ")\n", - "model_parq.run(10)\n", - "model_parq.dc.flush()\n", - "os.listdir(\"./data_parquet\")" - ] - }, - { - "cell_type": "markdown", - "id": "7623eae2785240b9bd12b16a66d81610", - "metadata": {}, - "source": [ - "## Writing to Amazon S3 (CSV or Parquet)\n", - "\n", - "Set AWS credentials via environment variables or your usual config. Then choose `S3-csv` or `S3-parquet` and pass an S3 URI (e.g., `s3://my-bucket/experiments/run-1`).\n", - "\n", - "> **Note:** This cell requires network access & credentials when actually run.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7cdc8c89c7104fffa095e18ddfef8986", - "metadata": { - "editable": true - }, - "outputs": [], - "source": [ - "model_s3 = MoneyModel(1000)\n", - "model_s3.dc = DataCollector(\n", - " model=model_s3,\n", - " model_reporters={\n", - " \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n", - " \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n", - " },\n", - " agent_reporters={\n", - " \"wealth\": \"wealth\",\n", - " },\n", - " storage=\"S3-csv\", # save as csv in S3\n", - " storage_uri=\"s3://my-bucket/experiments/run-1\", # change it to required path\n", - " trigger=lambda m: m._steps % 2 == 0,\n", - " reset_memory=True,\n", - ")\n", - "model_s3.run(10)\n", - "model_s3.dc.flush()" - ] - }, - { - "cell_type": "markdown", - "id": "b118ea5561624da68c537baed56e602f", - "metadata": {}, - "source": [ - "## Writing to PostgreSQL\n", - "\n", - "PostgreSQL requires that the target tables exist and that the expected reporter columns are present. The collector will validate tables/columns up front and raise descriptive errors if something is missing.\n", - "\n", - "Below is a minimal schema example. Adjust columns to your configured reporters.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "938c804e27f84196a10c8828c723f798", - "metadata": { - "editable": true - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "CREATE SCHEMA IF NOT EXISTS public;\n", - "CREATE TABLE IF NOT EXISTS public.model_data (\n", - " step INTEGER,\n", - " seed VARCHAR,\n", - " total_wealth BIGINT,\n", - " n_agents INTEGER\n", - ");\n", - "\n", - "\n", - "CREATE TABLE IF NOT EXISTS public.agent_data (\n", - " step INTEGER,\n", - " seed VARCHAR,\n", - " unique_id BIGINT,\n", - " wealth BIGINT\n", - ");\n", - "\n" - ] - } - ], - "source": [ - "DDL_MODEL = r\"\"\"\n", - "CREATE SCHEMA IF NOT EXISTS public;\n", - "CREATE TABLE IF NOT EXISTS public.model_data (\n", - " step INTEGER,\n", - " seed VARCHAR,\n", - " total_wealth BIGINT,\n", - " n_agents INTEGER\n", - ");\n", - "\"\"\"\n", - "DDL_AGENT = r\"\"\"\n", - "CREATE TABLE IF NOT EXISTS public.agent_data (\n", - " step INTEGER,\n", - " seed VARCHAR,\n", - " unique_id BIGINT,\n", - " wealth BIGINT\n", - ");\n", - "\"\"\"\n", - "print(DDL_MODEL)\n", - "print(DDL_AGENT)" - ] - }, - { - "cell_type": "markdown", - "id": "504fb2a444614c0babb325280ed9130a", - "metadata": {}, - "source": [ - "After creating the tables (outside this notebook or via a DB connection cell), configure and flush:\n" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "59bbdb311c014d738909a11f9e486628", - "metadata": { - "editable": true - }, - "outputs": [], - "source": [ - "POSTGRES_URI = \"postgresql://user:pass@localhost:5432/mydb\"\n", - "m_pg = MoneyModel(300)\n", - "m_pg.dc._storage = \"postgresql\"\n", - "m_pg.dc._storage_uri = POSTGRES_URI\n", - "m_pg.run(6)\n", - "m_pg.dc.flush()" - ] - }, - { - "cell_type": "markdown", - "id": "b43b363d81ae4b689946ece5c682cd59", - "metadata": {}, - "source": [ - "## Triggers & Conditional Collection\n", - "\n", - "The collector accepts a `trigger: Callable[[Model], bool]`. When using `conditional_collect()`, the collector checks the trigger and collects only if it returns `True`.\n", - "\n", - "You can always call `collect()` to gather data unconditionally.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "8a65eabff63a45729fe45fb5ade58bdc", - "metadata": { - "editable": true - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "shape: (5, 5)
stepseedbatchtotal_wealthn_agents
i64stri64f64i64
2"540832786058427425452319829502โ€ฆ0100.0100
4"540832786058427425452319829502โ€ฆ0100.0100
6"540832786058427425452319829502โ€ฆ0100.0100
8"540832786058427425452319829502โ€ฆ0100.0100
10"540832786058427425452319829502โ€ฆ0100.0100
" - ], - "text/plain": [ - "shape: (5, 5)\n", - "โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”\n", - "โ”‚ step โ”† seed โ”† batch โ”† total_wealth โ”† n_agents โ”‚\n", - "โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”‚\n", - "โ”‚ i64 โ”† str โ”† i64 โ”† f64 โ”† i64 โ”‚\n", - "โ•žโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก\n", - "โ”‚ 2 โ”† 540832786058427425452319829502โ€ฆ โ”† 0 โ”† 100.0 โ”† 100 โ”‚\n", - "โ”‚ 4 โ”† 540832786058427425452319829502โ€ฆ โ”† 0 โ”† 100.0 โ”† 100 โ”‚\n", - "โ”‚ 6 โ”† 540832786058427425452319829502โ€ฆ โ”† 0 โ”† 100.0 โ”† 100 โ”‚\n", - "โ”‚ 8 โ”† 540832786058427425452319829502โ€ฆ โ”† 0 โ”† 100.0 โ”† 100 โ”‚\n", - "โ”‚ 10 โ”† 540832786058427425452319829502โ€ฆ โ”† 0 โ”† 100.0 โ”† 100 โ”‚\n", - "โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "m = MoneyModel(100)\n", - "m.dc.trigger = lambda model: model._steps % 3 == 0 # every 3rd step\n", - "m.run(10, conditional=True)\n", - "m.dc.data[\"model\"].head()" - ] - }, - { - "cell_type": "markdown", - "id": "c3933fab20d04ec698c2621248eb3be0", - "metadata": {}, - "source": [ - "## Troubleshooting\n", - "\n", - "- **ValueError: Please define a storage_uri** โ€” for non-memory backends you must set `_storage_uri`.\n", - "- **Missing columns in table** โ€” check the PostgreSQL error text; create/alter the table to include the columns for your configured `model_reporters` and `agent_reporters`, plus required `step` and `seed`.\n", - "- **Permissions/credentials errors** (S3/PostgreSQL) โ€” ensure correct IAM/credentials or database permissions.\n" - ] - }, - { - "cell_type": "markdown", - "id": "4dd4641cc4064e0191573fe9c69df29b", - "metadata": {}, - "source": [ - "---\n", - "*Generated on 2025-08-30.*\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "mesa-frames (3.12.3)", - "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.3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/general/user-guide/4_datacollector.py b/docs/general/user-guide/4_datacollector.py new file mode 100644 index 00000000..16d9837b --- /dev/null +++ b/docs/general/user-guide/4_datacollector.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +# %% [markdown] +"""# Data Collector Tutorial + +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa-frames/blob/main/docs/general/user-guide/4_datacollector.ipynb) + +This notebook walks you through using the concrete `DataCollector` in `mesa-frames` to collect model- and agent-level data and write it to different storage backends: **memory, CSV, Parquet, S3, and PostgreSQL**. + +It also shows how to use **conditional triggers** and how the **schema validation** behaves for PostgreSQL.""" + +# %% [markdown] +"""## Installation (Colab or fresh env) + +Uncomment and run the next cell if you're in Colab or a clean environment.""" + +# %% +# !pip install git+https://github.com/projectmesa/mesa-frames mesa + +# %% [markdown] +"""## Minimal Example Model + +We create a tiny model using the `Model` and an `AgentSet`-style agent container. This is just to demonstrate collection APIs.""" + +# %% +from mesa_frames import Model, AgentSet, DataCollector +import polars as pl + + +class MoneyAgents(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + # one column, one unit of wealth each + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) + + def step(self) -> None: + self.select(self.wealth > 0) + receivers = self.df.sample(n=len(self.active_agents), with_replacement=True) + self["active", "wealth"] -= 1 + income = receivers.group_by("unique_id").len() + self[income["unique_id"], "wealth"] += income["len"] + + +class MoneyModel(Model): + def __init__(self, n: int): + super().__init__() + self.sets.add(MoneyAgents(n, self)) + self.dc = DataCollector( + model=self, + model_reporters={ + "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum(), + "n_agents": lambda m: len(m.sets["MoneyAgents"]), + }, + agent_reporters={ + "wealth": "wealth", # pull existing column + }, + storage="memory", # we'll switch this per example + storage_uri=None, + trigger=lambda m: m.steps % 2 + == 0, # collect every 2 steps via conditional_collect + reset_memory=True, + ) + + def step(self): + self.sets.do("step") + + def run(self, steps: int, conditional: bool = True): + for _ in range(steps): + self.step() + self.dc.conditional_collect() # or .collect if you want to collect every step regardless of trigger + + +model = MoneyModel(1000) +model.run(10) +model.dc.data # peek in-memory dataframes + +# %% [markdown] +"""## Saving the data for later use + +`DataCollector` supports multiple storage backends. +Files are saved with **step number** and **batch number** (e.g., `model_step10_batch2.csv`) so multiple collects at the same step donโ€™t overwrite. + +- **CSV:** `storage="csv"` โ†’ writes `model_step{n}_batch{k}.csv`, easy to open anywhere. +- **Parquet:** `storage="parquet"` โ†’ compressed, efficient for large datasets. +- **S3:** `storage="S3-csv"`/`storage="S3-parquet"` โ†’ saves CSV/Parquet directly to Amazon S3. +- **PostgreSQL:** `storage="postgresql"` โ†’ inserts results into `model_data` and `agent_data` tables for querying.""" + +# %% [markdown] +"""## Writing to Local CSV + +Switch the storage to `csv` and provide a folder path. Files are written as `model_step{n}.csv` and `agent_step{n}.csv`.""" + +# %% +import os + +os.makedirs("./data_csv", exist_ok=True) +model_csv = MoneyModel(1000) +model_csv.dc = DataCollector( + model=model_csv, + model_reporters={ + "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum(), + "n_agents": lambda m: len(m.sets["MoneyAgents"]), + }, + agent_reporters={ + "wealth": "wealth", + }, + storage="csv", # saving as csv + storage_uri="./data_csv", + trigger=lambda m: m._steps % 2 == 0, + reset_memory=True, +) +model_csv.run(10) +model_csv.dc.flush() +os.listdir("./data_csv") + +# %% [markdown] +"""## Writing to Local Parquet + +Use `parquet` for columnar output.""" + +# %% +os.makedirs("./data_parquet", exist_ok=True) +model_parq = MoneyModel(1000) +model_parq.dc = DataCollector( + model=model_parq, + model_reporters={ + "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum(), + "n_agents": lambda m: len(m.sets["MoneyAgents"]), + }, + agent_reporters={ + "wealth": "wealth", + }, + storage="parquet", # save as parquet + storage_uri="data_parquet", + trigger=lambda m: m._steps % 2 == 0, + reset_memory=True, +) +model_parq.run(10) +model_parq.dc.flush() +os.listdir("./data_parquet") + +# %% [markdown] +"""## Writing to Amazon S3 (CSV or Parquet) + +Set AWS credentials via environment variables or your usual config. Then choose `S3-csv` or `S3-parquet` and pass an S3 URI (e.g., `s3://my-bucket/experiments/run-1`). + +> **Note:** This cell requires network access & credentials when actually run.""" + +# %% +model_s3 = MoneyModel(1000) +model_s3.dc = DataCollector( + model=model_s3, + model_reporters={ + "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum(), + "n_agents": lambda m: len(m.sets["MoneyAgents"]), + }, + agent_reporters={ + "wealth": "wealth", + }, + storage="S3-csv", # save as csv in S3 + storage_uri="s3://my-bucket/experiments/run-1", # change it to required path + trigger=lambda m: m._steps % 2 == 0, + reset_memory=True, +) +model_s3.run(10) +model_s3.dc.flush() + +# %% [markdown] +"""## Writing to PostgreSQL + +PostgreSQL requires that the target tables exist and that the expected reporter columns are present. The collector will validate tables/columns up front and raise descriptive errors if something is missing. + +Below is a minimal schema example. Adjust columns to your configured reporters.""" + +# %% +DDL_MODEL = r""" +CREATE SCHEMA IF NOT EXISTS public; +CREATE TABLE IF NOT EXISTS public.model_data ( + step INTEGER, + seed VARCHAR, + total_wealth BIGINT, + n_agents INTEGER +); +""" +DDL_AGENT = r""" +CREATE TABLE IF NOT EXISTS public.agent_data ( + step INTEGER, + seed VARCHAR, + unique_id BIGINT, + wealth BIGINT +); +""" +print(DDL_MODEL) +print(DDL_AGENT) + +# %% [markdown] +"""After creating the tables (outside this notebook or via a DB connection cell), configure and flush:""" + +# %% +POSTGRES_URI = "postgresql://user:pass@localhost:5432/mydb" +m_pg = MoneyModel(300) +m_pg.dc._storage = "postgresql" +m_pg.dc._storage_uri = POSTGRES_URI +m_pg.run(6) +m_pg.dc.flush() + +# %% [markdown] +"""## Triggers & Conditional Collection + +The collector accepts a `trigger: Callable[[Model], bool]`. When using `conditional_collect()`, the collector checks the trigger and collects only if it returns `True`. + +You can always call `collect()` to gather data unconditionally.""" + +# %% +m = MoneyModel(100) +m.dc.trigger = lambda model: model._steps % 3 == 0 # every 3rd step +m.run(10, conditional=True) +m.dc.data["model"].head() + +# %% [markdown] +"""## Troubleshooting + +- **ValueError: Please define a storage_uri** โ€” for non-memory backends you must set `_storage_uri`. +- **Missing columns in table** โ€” check the PostgreSQL error text; create/alter the table to include the columns for your configured `model_reporters` and `agent_reporters`, plus required `step` and `seed`. +- **Permissions/credentials errors** (S3/PostgreSQL) โ€” ensure correct IAM/credentials or database permissions.""" + +# %% [markdown] +"""--- +*Generated on 2025-08-30.*""" diff --git a/docs/general/user-guide/5_benchmarks.md b/docs/general/user-guide/5_benchmarks.md index 61fca87b..233c394c 100644 --- a/docs/general/user-guide/5_benchmarks.md +++ b/docs/general/user-guide/5_benchmarks.md @@ -8,11 +8,11 @@ mesa-frames offers significant performance improvements over the original mesa f ### Comparison with mesa -![Performance Graph BW](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_no_mesa.png) +![Performance Graph BW](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_with_mesa.png) ### Comparison of mesa-frames implementations -![Performance Graph BW without Mesa](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_with_mesa.png) +![Performance Graph BW without Mesa](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_no_mesa.png) ## SugarScape with Instantaneous Growback ๐Ÿฌ diff --git a/examples/sugarscape_ig/ss_polars/__init__.py b/docs/stylesheets/brand-material.css similarity index 100% rename from examples/sugarscape_ig/ss_polars/__init__.py rename to docs/stylesheets/brand-material.css diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 00000000..069e9dc5 --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1,6 @@ +"""Examples package for the repository.""" + +__all__ = [ + "boltzmann_wealth", + "sugarscape_ig", +] diff --git a/examples/boltzmann_wealth/__init__.py b/examples/boltzmann_wealth/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/boltzmann_wealth/backend_frames.py b/examples/boltzmann_wealth/backend_frames.py new file mode 100644 index 00000000..23efac92 --- /dev/null +++ b/examples/boltzmann_wealth/backend_frames.py @@ -0,0 +1,163 @@ +"""Mesa-frames implementation of the Boltzmann wealth model with Typer CLI.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path +from typing import Annotated + +import numpy as np +import polars as pl +import typer +from time import perf_counter + +from mesa_frames import AgentSet, DataCollector, Model +from examples.utils import FramesSimulationResult +from examples.plotting import plot_model_metrics + + +# Note: by default we create a timestamped results directory under `results/`. +# The CLI will accept optional `results_dir` and `plots_dir` arguments to override. + + +def gini(frame: pl.DataFrame) -> float: + wealth = frame["wealth"] if "wealth" in frame.columns else pl.Series([]) + if wealth.is_empty(): + return float("nan") + values = wealth.to_numpy().astype(np.float64) + if values.size == 0: + return float("nan") + if np.allclose(values, 0.0): + return 0.0 + if np.allclose(values, values[0]): + return 0.0 + sorted_vals = np.sort(values) + n = sorted_vals.size + cumulative = np.cumsum(sorted_vals) + total = cumulative[-1] + if total == 0: + return 0.0 + index = np.arange(1, n + 1, dtype=np.float64) + return float((2.0 * np.dot(index, sorted_vals) / (n * total)) - (n + 1) / n) + + +class MoneyAgents(AgentSet): + """Vectorised agent set for the Boltzmann wealth exchange model.""" + + def __init__(self, model: Model, agents: int) -> None: + super().__init__(model) + self += pl.DataFrame({"wealth": pl.Series(np.ones(agents, dtype=np.int64))}) + + def step(self) -> None: + self.select(pl.col("wealth") > 0) + if len(self.active_agents) == 0: + return + # Use the model RNG to seed Polars sampling so results are reproducible + recipients = self.df.sample( + n=len(self.active_agents), + with_replacement=True, + seed=self.random.integers(np.iinfo(np.int32).max), + ) + # donors lose one unit + self["active", "wealth"] -= 1 + gains = recipients.group_by("unique_id").len() + self[gains, "wealth"] += gains["len"] + + +class MoneyModel(Model): + """Mesa-frames model that mirrors the Mesa implementation.""" + + def __init__( + self, agents: int, *, seed: int | None = None, results_dir: Path | None = None + ) -> None: + super().__init__(seed) + self.sets += MoneyAgents(self, agents) + storage_uri = str(results_dir) if results_dir is not None else None + self.datacollector = DataCollector( + model=self, + model_reporters={ + "gini": lambda m: gini(m.sets[0].df), + }, + storage="csv", + storage_uri=storage_uri, + ) + + def step(self) -> None: + self.sets.do("step") + self.datacollector.collect() + + def run(self, steps: int) -> None: + for _ in range(steps): + self.step() + + +def simulate( + agents: int, + steps: int, + seed: int | None = None, + results_dir: Path | None = None, +) -> FramesSimulationResult: + model = MoneyModel(agents, seed=seed, results_dir=results_dir) + model.run(steps) + # collect data from datacollector into memory first + return FramesSimulationResult(datacollector=model.datacollector) + + +app = typer.Typer(add_completion=False) + + +@app.command() +def run( + agents: Annotated[int, typer.Option(help="Number of agents to simulate.")] = 5000, + steps: Annotated[int, typer.Option(help="Number of model steps to run.")] = 100, + seed: Annotated[int | None, typer.Option(help="Optional RNG seed.")] = None, + plot: Annotated[bool, typer.Option(help="Render Seaborn plots.")] = True, + save_results: Annotated[bool, typer.Option(help="Persist metrics as CSV.")] = True, + results_dir: Annotated[ + Path | None, + typer.Option( + help="Directory to write CSV results and plots into. If omitted a timestamped subdir under `results/` is used." + ), + ] = None, +) -> None: + typer.echo( + f"Running Boltzmann wealth model (mesa-frames) with {agents} agents for {steps} steps" + ) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + if results_dir is None: + results_dir = ( + Path(__file__).resolve().parent / "results" / timestamp + ).resolve() + results_dir.mkdir(parents=True, exist_ok=True) + start_time = perf_counter() + result = simulate(agents=agents, steps=steps, seed=seed, results_dir=results_dir) + + typer.echo(f"Simulation complete in {perf_counter() - start_time:.2f} seconds") + + model_metrics = result.datacollector.data["model"].select("step", "gini") + + typer.echo(f"Metrics in the final 5 steps: {model_metrics.tail(5)}") + + if save_results: + result.datacollector.flush() + + if plot: + stem = f"gini_{timestamp}" + # write plots into the results directory so outputs are colocated + plot_model_metrics( + model_metrics, + results_dir, + stem, + title="Boltzmann wealth โ€” Gini", + subtitle=f"mesa-frames backend; seed={result.datacollector.seed}", + agents=agents, + steps=steps, + ) + typer.echo(f"Saved plots under {results_dir}") + + # Inform user where CSVs were saved + typer.echo(f"Saved CSV results under {results_dir}") + + +if __name__ == "__main__": + app() diff --git a/examples/boltzmann_wealth/backend_mesa.py b/examples/boltzmann_wealth/backend_mesa.py new file mode 100644 index 00000000..8b86ad3e --- /dev/null +++ b/examples/boltzmann_wealth/backend_mesa.py @@ -0,0 +1,181 @@ +"""Mesa implementation of the Boltzmann wealth model with Typer CLI.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path +from typing import Annotated +from collections.abc import Iterable +import pandas as pd + +import matplotlib.pyplot as plt +import mesa +from mesa.datacollection import DataCollector +import numpy as np +import polars as pl +import seaborn as sns +import typer +from time import perf_counter + +from examples.utils import MesaSimulationResult +from examples.plotting import plot_model_metrics + + +def gini(values: Iterable[float]) -> float: + """Compute the Gini coefficient from an iterable of wealth values.""" + array = np.fromiter(values, dtype=float) + if array.size == 0: + return float("nan") + if np.allclose(array, 0.0): + return 0.0 + if np.allclose(array, array[0]): + return 0.0 + sorted_vals = np.sort(array) + n = sorted_vals.size + cumulative = np.cumsum(sorted_vals) + total = cumulative[-1] + if total == 0: + return 0.0 + index = np.arange(1, n + 1, dtype=float) + return float((2.0 * np.dot(index, sorted_vals) / (n * total)) - (n + 1) / n) + + +class MoneyAgent(mesa.Agent): + """Agent that passes one unit of wealth to a random neighbour.""" + + def __init__(self, model: MoneyModel) -> None: + super().__init__(model) + self.wealth = 1 + + def step(self) -> None: + if self.wealth <= 0: + return + other = self.random.choice(self.model.agent_list) + if other is None: + return + other.wealth += 1 + self.wealth -= 1 + + +class MoneyModel(mesa.Model): + """Mesa backend that mirrors the mesa-frames Boltzmann wealth example.""" + + def __init__(self, agents: int, *, seed: int | None = None) -> None: + super().__init__() + if seed is None: + seed = self.random.randint(0, np.iinfo(np.int32).max) + self.reset_randomizer(seed) + self.agent_list: list[MoneyAgent] = [] + for _ in range(agents): + # NOTE: storing agents in a Python list keeps iteration fast for benchmarks. + agent = MoneyAgent(self) + self.agent_list.append(agent) + self.datacollector = DataCollector( + model_reporters={ + "gini": lambda m: gini(a.wealth for a in m.agent_list), + "seed": lambda m: seed, + } + ) + self.datacollector.collect(self) + + def step(self) -> None: + self.random.shuffle(self.agent_list) + for agent in self.agent_list: + agent.step() + self.datacollector.collect(self) + + def run(self, steps: int) -> None: + for _ in range(steps): + self.step() + + +def simulate(agents: int, steps: int, seed: int | None = None) -> MesaSimulationResult: + """Run the Mesa Boltzmann wealth model.""" + model = MoneyModel(agents, seed=seed) + model.run(steps) + + return MesaSimulationResult(datacollector=model.datacollector) + + +app = typer.Typer(add_completion=False) + + +@app.command() +def run( + agents: Annotated[int, typer.Option(help="Number of agents to simulate.")] = 5000, + steps: Annotated[int, typer.Option(help="Number of model steps to run.")] = 100, + seed: Annotated[int | None, typer.Option(help="Optional RNG seed.")] = None, + plot: Annotated[bool, typer.Option(help="Render plots.")] = True, + save_results: Annotated[ + bool, + typer.Option(help="Persist metrics as CSV."), + ] = True, + results_dir: Annotated[ + Path | None, + typer.Option( + help=( + "Directory to write CSV results and plots into. If omitted a " + "timestamped subdir under `results/` is used." + ) + ), + ] = None, +) -> None: + """Execute the Mesa Boltzmann wealth simulation.""" + + typer.echo( + f"Running Boltzmann wealth model (mesa) with {agents} agents for {steps} steps" + ) + + # Resolve output folder + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + if results_dir is None: + results_dir = ( + Path(__file__).resolve().parent / "results" / timestamp + ).resolve() + results_dir.mkdir(parents=True, exist_ok=True) + + start_time = perf_counter() + # Run simulation (Mesaโ€‘idiomatic): we only use DataCollector's public API + result = simulate(agents=agents, steps=steps, seed=seed) + typer.echo(f"Simulation completed in {perf_counter() - start_time:.3f} seconds") + dc = result.datacollector + + # ---- Extract metrics (no helper, no monkeyโ€‘patch): + # DataCollector returns a pandas DataFrame with the index as the step. + model_pd = dc.get_model_vars_dataframe() + model_pd = model_pd.reset_index() + # The first column is the step index; normalize name to "step". + model_pd = model_pd.rename(columns={model_pd.columns[0]: "step"}) + seed = model_pd["seed"].iloc[0] + model_pd = model_pd[["step", "gini"]] + + # Show a short tail in console for quick inspection + tail_str = model_pd.tail(5).to_string(index=False) + typer.echo(f"Metrics in the final 5 steps:\n{tail_str}") + + # ---- Save CSV (same filename/layout as frames backend expects) + if save_results: + csv_path = results_dir / "model.csv" + model_pd.to_csv(csv_path, index=False) + + # ---- Plot (convert to Polars to reuse the shared plotting helper) + if plot and not model_pd.empty: + model_pl = pl.from_pandas(model_pd) + stem = f"gini_{timestamp}" + plot_model_metrics( + model_pl, + results_dir, + stem, + title="Boltzmann wealth โ€” Gini", + subtitle=f"mesa backend; seed={seed}", + agents=agents, + steps=steps, + ) + typer.echo(f"Saved plots under {results_dir}") + + if save_results: + typer.echo(f"Saved CSV results under {results_dir}") + + +if __name__ == "__main__": + app() diff --git a/examples/boltzmann_wealth/boltzmann_no_mesa.png b/examples/boltzmann_wealth/boltzmann_no_mesa.png deleted file mode 100644 index 369597e2..00000000 Binary files a/examples/boltzmann_wealth/boltzmann_no_mesa.png and /dev/null differ diff --git a/examples/boltzmann_wealth/boltzmann_with_mesa.png b/examples/boltzmann_wealth/boltzmann_with_mesa.png deleted file mode 100644 index 257d5d18..00000000 Binary files a/examples/boltzmann_wealth/boltzmann_with_mesa.png and /dev/null differ diff --git a/examples/boltzmann_wealth/performance_plot.py b/examples/boltzmann_wealth/performance_plot.py deleted file mode 100644 index e565bda3..00000000 --- a/examples/boltzmann_wealth/performance_plot.py +++ /dev/null @@ -1,239 +0,0 @@ -import importlib.metadata - -import matplotlib.pyplot as plt -import mesa -import numpy as np -import perfplot -import polars as pl -import seaborn as sns -from packaging import version - -from mesa_frames import AgentSet, Model - - -### ---------- Mesa implementation ---------- ### -def mesa_implementation(n_agents: int) -> None: - model = MesaMoneyModel(n_agents) - model.run_model(100) - - -class MesaMoneyAgent(mesa.Agent): - """An agent with fixed initial wealth.""" - - def __init__(self, model): - # Pass the parameters to the parent class. - super().__init__(model) - - # Create the agent's variable and set the initial values. - self.wealth = 1 - - def step(self): - # Verify agent has some wealth - if self.wealth > 0: - other_agent = self.random.choice(self.model.agents) - if other_agent is not None: - other_agent.wealth += 1 - self.wealth -= 1 - - -class MesaMoneyModel(mesa.Model): - """A model with some number of agents.""" - - def __init__(self, N): - super().__init__() - self.num_agents = N - for _ in range(self.num_agents): - self.agents.add(MesaMoneyAgent(self)) - - def step(self): - """Advance the model by one step.""" - self.agents.shuffle_do("step") - - def run_model(self, n_steps) -> None: - for _ in range(n_steps): - self.step() - - -"""def compute_gini(model): - agent_wealths = model.sets.get("wealth") - x = sorted(agent_wealths) - N = model.num_agents - B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x)) - return 1 + (1 / N) - 2 * B""" - - -### ---------- Mesa-frames implementation ---------- ### - - -class MoneyAgentsConcise(AgentSet): - def __init__(self, n: int, model: Model): - super().__init__(model) - ## Adding the agents to the agent set - # 1. Changing the agents attribute directly (not recommended, if other agents were added before, they will be lost) - """self.sets = pl.DataFrame( - "wealth": pl.ones(n, eager=True)} - )""" - # 2. Adding the dataframe with add - """self.add( - pl.DataFrame( - { - "wealth": pl.ones(n, eager=True), - } - ) - )""" - # 3. Adding the dataframe with __iadd__ - self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) - - def step(self) -> None: - # The give_money method is called - # self.give_money() - self.do("give_money") - - def give_money(self): - ## Active agents are changed to wealthy agents - # 1. Using the __getitem__ method - # self.select(self["wealth"] > 0) - # 2. Using the fallback __getattr__ method - self.select(self.wealth > 0) - - # Receiving agents are sampled (only native expressions currently supported) - other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) - - # Wealth of wealthy is decreased by 1 - # 1. Using the __setitem__ method with self.active_agents mask - # self[self.active_agents, "wealth"] -= 1 - # 2. Using the __setitem__ method with "active" mask - self["active", "wealth"] -= 1 - - # Compute the income of the other agents (only native expressions currently supported) - new_wealth = other_agents.group_by("unique_id").len() - - # Add the income to the other agents - # 1. Using the set method - """self.set( - attr_names="wealth", - values=pl.col("wealth") + new_wealth["len"], - mask=new_wealth, - )""" - - # 2. Using the __setitem__ method - self[new_wealth, "wealth"] += new_wealth["len"] - - -class MoneyAgentsNative(AgentSet): - def __init__(self, n: int, model: Model): - super().__init__(model) - self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) - - def step(self) -> None: - self.do("give_money") - - def give_money(self): - ## Active agents are changed to wealthy agents - self.select(pl.col("wealth") > 0) - - other_agents = self.df.sample(n=len(self.active_agents), with_replacement=True) - - # Wealth of wealthy is decreased by 1 - self.df = self.df.with_columns( - wealth=pl.when( - pl.col("unique_id").is_in(self.active_agents["unique_id"].implode()) - ) - .then(pl.col("wealth") - 1) - .otherwise(pl.col("wealth")) - ) - - new_wealth = other_agents.group_by("unique_id").len() - - # Add the income to the other agents - self.df = ( - self.df.join(new_wealth, on="unique_id", how="left") - .fill_null(0) - .with_columns(wealth=pl.col("wealth") + pl.col("len")) - .drop("len") - ) - - -class MoneyModel(Model): - def __init__(self, N: int, agents_cls): - super().__init__() - self.n_agents = N - self.sets += agents_cls(N, self) - - def step(self): - # Executes the step method for every agentset in self.sets - self.sets.do("step") - - def run_model(self, n): - for _ in range(n): - self.step() - - -def mesa_frames_polars_concise(n_agents: int) -> None: - model = MoneyModel(n_agents, MoneyAgentsConcise) - model.run_model(100) - - -def mesa_frames_polars_native(n_agents: int) -> None: - model = MoneyModel(n_agents, MoneyAgentsNative) - model.run_model(100) - - -def plot_and_print_benchmark(labels, kernels, n_range, title, image_path): - out = perfplot.bench( - setup=lambda n: n, - kernels=kernels, - labels=labels, - n_range=n_range, - xlabel="Number of agents", - equality_check=None, - title=title, - ) - - plt.ylabel("Execution time (s)") - out.save(image_path, transparent=False) - - print("\nExecution times:") - for i, label in enumerate(labels): - print(f"---------------\n{label}:") - for n, t in zip(out.n_range, out.timings_s[i]): - print(f" Number of agents: {n}, Time: {t:.2f} seconds") - print("---------------") - - -def main(): - sns.set_theme(style="whitegrid") - - labels_0 = [ - "mesa", - "mesa-frames (pl concise)", - "mesa-frames (pl native)", - ] - kernels_0 = [ - mesa_implementation, - mesa_frames_polars_concise, - mesa_frames_polars_native, - ] - n_range_0 = [k for k in range(0, 100001, 10000)] - title_0 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_0) - image_path_0 = "boltzmann_with_mesa.png" - - plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0) - - labels_1 = [ - "mesa-frames (pl concise)", - "mesa-frames (pl native)", - ] - kernels_1 = [ - mesa_frames_polars_concise, - mesa_frames_polars_native, - ] - n_range_1 = [k for k in range(100000, 1000001, 100000)] - title_1 = "100 steps of the Boltzmann Wealth model:\n" + " vs ".join(labels_1) - image_path_1 = "boltzmann_no_mesa.png" - - plot_and_print_benchmark(labels_1, kernels_1, n_range_1, title_1, image_path_1) - - -if __name__ == "__main__": - main() diff --git a/examples/plotting.py b/examples/plotting.py new file mode 100644 index 00000000..0cf002c4 --- /dev/null +++ b/examples/plotting.py @@ -0,0 +1,285 @@ +# examples/plotting.py +from __future__ import annotations + +from pathlib import Path +from collections.abc import Sequence +import re + +import polars as pl +import seaborn as sns +import matplotlib.pyplot as plt +from matplotlib.ticker import FormatStrFormatter +from matplotlib.figure import Figure +from matplotlib.axes import Axes + +# ----------------------------- Shared theme ---------------------------------- + +_THEMES = { + "light": dict( + style="whitegrid", + rc={ + "axes.spines.top": False, + "axes.spines.right": False, + }, + ), + "dark": dict( + style="whitegrid", + rc={ + # real dark background + readable foreground + "figure.facecolor": "#0b1021", + "axes.facecolor": "#0b1021", + "axes.edgecolor": "#d6d6d7", + "axes.labelcolor": "#e8e8ea", + "text.color": "#e8e8ea", + "xtick.color": "#c9c9cb", + "ytick.color": "#c9c9cb", + "grid.color": "#2a2f4a", + "grid.alpha": 0.35, + "axes.spines.top": False, + "axes.spines.right": False, + "legend.facecolor": "#121734", + "legend.edgecolor": "#3b3f5a", + }, + ), +} + + +def _shorten_seed(text: str | None) -> str | None: + """Turn '... seed=1234567890123' into '... seed=12345678โ€ฆ' if present.""" + if not text: + return text + m = re.search(r"seed=([^;,\s]+)", text) + if not m: + return text + raw = m.group(1) + short = (raw[:8] + "โ€ฆ") if len(raw) > 10 else raw + return re.sub(r"seed=[^;,\s]+", f"seed={short}", text) + + +def _apply_titles(fig: Figure, ax: Axes, title: str, subtitle: str | None) -> None: + """Consistent title placement: figure-level title + small italic subtitle.""" + fig.suptitle(title, fontsize=18, y=0.98) + ax.set_title(_shorten_seed(subtitle) or "", fontsize=12, fontstyle="italic", pad=4) + + +def _finalize_and_save(fig: Figure, output_dir: Path, stem: str, theme: str) -> None: + """Tight layout with space for suptitle, export PNG + (optional) SVG.""" + output_dir.mkdir(parents=True, exist_ok=True) + fig.tight_layout(rect=[0, 0, 1, 0.94]) + png = output_dir / f"{stem}_{theme}.png" + fig.savefig(png, dpi=300) + try: + fig.savefig(output_dir / f"{stem}_{theme}.svg", bbox_inches="tight") + except Exception: + pass # SVG is a nice-to-have + plt.close(fig) + + +# -------------------------- Public: model metrics ---------------------------- + + +def plot_model_metrics( + metrics: pl.DataFrame, + output_dir: Path, + stem: str, + title: str, + *, + subtitle: str = "", + figsize: tuple[int, int] | None = None, + agents: int | None = None, + steps: int | None = None, +) -> None: + """ + Plot time-series metrics from a Polars DataFrame and export light/dark PNG/SVG. + + - Auto-detects `step` or adds one if missing. + - Melts all non-`step` columns into long form. + - If there's a single metric (e.g., 'gini'), removes legend and uses a + descriptive y-axis label (e.g., 'Gini coefficient'). + - Optional `agents` and `steps` will be appended to the suptitle as + "(N=, T=)"; if `steps` is omitted it will be inferred + from the `step` column when available. + """ + if metrics.is_empty(): + return + + if "step" not in metrics.columns: + metrics = metrics.with_row_index("step") + + # If steps not provided, try to infer from the data (max step + 1). Keep it None if we can't determine it. + if steps is None: + try: + steps = int(metrics.select(pl.col("step").max()).item()) + 1 + except Exception: + steps = None + + value_cols: Sequence[str] = [c for c in metrics.columns if c != "step"] + if not value_cols: + return + + long = ( + metrics.select(["step", *value_cols]) + .unpivot( + index="step", on=value_cols, variable_name="metric", value_name="value" + ) + .to_pandas() + ) + + # Compose informative title with optional (N, T) + if agents is not None and steps is not None: + full_title = f"{title} (N={agents}, T={steps})" + elif agents is not None: + full_title = f"{title} (N={agents})" + elif steps is not None: + full_title = f"{title} (T={steps})" + else: + full_title = title + + for theme, cfg in _THEMES.items(): + sns.set_theme(**cfg) + sns.set_context("talk") + fig, ax = plt.subplots(figsize=figsize or (10, 6)) + + sns.lineplot(data=long, x="step", y="value", hue="metric", linewidth=2, ax=ax) + + _apply_titles(fig, ax, full_title, subtitle) + + ax.set_xlabel("Step") + unique_metrics = long["metric"].unique() + + if len(unique_metrics) == 1: + name = unique_metrics[0] + ax.set_ylabel(name.capitalize()) + leg = ax.get_legend() + if leg is not None: + leg.remove() + vals = long.loc[long["metric"] == name, "value"] + if not vals.empty: + vmin, vmax = float(vals.min()), float(vals.max()) + pad = max(0.005, (vmax - vmin) * 0.05) + ax.set_ylim(vmin - pad, vmax + pad) + else: + ax.set_ylabel("Value") + leg = ax.get_legend() + if theme == "dark" and leg is not None: + leg.set_title(None) + leg.get_frame().set_alpha(0.8) + + ax.yaxis.set_major_formatter(FormatStrFormatter("%.3f")) + ax.margins(x=0.01) + + _finalize_and_save(fig, output_dir, stem, theme) + + +# -------------------------- Public: agent metrics ---------------------------- + + +def plot_agent_metrics( + agent_metrics: pl.DataFrame, + output_dir: Path, + stem: str, + *, + title: str = "Agent metrics", + subtitle: str = "", + figsize: tuple[int, int] | None = None, +) -> None: + """ + Plot agent-level metrics (multi-series) and export light/dark PNG/SVG. + + - Preserves common id vars if present: `step`, `seed`, `batch`. + - Uses the first column as id if none of the preferred ids exist. + """ + if agent_metrics is None or agent_metrics.is_empty(): + return + + preferred = ["step", "seed", "batch"] + id_vars = [c for c in preferred if c in agent_metrics.columns] or [ + agent_metrics.columns[0] + ] + + # Determine which columns to unpivot (all columns except the id vars). + value_cols = [c for c in agent_metrics.columns if c not in id_vars] + if not value_cols: + return + + melted = agent_metrics.unpivot( + index=id_vars, on=value_cols, variable_name="metric", value_name="value" + ).to_pandas() + + xcol = id_vars[0] + + for theme, cfg in _THEMES.items(): + sns.set_theme(**cfg) + sns.set_context("talk") + fig, ax = plt.subplots(figsize=figsize or (10, 6)) + + sns.lineplot(data=melted, x=xcol, y="value", hue="metric", linewidth=1.8, ax=ax) + + _apply_titles(fig, ax, title, subtitle) + ax.set_xlabel(xcol.capitalize()) + ax.set_ylabel("Value") + + if theme == "dark": + leg = ax.get_legend() + if leg is not None: + leg.set_title(None) + leg.get_frame().set_alpha(0.8) + + _finalize_and_save(fig, output_dir, f"{stem}_agents", theme) + + +# -------------------------- Public: performance ------------------------------ + + +def plot_performance( + df: pl.DataFrame, + output_dir: Path, + stem: str, + *, + title: str = "Runtime vs agents", + subtitle: str = "", + figsize: tuple[int, int] | None = None, +) -> None: + """ + Plot backend performance (runtime vs agents) with meanยฑsd error bars. + Expected columns: `agents`, `runtime_seconds`, `backend`. + """ + if df.is_empty(): + return + + pdf = df.to_pandas() + + for theme, cfg in _THEMES.items(): + sns.set_theme(**cfg) + sns.set_context("talk") + fig, ax = plt.subplots(figsize=figsize or (10, 6)) + + sns.lineplot( + data=pdf, + x="agents", + y="runtime_seconds", + hue="backend", + estimator="mean", + errorbar="sd", + marker="o", + ax=ax, + ) + + _apply_titles(fig, ax, title, subtitle) + ax.set_xlabel("Agents") + ax.set_ylabel("Runtime (seconds)") + + if theme == "dark": + leg = ax.get_legend() + if leg is not None: + leg.set_title(None) + leg.get_frame().set_alpha(0.8) + + _finalize_and_save(fig, output_dir, stem, theme) + + +__all__ = [ + "plot_model_metrics", + "plot_agent_metrics", + "plot_performance", +] diff --git a/examples/sugarscape_ig/backend_frames/__init__.py b/examples/sugarscape_ig/backend_frames/__init__.py new file mode 100644 index 00000000..614fa64d --- /dev/null +++ b/examples/sugarscape_ig/backend_frames/__init__.py @@ -0,0 +1 @@ +"""mesa-frames backend package for Sugarscape IG examples.""" diff --git a/examples/sugarscape_ig/backend_mesa/__init__.py b/examples/sugarscape_ig/backend_mesa/__init__.py new file mode 100644 index 00000000..463099c0 --- /dev/null +++ b/examples/sugarscape_ig/backend_mesa/__init__.py @@ -0,0 +1 @@ +"""Mesa backend package for Sugarscape IG examples.""" diff --git a/examples/sugarscape_ig/mesa_comparison.png b/examples/sugarscape_ig/mesa_comparison.png deleted file mode 100644 index c619ae28..00000000 Binary files a/examples/sugarscape_ig/mesa_comparison.png and /dev/null differ diff --git a/examples/sugarscape_ig/performance_comparison.py b/examples/sugarscape_ig/performance_comparison.py deleted file mode 100644 index d8d2f196..00000000 --- a/examples/sugarscape_ig/performance_comparison.py +++ /dev/null @@ -1,224 +0,0 @@ -import math - -import matplotlib.pyplot as plt -import numpy as np -import perfplot -import polars as pl -import seaborn as sns -from polars.testing import assert_frame_equal -from ss_mesa.model import SugarscapeMesa -from ss_polars.agents import ( - AntPolarsLoopDF, - AntPolarsLoopNoVec, - AntPolarsNumbaCPU, - AntPolarsNumbaGPU, - AntPolarsNumbaParallel, -) -from ss_polars.model import SugarscapePolars -from collections.abc import Callable - - -class SugarScapeSetup: - def __init__(self, n: int): - if n >= 10**6: - density = 0.17 # FLAME2-GPU - else: - density = 0.04 # mesa - self.n = n - self.seed = 42 - dimension = math.ceil(math.sqrt(n / density)) - random_gen = np.random.default_rng(self.seed) - self.sugar_grid = random_gen.integers(0, 4, (dimension, dimension)) - self.initial_sugar = random_gen.integers(6, 25, n) - self.metabolism = random_gen.integers(2, 4, n) - self.vision = random_gen.integers(1, 6, n) - self.initial_positions = pl.DataFrame( - schema={"dim_0": pl.Int64, "dim_1": pl.Int64} - ) - while self.initial_positions.shape[0] < n: - initial_pos_0 = random_gen.integers( - 0, dimension, n - self.initial_positions.shape[0] - ) - initial_pos_1 = random_gen.integers( - 0, dimension, n - self.initial_positions.shape[0] - ) - self.initial_positions = self.initial_positions.vstack( - pl.DataFrame( - { - "dim_0": initial_pos_0, - "dim_1": initial_pos_1, - } - ) - ).unique(maintain_order=True) - return - - -def mesa_implementation(setup: SugarScapeSetup): - model = SugarscapeMesa( - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def mesa_frames_polars_loop_DF(setup: SugarScapeSetup): - model = SugarscapePolars( - AntPolarsLoopDF, - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def mesa_frames_polars_loop_no_vec(setup: SugarScapeSetup): - model = SugarscapePolars( - AntPolarsLoopNoVec, - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def mesa_frames_polars_numba_cpu(setup: SugarScapeSetup): - model = SugarscapePolars( - AntPolarsNumbaCPU, - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def mesa_frames_polars_numba_gpu(setup: SugarScapeSetup): - model = SugarscapePolars( - AntPolarsNumbaGPU, - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def mesa_frames_polars_numba_parallel(setup: SugarScapeSetup): - model = SugarscapePolars( - AntPolarsNumbaParallel, - setup.n, - setup.sugar_grid, - setup.initial_sugar, - setup.metabolism, - setup.vision, - setup.initial_positions, - setup.seed, - ) - model.run_model(100) - return model - - -def plot_and_print_benchmark( - labels: list[str], - kernels: list[Callable], - n_range: list[int], - title: str, - image_path: str, - equality_check: Callable | None = None, -): - out = perfplot.bench( - setup=SugarScapeSetup, - kernels=kernels, - labels=labels, - n_range=n_range, - xlabel="Number of agents", - equality_check=equality_check, - title=title, - ) - plt.ylabel("Execution time (s)") - out.save(image_path) - print("\nExecution times:") - for i, label in enumerate(labels): - print(f"---------------\n{label}:") - for n, t in zip(out.n_range, out.timings_s[i]): - print(f" Number of agents: {n}, Time: {t:.2f} seconds") - print("---------------") - - -def polars_equality_check(a: SugarscapePolars, b: SugarscapePolars): - assert_frame_equal(a.space.agents, b.space.agents, check_row_order=False) - assert_frame_equal(a.space.cells, b.space.cells, check_row_order=False) - return True - - -def main(): - # Mesa comparison - sns.set_theme(style="whitegrid") - labels_0 = [ - "mesa-frames (pl numba parallel)", - "mesa", - ] - kernels_0 = [ - mesa_frames_polars_numba_parallel, - mesa_implementation, - ] - n_range_0 = [k for k in range(10**5, 5 * 10**5 + 2, 10**5)] - title_0 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_0) - image_path_0 = "mesa_comparison.png" - plot_and_print_benchmark(labels_0, kernels_0, n_range_0, title_0, image_path_0) - - # mesa-frames comparison - labels_1 = [ - "mesa-frames (pl loop DF)", - "mesa-frames (pl loop no vec)", - "mesa-frames (pl numba CPU)", - "mesa-frames (pl numba parallel)", - "mesa-frames (pl numba GPU)", - ] - # Polars best_moves (non-vectorized loop vs DF loop vs numba loop) - kernels_1 = [ - mesa_frames_polars_loop_DF, - mesa_frames_polars_loop_no_vec, - mesa_frames_polars_numba_cpu, - mesa_frames_polars_numba_parallel, - mesa_frames_polars_numba_gpu, - ] - n_range_1 = [k for k in range(10**6, 3 * 10**6 + 2, 10**6)] - title_1 = "100 steps of the SugarScape IG model:\n" + " vs ".join(labels_1) - image_path_1 = "polars_comparison.png" - plot_and_print_benchmark( - labels_1, - kernels_1, - n_range_1, - title_1, - image_path_1, - equality_check=polars_equality_check, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/sugarscape_ig/polars_comparison.png b/examples/sugarscape_ig/polars_comparison.png deleted file mode 100644 index 1b211261..00000000 Binary files a/examples/sugarscape_ig/polars_comparison.png and /dev/null differ diff --git a/examples/sugarscape_ig/ss_mesa/agents.py b/examples/sugarscape_ig/ss_mesa/agents.py deleted file mode 100644 index e4d1a700..00000000 --- a/examples/sugarscape_ig/ss_mesa/agents.py +++ /dev/null @@ -1,83 +0,0 @@ -import math - -import mesa - - -def get_distance(pos_1, pos_2): - """Get the distance between two point - - Args: - pos_1, pos_2: Coordinate tuples for both points. - """ - x1, y1 = pos_1 - x2, y2 = pos_2 - dx = x1 - x2 - dy = y1 - y2 - return math.sqrt(dx**2 + dy**2) - - -class AntMesa(mesa.Agent): - def __init__(self, model, moore=False, sugar=0, metabolism=0, vision=0): - super().__init__(model) - self.moore = moore - self.sugar = sugar - self.metabolism = metabolism - self.vision = vision - - def get_sugar(self, pos): - this_cell = self.model.space.get_cell_list_contents([pos]) - for agent in this_cell: - if type(agent) is Sugar: - return agent - - def is_occupied(self, pos): - this_cell = self.model.space.get_cell_list_contents([pos]) - return any(isinstance(agent, AntMesa) for agent in this_cell) - - def move(self): - # Get neighborhood within vision - neighbors = [ - i - for i in self.model.space.get_neighborhood( - self.pos, self.moore, False, radius=self.vision - ) - if not self.is_occupied(i) - ] - neighbors.append(self.pos) - # Look for location with the most sugar - max_sugar = max(self.get_sugar(pos).amount for pos in neighbors) - candidates = [ - pos for pos in neighbors if self.get_sugar(pos).amount == max_sugar - ] - # Narrow down to the nearest ones - min_dist = min(get_distance(self.pos, pos) for pos in candidates) - final_candidates = [ - pos for pos in candidates if get_distance(self.pos, pos) == min_dist - ] - self.random.shuffle(final_candidates) - self.model.space.move_agent(self, final_candidates[0]) - - def eat(self): - sugar_patch = self.get_sugar(self.pos) - self.sugar = self.sugar - self.metabolism + sugar_patch.amount - sugar_patch.amount = 0 - - def step(self): - self.move() - self.eat() - if self.sugar <= 0: - self.model.space.remove_agent(self) - self.model.agents.remove(self) - - -class Sugar(mesa.Agent): - def __init__(self, model, max_sugar): - super().__init__(model) - self.amount = max_sugar - self.max_sugar = max_sugar - - def step(self): - if self.model.space.is_cell_empty(self.pos): - self.amount = self.max_sugar - else: - self.amount = 0 diff --git a/examples/sugarscape_ig/ss_mesa/model.py b/examples/sugarscape_ig/ss_mesa/model.py deleted file mode 100644 index 43114413..00000000 --- a/examples/sugarscape_ig/ss_mesa/model.py +++ /dev/null @@ -1,77 +0,0 @@ -import mesa -import numpy as np -import polars as pl - -from .agents import AntMesa, Sugar - - -class SugarscapeMesa(mesa.Model): - """ - Sugarscape 2 Instant Growback - """ - - def __init__( - self, - n_agents: int, - sugar_grid: np.ndarray | None = None, - initial_sugar: np.ndarray | None = None, - metabolism: np.ndarray | None = None, - vision: np.ndarray | None = None, - initial_positions: pl.DataFrame | None = None, - seed: int | None = None, - width: int | None = None, - height: int | None = None, - ): - """ - Create a new Instant Growback model with the given parameters. - - """ - super().__init__() - - # Set parameters - if sugar_grid is None: - sugar_grid = np.random.randint(0, 4, (width, height)) - if initial_sugar is None: - initial_sugar = np.random.randint(6, 25, n_agents) - if metabolism is None: - metabolism = np.random.randint(2, 4, n_agents) - if vision is None: - vision = np.random.randint(1, 6, n_agents) - if seed is not None: - self.reset_randomizer(seed) - - self.width, self.height = sugar_grid.shape - self.n_agents = n_agents - self.space = mesa.space.MultiGrid(self.width, self.height, torus=False) - - self.sugars = [] - - for _, (x, y) in self.space.coord_iter(): - max_sugar = sugar_grid[x, y] - sugar = Sugar(self, max_sugar) - self.space.place_agent(sugar, (x, y)) - self.sugars.append(sugar) - - # Create agent: - for i in range(self.n_agents): - if initial_positions is not None: - x = initial_positions["dim_0"][i] - y = initial_positions["dim_1"][i] - else: - x = self.random.randrange(self.width) - y = self.random.randrange(self.height) - ssa = AntMesa(self, False, initial_sugar[i], metabolism[i], vision[i]) - self.agents.add(ssa) - self.space.place_agent(ssa, (x, y)) - - self.running = True - - def step(self): - self.agents.shuffle_do("step") - [sugar.step() for sugar in self.sugars] - - def run_model(self, step_count=200): - for _ in range(step_count): - if len(self.agents) == 0: - return - self.step() diff --git a/examples/sugarscape_ig/ss_polars/agents.py b/examples/sugarscape_ig/ss_polars/agents.py deleted file mode 100644 index 32ca91f5..00000000 --- a/examples/sugarscape_ig/ss_polars/agents.py +++ /dev/null @@ -1,406 +0,0 @@ -from abc import abstractmethod - -import numpy as np -import polars as pl -from numba import b1, guvectorize, int32 - -from mesa_frames import AgentSet, Model - - -class AntDFBase(AgentSet): - def __init__( - self, - model: Model, - n_agents: int, - initial_sugar: np.ndarray | None = None, - metabolism: np.ndarray | None = None, - vision: np.ndarray | None = None, - ): - super().__init__(model) - - if initial_sugar is None: - initial_sugar = model.random.integers(6, 25, n_agents) - if metabolism is None: - metabolism = model.random.integers(2, 4, n_agents) - if vision is None: - vision = model.random.integers(1, 6, n_agents) - - agents = pl.DataFrame( - { - "sugar": initial_sugar, - "metabolism": metabolism, - "vision": vision, - } - ) - self.add(agents) - - def eat(self): - # Only consider cells currently occupied by agents of this set - cells = self.space.cells.filter(pl.col("agent_id").is_not_null()) - mask_in_set = cells["agent_id"].is_in(self.index) - if mask_in_set.any(): - cells = cells.filter(mask_in_set) - ids = cells["agent_id"] - self[ids, "sugar"] = ( - self[ids, "sugar"] + cells["sugar"] - self[ids, "metabolism"] - ) - - def step(self): - self.shuffle().do("move").do("eat") - self.discard(self.df.filter(pl.col("sugar") <= 0)) - - def move(self): - neighborhood = self._get_neighborhood() - agent_order = self._get_agent_order(neighborhood) - neighborhood = self._prepare_neighborhood(neighborhood, agent_order) - best_moves = self.get_best_moves(neighborhood) - self.space.move_agents(agent_order["agent_id_center"], best_moves) - - def _get_neighborhood(self) -> pl.DataFrame: - """Get the neighborhood of each agent, completed with the sugar of the cell and the agent_id of the center cell - - NOTE: This method should be unnecessary if get_neighborhood/get_neighbors return the agent_id of the center cell and the properties of the cells - - Returns - ------- - pl.DataFrame - Neighborhood DataFrame - """ - neighborhood: pl.DataFrame = self.space.get_neighborhood( - radius=self["vision"], agents=self, include_center=True - ) - # Join self.space.cells to obtain properties ('sugar') per cell - - neighborhood = neighborhood.join(self.space.cells, on=["dim_0", "dim_1"]) - - # Join self.pos to obtain the agent_id of the center cell - # TODO: get_neighborhood/get_neighbors should return 'agent_id_center' instead of center position when input is AgentLike - - neighborhood = neighborhood.with_columns( - agent_id_center=neighborhood.join( - self.pos, - left_on=["dim_0_center", "dim_1_center"], - right_on=["dim_0", "dim_1"], - )["unique_id"] - ) - return neighborhood - - def _get_agent_order(self, neighborhood: pl.DataFrame) -> pl.DataFrame: - """Get the order of agents based on the original order of agents - - Parameters - ---------- - neighborhood : pl.DataFrame - Neighborhood DataFrame - - Returns - ------- - pl.DataFrame - DataFrame with 'agent_id_center' and 'agent_order' columns - """ - # Order of agents moves based on the original order of agents. - # The agent in his cell has order 0 (highest) - - return ( - neighborhood.unique( - subset=["agent_id_center"], keep="first", maintain_order=True - ) - .with_row_count("agent_order") - .select(["agent_id_center", "agent_order"]) - ) - - def _prepare_neighborhood( - self, neighborhood: pl.DataFrame, agent_order: pl.DataFrame - ) -> pl.DataFrame: - """Prepare the neighborhood DataFrame to find the best moves - - Parameters - ---------- - neighborhood : pl.DataFrame - Neighborhood DataFrame - agent_order : pl.DataFrame - DataFrame with 'agent_id_center' and 'agent_order' columns - - Returns - ------- - pl.DataFrame - Prepared neighborhood DataFrame - """ - neighborhood = neighborhood.join(agent_order, on="agent_id_center") - - # Add blocking agent order - neighborhood = neighborhood.join( - agent_order.select( - pl.col("agent_id_center").alias("agent_id"), - pl.col("agent_order").alias("blocking_agent_order"), - ), - on="agent_id", - how="left", - ).rename({"agent_id": "blocking_agent_id"}) - - # Filter only possible moves (agent is in his cell, blocking agent has moved before him or there is no blocking agent) - neighborhood = neighborhood.filter( - (pl.col("agent_order") >= pl.col("blocking_agent_order")) - | pl.col("blocking_agent_order").is_null() - ) - - # Sort neighborhood by agent_order & max_sugar (max_sugar because we will check anyway if the cell is empty) - # However, we need to make sure that the current agent cell is ordered by current sugar (since it's 0 until agent hasn't moved) - neighborhood = neighborhood.with_columns( - max_sugar=pl.when(pl.col("blocking_agent_id") == pl.col("agent_id_center")) - .then(pl.lit(0)) - .otherwise(pl.col("max_sugar")) - ).sort( - ["agent_order", "max_sugar", "radius", "dim_0"], - descending=[False, True, False, False], - ) - return neighborhood - - def get_best_moves(self, neighborhood: pl.DataFrame) -> pl.DataFrame: - """Get the best moves for each agent - - Parameters - ---------- - neighborhood : pl.DataFrame - Neighborhood DataFrame - - Returns - ------- - pl.DataFrame - DataFrame with the best moves for each agent - """ - raise NotImplementedError("Subclasses must implement this method") - - -class AntPolarsLoopDF(AntDFBase): - def get_best_moves(self, neighborhood: pl.DataFrame): - best_moves = pl.DataFrame() - - # While there are agents that do not have a best move, keep looking for one - while len(best_moves) < len(self.df): - # Check if there are previous agents that might make the same move (priority for the given move is > 1) - neighborhood = neighborhood.with_columns( - priority=pl.col("agent_order").cum_count().over(["dim_0", "dim_1"]) - ) - - # Get the best moves for each agent: - # If duplicates are found, select the one with the highest order - new_best_moves = ( - neighborhood.group_by("agent_id_center", maintain_order=True) - .first() - .unique(subset=["dim_0", "dim_1"], keep="first", maintain_order=True) - ) - # Agents can make the move if: - # - There is no blocking agent - # - The agent is in its own cell - # - The blocking agent has moved before him - # - There isn't a higher priority agent that might make the same move - - condition = pl.col("blocking_agent_id").is_null() | ( - pl.col("blocking_agent_id") == pl.col("agent_id_center") - ) - if len(best_moves) > 0: - condition = condition | pl.col("blocking_agent_id").is_in( - best_moves["agent_id_center"] - ) - - condition = condition & (pl.col("priority") == 1) - - new_best_moves = new_best_moves.filter(condition) - - best_moves = pl.concat([best_moves, new_best_moves]) - - # Remove agents that have already moved - neighborhood = neighborhood.filter( - ~pl.col("agent_id_center").is_in(best_moves["agent_id_center"]) - ) - - # Remove cells that have been already selected - neighborhood = neighborhood.join( - best_moves.select(["dim_0", "dim_1"]), on=["dim_0", "dim_1"], how="anti" - ) - - # Check if there are previous agents that might make the same move (priority for the given move is > 1) - neighborhood = neighborhood.with_columns( - priority=pl.col("agent_order").cum_count().over(["dim_0", "dim_1"]) - ) - return best_moves.sort("agent_order").select(["dim_0", "dim_1"]) - - -class AntPolarsLoop(AntDFBase): - numba_target = None - - def get_best_moves(self, neighborhood: pl.DataFrame): - occupied_cells, free_cells, target_cells = self._prepare_cells(neighborhood) - best_moves_func = self._get_best_moves() - - processed_agents = np.zeros(len(self.df), dtype=np.bool_) - - if self.numba_target is None: - # Non-vectorized case: we need to create and pass the best_moves array - map_batches_func = lambda df: best_moves_func( - occupied_cells, - free_cells, - target_cells, - df.struct.field("agent_order"), - df.struct.field("blocking_agent_order"), - processed_agents, - best_moves=np.full(len(self.df), -1, dtype=np.int32), - ) - else: - # Vectorized case: Polars will create the output array (best_moves) automatically - map_batches_func = lambda df: best_moves_func( - occupied_cells.astype(np.int32), - free_cells.astype(np.bool_), - target_cells.astype(np.int32), - df.struct.field("agent_order"), - df.struct.field("blocking_agent_order"), - processed_agents.astype(np.bool_), - ) - - best_moves = ( - # Only fill nulls for the column we need as an int sentinel; - # avoid touching UInt columns like 'blocking_agent_id'. - neighborhood.with_columns(pl.col("blocking_agent_order").fill_null(-1)) - .cast({"agent_order": pl.Int32, "blocking_agent_order": pl.Int32}) - .select( - pl.struct(["agent_order", "blocking_agent_order"]).map_batches( - map_batches_func, - return_dtype=pl.Int32, - ) - ) - .with_columns( - dim_0=pl.col("agent_order") // self.space.dimensions[1], - dim_1=pl.col("agent_order") % self.space.dimensions[1], - ) - .drop("agent_order") - ) - return best_moves - - def _prepare_cells( - self, neighborhood: pl.DataFrame - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Get the occupied and free cells and the target cells for each agent, - based on the neighborhood DataFrame such that the arrays refer to a flattened version of the grid - - Parameters - ---------- - neighborhood : pl.DataFrame - Neighborhood DataFrame - - Returns - ------- - tuple[np.ndarray, np.ndarray, np.ndarray] - occupied_cells, free_cells, target_cells - """ - occupied_cells = ( - neighborhood[["agent_id_center", "agent_order"]] - .unique() - .join(self.pos, left_on="agent_id_center", right_on="unique_id") - .with_columns( - flattened=(pl.col("dim_0") * self.space.dimensions[1] + pl.col("dim_1")) - ) - .sort("agent_order")["flattened"] - .to_numpy() - ) - free_cells = np.ones( - self.space.dimensions[0] * self.space.dimensions[1], dtype=np.bool_ - ) - free_cells[occupied_cells] = False - - target_cells = ( - neighborhood["dim_0"] * self.space.dimensions[1] + neighborhood["dim_1"] - ).to_numpy() - return occupied_cells, free_cells, target_cells - - def _get_best_moves(self): - raise NotImplementedError("Subclasses must implement this method") - - -class AntPolarsLoopNoVec(AntPolarsLoop): - # Non-vectorized case - def _get_best_moves(self): - def inner_get_best_moves( - occupied_cells: np.ndarray, - free_cells: np.ndarray, - target_cells: np.ndarray, - agent_id_center: np.ndarray, - blocking_agent: np.ndarray, - processed_agents: np.ndarray, - best_moves: np.ndarray, - ) -> np.ndarray: - for i, agent in enumerate(agent_id_center): - # If the agent has not moved yet - if not processed_agents[agent]: - # If the target cell is free - if free_cells[target_cells[i]] or blocking_agent[i] == agent: - best_moves[agent] = target_cells[i] - # Free current cell - free_cells[occupied_cells[agent]] = True - # Occupy target cell - free_cells[target_cells[i]] = False - processed_agents[agent] = True - return best_moves - - return inner_get_best_moves - - -class AntPolarsNumba(AntPolarsLoop): - # Vectorized case - def _get_best_moves(self): - @guvectorize( - [ - ( - int32[:], - b1[:], - int32[:], - int32[:], - int32[:], - b1[:], - int32[:], - ) - ], - "(n), (m), (p), (p), (p), (n)->(n)", - nopython=True, - target=self.numba_target, - # Writable inputs should be declared according to https://numba.pydata.org/numba-doc/dev/user/vectorize.html#overwriting-input-values - # In this case, there doesn't seem to be a difference. I will leave it commented for reference so that we can use CUDA target (which doesn't support writable_args) - # writable_args=( - # "free_cells", - # "processed_agents", - # ), - ) - def vectorized_get_best_moves( - occupied_cells, - free_cells, - target_cells, - agent_id_center, - blocking_agent, - processed_agents, - best_moves, - ): - for i, agent in enumerate(agent_id_center): - # If the agent has not moved yet - if not processed_agents[agent]: - # If the target cell is free - if free_cells[target_cells[i]] or blocking_agent[i] == agent: - best_moves[agent] = target_cells[i] - # Free current cell - free_cells[occupied_cells[agent]] = True - # Occupy target cell - free_cells[target_cells[i]] = False - processed_agents[agent] = True - - return vectorized_get_best_moves - - -class AntPolarsNumbaCPU(AntPolarsNumba): - numba_target = "cpu" - - -class AntPolarsNumbaParallel(AntPolarsNumba): - numba_target = "parallel" - - -class AntPolarsNumbaGPU(AntPolarsNumba): - numba_target = "cuda" diff --git a/examples/sugarscape_ig/ss_polars/model.py b/examples/sugarscape_ig/ss_polars/model.py deleted file mode 100644 index 36b2718e..00000000 --- a/examples/sugarscape_ig/ss_polars/model.py +++ /dev/null @@ -1,57 +0,0 @@ -import numpy as np -import polars as pl - -from mesa_frames import Grid, Model - -from .agents import AntDFBase - - -class SugarscapePolars(Model): - def __init__( - self, - agent_type: type[AntDFBase], - n_agents: int, - sugar_grid: np.ndarray | None = None, - initial_sugar: np.ndarray | None = None, - metabolism: np.ndarray | None = None, - vision: np.ndarray | None = None, - initial_positions: pl.DataFrame | None = None, - seed: int | None = None, - width: int | None = None, - height: int | None = None, - ): - super().__init__(seed) - if sugar_grid is None: - sugar_grid = self.random.integers(0, 4, (width, height)) - grid_dimensions = sugar_grid.shape - self.space = Grid( - self, grid_dimensions, neighborhood_type="von_neumann", capacity=1 - ) - dim_0 = pl.Series("dim_0", pl.arange(grid_dimensions[0], eager=True)).to_frame() - dim_1 = pl.Series("dim_1", pl.arange(grid_dimensions[1], eager=True)).to_frame() - sugar_grid = dim_0.join(dim_1, how="cross").with_columns( - sugar=sugar_grid.flatten(), max_sugar=sugar_grid.flatten() - ) - self.space.set_cells(sugar_grid) - # Create and register the main agent set; keep its name for later lookups - main_set = agent_type(self, n_agents, initial_sugar, metabolism, vision) - self.sets += main_set - self._main_set_name = main_set.name - if initial_positions is not None: - self.space.place_agents(self.sets, initial_positions) - else: - self.space.place_to_empty(self.sets) - - def run_model(self, steps: int) -> list[int]: - for _ in range(steps): - # Stop if the main agent set is empty - if len(self.sets[self._main_set_name]) == 0: # type: ignore[index] - return - empty_cells = self.space.empty_cells - full_cells = self.space.full_cells - max_sugar = self.space.cells.join( - empty_cells, on=["dim_0", "dim_1"] - ).select(pl.col("max_sugar")) - self.space.set_cells(full_cells, {"sugar": 0}) - self.space.set_cells(empty_cells, {"sugar": max_sugar}) - self.step() diff --git a/examples/utils.py b/examples/utils.py new file mode 100644 index 00000000..4d075dc4 --- /dev/null +++ b/examples/utils.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass +import mesa_frames +import mesa + + +@dataclass +class FramesSimulationResult: + """Container for example simulation outputs. + + The dataclass is intentionally permissive: some backends only provide + `metrics`, while others also return `agent_metrics`. + """ + + datacollector: mesa_frames.DataCollector + + +@dataclass +class MesaSimulationResult: + """Container for example simulation outputs. + + The dataclass is intentionally permissive: some backends only provide + `metrics`, while others also return `agent_metrics`. + """ + + datacollector: mesa.DataCollector diff --git a/mkdocs.yml b/mkdocs.yml index 8a462881..a1caa258 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ # Project information -site_name: mesa-frames +site_name: mesa-frames documentation site_url: https://projectmesa.github.io/mesa-frames repo_url: https://github.com/projectmesa/mesa-frames repo_name: projectmesa/mesa-frames @@ -40,6 +40,10 @@ theme: code: Roboto Mono icon: repo: fontawesome/brands/github + # Logo (PNG) + logo: https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.png + # Favicon (ICO) + favicon: https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.ico # Plugins plugins: @@ -96,6 +100,11 @@ extra_javascript: - https://polyfill.io/v3/polyfill.min.js?features=es6 - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js +# Custom CSS for branding (brand-core then material adapter) +extra_css: + - stylesheets/brand-core.css + - stylesheets/brand-material.css + # Customization extra: social: @@ -112,8 +121,8 @@ nav: - Classes: user-guide/1_classes.md - Introductory Tutorial: user-guide/2_introductory-tutorial.ipynb - Data Collector Tutorial: user-guide/4_datacollector.ipynb - - Advanced Tutorial: user-guide/3_advanced-tutorial.md - - Benchmarks: user-guide/4_benchmarks.md + - Advanced Tutorial: user-guide/3_advanced-tutorial.ipynb + - Benchmarks: user-guide/5_benchmarks.md - API Reference: api/index.html - Contributing: - Contribution Guide: contributing.md diff --git a/pyproject.toml b/pyproject.toml index 99b15899..addcc239 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ test = [ docs = [ { include-group = "typechecking" }, + "typer>=0.9.0", "mkdocs-material>=9.6.14", "mkdocs-jupyter>=0.25.1", "mkdocs-git-revision-date-localized-plugin>=1.4.7", @@ -76,10 +77,10 @@ docs = [ "sphinx-copybutton>=0.5.2", "sphinx-design>=0.6.1", "autodocsumm>=0.2.14", - "perfplot>=0.10.2", "seaborn>=0.13.2", "sphinx-autobuild>=2025.8.25", "mesa>=3.2.0", + "jupytext>=1.17.3", ] # dev = test โˆช docs โˆช extra tooling diff --git a/uv.lock b/uv.lock index a72164c0..cc55da48 100644 --- a/uv.lock +++ b/uv.lock @@ -1161,19 +1161,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899, upload-time = "2024-04-15T13:44:43.265Z" }, ] -[[package]] -name = "matplotx" -version = "0.3.10" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "matplotlib" }, - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c3/01/0e6938bb717fa7722d6d81336c62de71b815ce73e382aa1873a1e68ccc93/matplotx-0.3.10.tar.gz", hash = "sha256:b6926ce5274cf5da966cb46b90a8c7fefb761478c6c85c8f7ed3ee8ec90e86e5", size = 24041, upload-time = "2022-08-22T14:22:56.374Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/95/ef/e8a30503ae0c26681a9610c7f0be58646bea8119b98cc65c47661abc27a3/matplotx-0.3.10-py3-none-any.whl", hash = "sha256:4d7adafdb001c771d66d9362bb8ca99fcaed15319259223a714f36793dfabbb8", size = 25099, upload-time = "2022-08-22T14:22:54.733Z" }, -] - [[package]] name = "mdit-py-plugins" version = "0.5.0" @@ -1234,6 +1221,7 @@ dependencies = [ dev = [ { name = "autodocsumm" }, { name = "beartype" }, + { name = "jupytext" }, { name = "mesa" }, { name = "mkdocs-git-revision-date-localized-plugin" }, { name = "mkdocs-include-markdown-plugin" }, @@ -1242,7 +1230,6 @@ dev = [ { name = "mkdocs-minify-plugin" }, { name = "numba" }, { name = "numpydoc" }, - { name = "perfplot" }, { name = "pre-commit" }, { name = "pydata-sphinx-theme" }, { name = "pytest" }, @@ -1254,10 +1241,12 @@ dev = [ { name = "sphinx-copybutton" }, { name = "sphinx-design" }, { name = "sphinx-rtd-theme" }, + { name = "typer" }, ] docs = [ { name = "autodocsumm" }, { name = "beartype" }, + { name = "jupytext" }, { name = "mesa" }, { name = "mkdocs-git-revision-date-localized-plugin" }, { name = "mkdocs-include-markdown-plugin" }, @@ -1265,7 +1254,6 @@ docs = [ { name = "mkdocs-material" }, { name = "mkdocs-minify-plugin" }, { name = "numpydoc" }, - { name = "perfplot" }, { name = "pydata-sphinx-theme" }, { name = "seaborn" }, { name = "sphinx" }, @@ -1273,6 +1261,7 @@ docs = [ { name = "sphinx-copybutton" }, { name = "sphinx-design" }, { name = "sphinx-rtd-theme" }, + { name = "typer" }, ] test = [ { name = "beartype" }, @@ -1296,6 +1285,7 @@ requires-dist = [ dev = [ { name = "autodocsumm", specifier = ">=0.2.14" }, { name = "beartype", specifier = ">=0.21.0" }, + { name = "jupytext", specifier = ">=1.17.3" }, { name = "mesa", specifier = ">=3.2.0" }, { name = "mkdocs-git-revision-date-localized-plugin", specifier = ">=1.4.7" }, { name = "mkdocs-include-markdown-plugin", specifier = ">=7.1.5" }, @@ -1304,7 +1294,6 @@ dev = [ { name = "mkdocs-minify-plugin", specifier = ">=0.8.0" }, { name = "numba", specifier = ">=0.60.0" }, { name = "numpydoc", specifier = ">=1.8.0" }, - { name = "perfplot", specifier = ">=0.10.2" }, { name = "pre-commit", specifier = ">=4.2.0" }, { name = "pydata-sphinx-theme", specifier = ">=0.16.1" }, { name = "pytest", specifier = ">=8.3.5" }, @@ -1316,10 +1305,12 @@ dev = [ { name = "sphinx-copybutton", specifier = ">=0.5.2" }, { name = "sphinx-design", specifier = ">=0.6.1" }, { name = "sphinx-rtd-theme", specifier = ">=3.0.2" }, + { name = "typer", extras = ["all"], specifier = ">=0.9.0" }, ] docs = [ { name = "autodocsumm", specifier = ">=0.2.14" }, { name = "beartype", specifier = ">=0.21.0" }, + { name = "jupytext", specifier = ">=1.17.3" }, { name = "mesa", specifier = ">=3.2.0" }, { name = "mkdocs-git-revision-date-localized-plugin", specifier = ">=1.4.7" }, { name = "mkdocs-include-markdown-plugin", specifier = ">=7.1.5" }, @@ -1327,7 +1318,6 @@ docs = [ { name = "mkdocs-material", specifier = ">=9.6.14" }, { name = "mkdocs-minify-plugin", specifier = ">=0.8.0" }, { name = "numpydoc", specifier = ">=1.8.0" }, - { name = "perfplot", specifier = ">=0.10.2" }, { name = "pydata-sphinx-theme", specifier = ">=0.16.1" }, { name = "seaborn", specifier = ">=0.13.2" }, { name = "sphinx", specifier = ">=7.4.7" }, @@ -1335,6 +1325,7 @@ docs = [ { name = "sphinx-copybutton", specifier = ">=0.5.2" }, { name = "sphinx-design", specifier = ">=0.6.1" }, { name = "sphinx-rtd-theme", specifier = ">=3.0.2" }, + { name = "typer", extras = ["all"], specifier = ">=0.9.0" }, ] test = [ { name = "beartype", specifier = ">=0.21.0" }, @@ -1728,21 +1719,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", size = 31191, upload-time = "2023-12-10T22:30:43.14Z" }, ] -[[package]] -name = "perfplot" -version = "0.10.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "matplotlib" }, - { name = "matplotx" }, - { name = "numpy" }, - { name = "rich" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/97/41/51d8b9caa150a050de16a229f627e4b37515dbff0075259e4e75aff7218b/perfplot-0.10.2.tar.gz", hash = "sha256:d76daa72334564b5c8825663f24d15db55ea33e938b34595a146e5e44ed87e41", size = 25044, upload-time = "2022-03-03T15:56:37.392Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/20/85/ffaf2c1f92d17916c089a5c860d23b3117398f19f467fd1de1026d03aebc/perfplot-0.10.2-py3-none-any.whl", hash = "sha256:545ce0f7f22509ad00092d79a794cdc6e9805383e6cedab2bfed3519a7ef4e19", size = 21198, upload-time = "2022-03-03T15:56:35.388Z" }, -] - [[package]] name = "pexpect" version = "4.9.0" @@ -2520,6 +2496,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/11/00d3c3dfc25ad54e731d91449895a79e4bf2384dc3ac01809010ba88f6d5/seaborn-0.13.2-py3-none-any.whl", hash = "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987", size = 294914, upload-time = "2024-01-25T13:21:49.598Z" }, ] +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + [[package]] name = "six" version = "1.17.0" @@ -2832,6 +2817,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" }, ] +[[package]] +name = "typer" +version = "0.19.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "rich" }, + { name = "shellingham" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/ca/950278884e2ca20547ff3eb109478c6baf6b8cf219318e6bc4f666fad8e8/typer-0.19.2.tar.gz", hash = "sha256:9ad824308ded0ad06cc716434705f691d4ee0bfd0fb081839d2e426860e7fdca", size = 104755, upload-time = "2025-09-23T09:47:48.256Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/22/35617eee79080a5d071d0f14ad698d325ee6b3bf824fc0467c03b30e7fa8/typer-0.19.2-py3-none-any.whl", hash = "sha256:755e7e19670ffad8283db353267cb81ef252f595aa6834a0d1ca9312d9326cb9", size = 46748, upload-time = "2025-09-23T09:47:46.777Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0"