diff --git a/.github/workflows/docs-gh-pages.yml b/.github/workflows/docs-gh-pages.yml index 435af957..220c646b 100644 --- a/.github/workflows/docs-gh-pages.yml +++ b/.github/workflows/docs-gh-pages.yml @@ -1,40 +1,100 @@ -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 } - - name: Install uv via GitHub Action - uses: astral-sh/setup-uv@v6 + - 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: Short SHA + id: sha + run: echo "short=$(git rev-parse --short HEAD)" >> "$GITHUB_OUTPUT" + + - name: Upload site artifact + uses: actions/upload-artifact@v4 + with: + name: site + path: site - - name: Build Sphinx docs (API documentation) - run: sphinx-build -b html docs/api site/api + 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 - - name: Deploy to GitHub Pages + 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 - force_orphan: true \ No newline at end of file + 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/.github/workflows/publish.yml b/.github/workflows/publish.yml index 105824de..7204a0bc 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -17,33 +17,27 @@ jobs: with: fetch-depth: 0 token: ${{ secrets.VERSION_PUSH_TOKEN }} - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.x' - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install hatch + - name: Setup uv + uses: astral-sh/setup-uv@v3 - name: Set release version run: | # Get the tag from the GitHub release TAG=${GITHUB_REF#refs/tags/} # Remove 'v' prefix if present VERSION=${TAG#v} - hatch version $VERSION + uvx hatch version $VERSION - name: Build package - run: hatch build + run: uvx hatch build - name: Run tests - run: hatch run test:pytest + run: uvx hatch run test:pytest - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@release/v1 - name: Verify PyPI Release run: | # Verify PyPI release PACKAGE_NAME="mesa_frames" - CURRENT_VERSION=$(hatch version) - pip install $PACKAGE_NAME==$CURRENT_VERSION + CURRENT_VERSION=$(uvx hatch version) + uv pip install --system $PACKAGE_NAME==$CURRENT_VERSION python -c "import mesa_frames; print(mesa_frames.__version__)" - name: Update GitHub Release uses: softprops/action-gh-release@v1 @@ -51,9 +45,50 @@ jobs: with: files: | dist/* + - name: Generate changelog from release notes + id: notes + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + const tag = (context.payload.release && context.payload.release.tag_name) + ? context.payload.release.tag_name + : (process.env.GITHUB_REF || '').replace('refs/tags/', ''); + + const body = (context.payload.release && context.payload.release.body) ? context.payload.release.body : ''; + if (!body || body.trim().length === 0) { + core.setFailed('Release body is empty. Ensure the GitHub Release is created with auto-generated notes configured by .github/release.yml or supply a body.'); + } + + fs.writeFileSync('RELEASE_BODY.md', body, 'utf8'); + core.setOutput('tag', tag); + - name: Prepend notes to CHANGELOG.md + env: + TAG: ${{ steps.notes.outputs.tag }} + run: | + VERSION_NO_V=${TAG#v} + DATE_UTC=$(date -u +%Y-%m-%d) + echo "## Version ${VERSION_NO_V} โ€” ${DATE_UTC}" > RELEASE_HEADER.md + echo "" >> RELEASE_HEADER.md + if [ -f CHANGELOG.md ]; then + cat RELEASE_HEADER.md RELEASE_BODY.md CHANGELOG.md > CHANGELOG.new + else + cat RELEASE_HEADER.md RELEASE_BODY.md > CHANGELOG.new + fi + mv CHANGELOG.new CHANGELOG.md + - name: Commit and push CHANGELOG update + env: + TAG: ${{ steps.notes.outputs.tag }} + run: | + git config user.name github-actions + git config user.email github-actions@github.com + git add CHANGELOG.md + # Avoid CI cycles + git commit -m "Changelog: add notes for ${TAG} [skip ci]" || echo "No changelog changes to commit" + git push origin main || true - name: Create or recreate version branch run: | - CURRENT_VERSION=$(hatch version) + CURRENT_VERSION=$(uvx hatch version) BRANCH_NAME="v$CURRENT_VERSION" git config user.name github-actions @@ -72,15 +107,15 @@ jobs: - name: Update to Next Version run: | # Bump to next development version - hatch version patch - hatch version dev + uvx hatch version patch + uvx hatch version dev # Get the new version - NEW_VERSION=$(hatch version) + NEW_VERSION=$(uvx hatch version) # Commit and push the version bump git config user.name github-actions git config user.email github-actions@github.com - git add mesa_frames/__init__.py + git add mesa_frames/__init__.py CHANGELOG.md git commit -m "Bump version to $NEW_VERSION [skip ci]" - git push origin main \ No newline at end of file + git push origin main diff --git a/.gitignore b/.gitignore index 4a189d56..11fa6a54 100644 --- a/.gitignore +++ b/.gitignore @@ -154,3 +154,12 @@ cython_debug/ *.code-workspace llm_rules.md .python-version +docs/site +docs/api/_build +docs/general/tutorials/data_csv +docs/general/tutorials/data_parquet +docs/general/tutorials/*.ipynb +docs/api/reference/**/mesa_frames.*.rst +examples/**/results +benchmarks/**/results +benchmarks/**/plots \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..b7595854 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,40 @@ +## Version 0.1.0-alpha โ€” 2024-08-28 + +## What's Changed + +* Refactoring mesa.Agent, mesa.AgentSet, mesa.Model -> AgentSetDF, AgentsDF, ModelDF by @adamamer20 in +* setup: Migrate from setup.py to pyproject.toml by @rht in +* ci: Add pre-commit configuration by @rht in +* Merge requirements.txt into pyproject.toml by @rht in +* ci: Add GA for tests by @rht in +* Changes to AgentSetDF and AgentsDF before time.py -> CopyMixin by @adamamer20 in +* benchmark: Split Polars agent into native and concise by @rht in +* benchmark: Split pandas agent into native and concise by @rht in +* speed up mesa readme_plot script by @adamamer20 in +* Adding DataFrameMixin for improved reusability/encapsulation by @adamamer20 in +* Abstract SpaceDF by @adamamer20 in +* Adding Abstract DiscreteSpaceDF by @adamamer20 in +* Adding abstract GridDF by @adamamer20 in +* Additional methods and fixes to DataFrameMixin by @adamamer20 in +* Concrete GridPandas by @adamamer20 in +* [pre-commit.ci] pre-commit autoupdate by @pre-commit-ci in +* Fixes and Tests for PolarsMixin by @adamamer20 in +* Adding Comparison and Indexing methods to DataFrameMixin by @adamamer20 in +* Concrete GridPolars by @adamamer20 in +* Sugarscape Instantaneous Growback (Pandas-with-loop implementation) by @adamamer20 in +* Adding pydoclint and properly format docstring by @adamamer20 in +* Docs with material-from-mkdocs by @adamamer20 in +* Enforce correct numpy docstring formatting with ruff.pydocstyle by @adamamer20 in +* API Documentation with Sphinx by @adamamer20 in +* Move images from docs to docs/general to make it available for mkdocs by @adamamer20 in +* Adding user guide by @adamamer20 in +* Adding SugarScape IG (polars with loops) by @adamamer20 in +* Automatic publishing on PyPI on new release by @adamamer20 in + +## New Contributors + +* @adamamer20 made their first contribution in +* @rht made their first contribution in +* @pre-commit-ci made their first contribution in + +**Full Changelog**: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 147b84d3..82a340c9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -15,16 +15,22 @@ Before contributing, we recommend reviewing our [roadmap](https://projectmesa.gi Before you begin contributing, ensure that you have the necessary tools installed: - **Install Python** (at least the version specified in `requires-python` of `pyproject.toml`). ๐Ÿ -- We recommend using a virtual environment manager like: - - [Astral's UV](https://docs.astral.sh/uv/#installation) ๐ŸŒŸ - - [Hatch](https://hatch.pypa.io/latest/install/) ๐Ÿ—๏ธ + +-- We recommend using a virtual environment manager like: + + - [Astral's UV](https://docs.astral.sh/uv/#installation) ๐ŸŒŸ + - [Hatch](https://hatch.pypa.io/latest/install/) ๐Ÿ—๏ธ + - Install **pre-commit** to enforce code quality standards before pushing changes: + - [Pre-commit installation guide](https://pre-commit.com/#install) โœ… - [More about pre-commit hooks](https://stackoverflow.com/collectives/articles/71270196/how-to-use-pre-commit-to-automatically-correct-commits-and-merge-requests-with-g) -- If using **VS Code**, consider installing these extensions to automatically enforce formatting: - - [Ruff](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) โ€“ Python linting & formatting ๐Ÿพ - - [Markdownlint](https://marketplace.visualstudio.com/items?itemName=DavidAnson.vscode-markdownlint) โ€“ Markdown linting (for documentation) โœ๏ธ - - [Git Hooks](https://marketplace.visualstudio.com/items?itemName=lakshmikanthayyadevara.githooks) โ€“ Automatically runs & visualizes pre-commit hooks ๐Ÿ”— + +-- If using **VS Code**, consider installing these extensions to automatically enforce formatting: + + - [Ruff](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) โ€“ Python linting & formatting ๐Ÿพ + - [Markdownlint](https://marketplace.visualstudio.com/items?itemName=DavidAnson.vscode-markdownlint) โ€“ Markdown linting (for documentation) โœ๏ธ + - [Git Hooks](https://marketplace.visualstudio.com/items?itemName=lakshmikanthayyadevara.githooks) โ€“ Automatically runs & visualizes pre-commit hooks ๐Ÿ”— --- @@ -58,28 +64,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. - -- **Using UV**: - - ```sh - uv add --dev .[dev] - ``` - -- **Using Hatch**: +We manage the development environment with [uv](https://docs.astral.sh/uv/): - ```sh - hatch env create dev - ``` +```sh +uv sync --all-extras +``` -- **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,33 +90,35 @@ 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: Runtime Type Checking (beartype)** ๐Ÿ” -- **Optional: Enable runtime type checking** during development for enhanced type safety: + You can enable stricter runtime validation of function arguments/returns with `beartype` during local development: ```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" - Runtime type checking is automatically enabled in these scenarios: + Quick facts: - - **Hatch development environment** (`hatch shell dev`) - - **VS Code debugging** (when using the debugger) - - **VS Code testing** (when running tests through VS Code's testing interface) +- Automatically enabled in: Hatch dev env (`hatch shell dev`), VS Code debugger, and VS Code test runs. +- Enable manually by exporting `MESA_FRAMES_RUNTIME_TYPECHECKING=1` (any of 1/true/yes). +- Use only for development/debugging; adds overheadโ€”disable for performance measurements or large simulations. +- Unset with your shell (e.g. `unset`/`Remove-Item Env:` depending on shell) to turn it off. - No manual setup needed in these environments! + Example for a one-off test run: - For more details on runtime type checking, see the [Development Guidelines](https://projectmesa.github.io/mesa-frames/development/). + ```sh + MESA_FRAMES_RUNTIME_TYPECHECKING=1 uv run pytest -q + ``` #### **Step 6: Documentation Updates (If Needed)** ๐Ÿ“– @@ -135,8 +128,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..b46acc8e 100644 --- a/README.md +++ b/README.md @@ -1,173 +1,140 @@ -# 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://github.com/projectmesa/mesa-frames/blob/main/benchmarks/README.md) -If you want to install it into an existing environment: +**mesa-frames consistently outperforms classic Mesa across both toy and canonical ABMs.** -```bash -conda activate myenv -``` +In the Boltzmann model, it maintains near-constant runtimes even as agent count rises, achieving **up to 10ร— faster execution** at scale. -Then, to install mesa-frames itself: +In the more computation-intensive Sugarscape model, **mesa-frames roughly halves total runtime**. -```bash -pip install -e . -``` +We still have room to optimize performance further (see [Roadmap](#roadmap)). -#### Installing in a Python Virtual Environment +![Benchmark: Boltzmann Wealth](docs/general/plots/boltzmann.svg) -If you want to install it into a new environment: +![Benchmark: Sugarscape IG](docs/general/plots/sugarscape.svg) -```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: +## Quick Start -```bash -source myenv/bin/activate # On Windows, use `myenv\Scripts\activate` -``` +[![Explore the Tutorials](https://img.shields.io/badge/Explore%20the%20Tutorials-๐Ÿ“š-blue?style=for-the-badge)](/mesa-frames/tutorials/2_introductory_tutorial/) -Then, to install mesa-frames itself: +1. **Install** ```bash -pip install -e . + pip install mesa-frames ``` -## Usage - -[![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) +Or for development: -**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. - -[You can find the API documentation here](https://projectmesa.github.io/mesa-frames/api). - -### Creation of an Agent - -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. - -```python -from mesa-frames import AgentSet - -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") - - def give_money(self): - # Active agents are changed to wealthy agents - self.select(self.wealth > 0) +```bash +git clone https://github.com/projectmesa/mesa-frames.git +cd mesa-frames +uv sync --all-extras +``` - # Receiving agents are sampled (only native expressions currently supported) - other_agents = self.df.sample( - n=len(self.active_agents), with_replacement=True - ) +1. **Create a model** - # Wealth of wealthy is decreased by 1 - self["active", "wealth"] -= 1 + ```python + from mesa_frames import AgentSet, Model + import polars as pl - # Compute the income of the other agents (only native expressions currently supported) - new_wealth = other_agents.group_by("unique_id").len() + class MoneyAgents(AgentSet): + def __init__(self, n: int, model: Model): + super().__init__(model) + self += pl.DataFrame({"wealth": pl.ones(n, eager=True)}) - # Add the income to the other agents - self[new_wealth, "wealth"] += new_wealth["len"] -``` + 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 the Model + def step(self): + self.do("give_money") -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. + class MoneyModelDF(Model): + def __init__(self, N: int): + super().__init__() + self.sets += MoneyAgents(N, self) -```python -from mesa-frames import Model + def step(self): + self.sets.do("step") + ``` -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") +## Roadmap - def run_model(self, n): - for _ in range(n): - self.step() -``` +> Community contributions welcome โ€” see the [full roadmap](mesa-frames/roadmap) -## What's Next? ๐Ÿ”ฎ +- Transition to LazyFrames for optimization and GPU support +- Auto-vectorize existing Mesa models via decorator +- Increase possible Spaces (Network, Continuous...) +- Refine the API to align to Mesa -- 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/ROADMAP.md b/ROADMAP.md index 7dd953f5..e70292cf 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,67 +1,63 @@ -# Roadmap ๐Ÿ—บ๏ธ +# Roadmap -This document outlines the development roadmap for the mesa-frames project. It provides insights into our current priorities, upcoming features, and long-term vision. +This document outlines the near-term roadmap for mesa-frames as of October 2025. -## 0.1.0 Stable Release Goals ๐ŸŽฏ +### 1) LazyFrames for Polars + GPU -### 1. Transitioning polars implementation from eager API to lazy API +Switch Polars usage from eager to `LazyFrame` to enable better query optimization and GPU acceleration. -One of our major priorities was to move from pandas to polars as the primary dataframe backend. This transition was motivated by performance considerations. -Now we should transition to using the lazily evaluated version of polars. +Related issues: -**Related issues:** [#10: GPU integration: Dask, cuda (cudf) and RAPIDS (Polars)](https://github.com/projectmesa/mesa-frames/issues/10), [#89: Investigate using Ibis for the common interface library to any DF backend](https://github.com/projectmesa/mesa-frames/issues/89), [#52: Use of LazyFrames for Polars implementation](https://github.com/projectmesa/mesa-frames/issues/52) +- [#52: Use of LazyFrames for Polars implementation](https://github.com/projectmesa/mesa-frames/issues/52) -#### Progress and Next Steps +- [#144: Switch to LazyFrame for Polars implementation (PR)](https://github.com/projectmesa/mesa-frames/pull/144) -- We are exploring [Ibis](https://ibis-project.org/) or [narwhals](https://github.com/narwhals-dev/narwhals) as a common interface library that could support multiple backends (Polars, DuckDB, Spark etc.), but since most of the development is currently in polars, we will currently continue using Polars. -- We're transitioning to the lazy API, mainly in order to use GPU acceleration +- [#89: Investigate Ibis or Narwhals for backend flexibility](https://github.com/projectmesa/mesa-frames/issues/89) -### 2. Handling Concurrency Management +- [#122: Deprecate DataFrameMixin (remove during LazyFrames refactor)](https://github.com/projectmesa/mesa-frames/issues/122) -A critical aspect of agent-based models is efficiently managing concurrent agent movements, especially when multiple agents attempt to move to the same location simultaneously. We aim to implement abstractions that handle these concurrency conditions automatically. +Progress and next steps: -**Related issues:** [#108: Adding abstraction of optimal agent movement](https://github.com/projectmesa/mesa-frames/issues/108), [#48: Emulate RandomActivation with DataFrame.rolling](https://github.com/projectmesa/mesa-frames/issues/48) +- Land [#144](https://github.com/projectmesa/mesa-frames/pull/144) and convert remaining eager paths to lazy. -#### Sugarscape Example of Concurrency Issues +- Validate GPU execution paths and benchmark improvements. -Testing with many potential collisions revealed a specific issue: +- Revisit Ibis/Narwhals after LazyFrame stabilization. -**Problem scenario:** +- Fold DataFrameMixin removal into the LazyFrames transition ([#122](https://github.com/projectmesa/mesa-frames/issues/122)). -- Consider two agents targeting the same cell: - - A mid-priority agent (higher in the agent order) - - A low-priority agent (lower in the agent order) -- The mid-priority agent has low preference for the cell -- The low-priority agent has high preference for the cell -- Without accounting for priority: - - The mid-priority agent's best moves kept getting "stolen" by higher priority agents - - This forced it to resort to lower preference target cells - - However, these lower preference cells were often already taken by lower priority agents in previous iterations +--- -**Solution approach:** +### 2) AgentSet Enhancements -- Implement a "priority" count to ensure that each action is "legal" -- This prevents race conditions but requires recomputing the priority at each iteration -- Current implementation may be slower than Numba due to this overhead -- After the Ibis refactoring, we can investigate if lazy evaluation can help mitigate this performance issue +Expose movement methods from `AgentContainer` and provide optimized utilities for "move to optimal" workflows. -The Sugarscape example demonstrates the need for this abstraction, as multiple agents often attempt to move to the same cell simultaneously. By generalizing this functionality, we can eliminate the need for users to implement complex conflict resolution logic repeatedly. +Related issues: -#### Progress and Next Steps +- [#108: Adding abstraction of optimal agent movement](https://github.com/projectmesa/mesa-frames/issues/108) -- Create utility functions in `DiscreteSpace` and `AgentSetRegistry` to move agents optimally based on specified attributes -- Provide built-in resolution strategies for common concurrency scenarios -- Ensure the implementation works efficiently with the vectorized approach of mesa-frames +- [#118: Adds move_to_optimal in DiscreteSpaceDF (PR)](https://github.com/projectmesa/mesa-frames/pull/118) -### Additional 0.1.0 Goals +- [#82: Add movement methods to AgentContainer](https://github.com/projectmesa/mesa-frames/issues/82) -- Complete core API stabilization -- Completely mirror mesa's functionality -- Improve documentation and examples -- Address outstanding bugs and performance issues +Next steps: -## Beyond 0.1.0 +- Consolidate movement APIs under `AgentContainer`. -Future roadmap items will be added as the project evolves and new priorities emerge. +- Keep conflict resolution simple, vectorized, and well-documented. -We welcome community feedback on our roadmap! Please open an issue if you have suggestions or would like to contribute to any of these initiatives. +--- + +### 3) Research & Publication + +JOSS paper preparation and submission. + +Related items: + +- [#90: JOSS paper for the package](https://github.com/projectmesa/mesa-frames/issues/90) + +- [#107: paper - Adding Statement of Need (PR)](https://github.com/projectmesa/mesa-frames/pull/107) + +--- + +See [our contribution guide](/mesa-frames/contributing/) and browse all open items at diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..56ba998a --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,88 @@ +# Benchmarks + +Performance benchmarks compare Mesa Frames backends ("frames") with classic Mesa ("mesa") +implementations for a small set of representative models. They help track runtime scaling +and regressions. + +Currently included models: + +- **boltzmann**: Simple wealth exchange ("Boltzmann wealth") model. +- **sugarscape**: Sugarscape Immediate Growback variant (square grid sized relative to agent count). + +## Quick start + +```bash +uv run benchmarks/cli.py +``` + +That command (with defaults) will: + +- Benchmark both models (`boltzmann`, `sugarscape`). +- Use agent counts 1000, 2000, 3000, 4000, 5000. +- Run 100 steps per simulation. +- Repeat each configuration once. +- Save CSV results and generate plots. + +## CLI options + +Invoke `uv run benchmarks/cli.py --help` to see full help. Key options: + +| Option | Default | Description | +| ------ | ------- | ----------- | +| `--models` | `all` | Comma list or `all`; accepted: `boltzmann`, `sugarscape`. | +| `--agents` | `1000:5000:1000` | Single int or range `start:stop:step`. | +| `--steps` | `100` | Steps per simulation run. | +| `--repeats` | `1` | How many repeats per (model, backend, agents) config. Seed increments per repeat. | +| `--seed` | `42` | Base RNG seed. Incremented by repeat index. | +| `--save / --no-save` | `--save` | Persist perโ€‘model CSVs. | +| `--plot / --no-plot` | `--plot` | Generate scaling plots (PNG + possibly other formats). | +| `--results-dir` | `benchmarks/results` | Root directory that will receive a timestamped subdirectory. | + +Range parsing: `A:B:S` includes `A, A+S, ... <= B`. Final value > B is dropped. + +## Output layout + +Each invocation uses a single UTC timestamp, e.g. `20251016_173702`: + +```text +benchmarks/ + results/ + 20251016_173702/ + boltzmann_perf_20251016_173702.csv + sugarscape_perf_20251016_173702.csv + plots/ + boltzmann_runtime_20251016_173702_dark.png + sugarscape_runtime_20251016_173702_dark.png + ... (other themed variants if enabled) +``` + +CSV schema (one row per completed run): + +| Column | Meaning | +| ------ | ------- | +| `model` | Model key (`boltzmann`, `sugarscape`). | +| `backend` | `mesa` or `frames`. | +| `agents` | Agent count for that run. | +| `steps` | Steps simulated. | +| `seed` | Seed used (base seed + repeat index). | +| `repeat_idx` | Repeat counter starting at 0. | +| `runtime_seconds` | Wall-clock runtime for that run. | +| `timestamp` | Shared timestamp identifier for the benchmark batch. | + +## Performance tips + +- Ensure the environment variable `MESA_FRAMES_RUNTIME_TYPECHECKING` is **unset** or set to `0` / `false` when collecting performance numbers. Enabling it adds runtime type validation overhead and the CLI will warn you. +- Run multiple repeats (`--repeats 5`) to smooth variance. + +## Extending benchmarks + +To benchmark an additional model: + +1. Add or import both a Mesa implementation and a Frames implementation exposing a `simulate(agents:int, steps:int, seed:int|None, ...)` function. +2. Register it in `benchmarks/cli.py` inside the `MODELS` dict with two backends (names must be `mesa` and `frames`). +3. Ensure any extra spatial parameters are derived from `agents` inside the runner lambda (see sugarscape example). +4. Run the CLI to verify new CSV columns still align. + +## Related documentation + +See `docs/user-guide/5_benchmarks.md` (user-facing narrative) and the main project `README.md` for overall context. diff --git a/benchmarks/cli.py b/benchmarks/cli.py new file mode 100644 index 00000000..c9beb7d2 --- /dev/null +++ b/benchmarks/cli.py @@ -0,0 +1,266 @@ +"""Typer CLI for running mesa vs mesa-frames performance benchmarks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timezone +import os +from pathlib import Path +from time import perf_counter +from typing import Literal, Annotated, Protocol, Optional + +import math +import polars as pl +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 +from examples.plotting import ( + plot_performance as _examples_plot_performance, +) + +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=lambda agents, steps, seed=None: sugarscape_mesa.simulate( + agents=agents, + steps=steps, + width=int(max(20, math.ceil((agents) ** 0.5) * 2)), + height=int(max(20, math.ceil((agents) ** 0.5) * 2)), + seed=seed, + ), + ), + Backend( + name="frames", + # Benchmarks expect a runner signature (agents:int, steps:int, seed:int|None) + # Sugarscape frames simulate requires width/height; choose square close to agent count. + runner=lambda agents, steps, seed=None: sugarscape_frames.simulate( + agents=agents, + steps=steps, + width=int(max(20, math.ceil((agents) ** 0.5) * 2)), + height=int(max(20, math.ceil((agents) ** 0.5) * 2)), + seed=seed, + ), + ), + ], + ), +} + + +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: + """Wrap examples.plotting.plot_performance to ensure consistent theming. + + The original benchmark implementation used simple seaborn styles (whitegrid / darkgrid). + Our example plotting utilities define a much darker, high-contrast *true* dark theme + (custom rc params overriding bg/fg colors). Reuse that logic here so the + benchmark dark plots match the example dark plots users see elsewhere. + """ + if df.is_empty(): + return + stem = f"{model_name}_runtime_{timestamp}" + _examples_plot_performance( + df.select(["agents", "runtime_seconds", "backend"]), + output_dir=output_dir, + stem=stem, + # Prefer more concise, publication-style wording + title=f"{model_name.title()} runtime scaling", + ) + + +@app.command() +def run( + models: Annotated[ + str, + typer.Option( + help="Models to benchmark: boltzmann, sugarscape, or all", + callback=_parse_models, + ), + ] = "all", + agents: Annotated[ + str, + 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=( + "Base directory for benchmark outputs. A timestamped subdirectory " + "(e.g. results/20250101_120000) is created with CSV files at the root " + "and a 'plots/' subfolder for images." + ), + ), + ] = Path(__file__).resolve().parent / "results", +) -> None: + """Run performance benchmarks for the models models.""" + runtime_typechecking = os.environ.get("MESA_FRAMES_RUNTIME_TYPECHECKING", "") + if runtime_typechecking and runtime_typechecking.lower() not in {"0", "false"}: + typer.secho( + "Warning: MESA_FRAMES_RUNTIME_TYPECHECKING is enabled; benchmarks may run significantly slower.", + fg=typer.colors.YELLOW, + ) + rows: list[dict[str, object]] = [] + # Single timestamp per CLI invocation so all model results are co-located. + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + # Create unified output layout: //{CSV files, plots/} + base_results_dir = results_dir + timestamp_dir = (base_results_dir / timestamp).resolve() + plots_subdir: Path = timestamp_dir / "plots" + 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, + } + ) + # Report completion of this run to the CLI + typer.echo( + f"Completed {backend.name} for model={model} agents={agents_count} steps={steps} seed={run_seed} repeat={repeat_idx} in {runtime:.3f}s" + ) + # Finished all runs for this model + typer.echo(f"Finished benchmarking model {model}") + + if not rows: + typer.echo("No benchmark data collected.") + return + df = pl.DataFrame(rows) + if save: + timestamp_dir.mkdir(parents=True, exist_ok=True) + for model in models: + model_df = df.filter(pl.col("model") == model) + csv_path = timestamp_dir / f"{model}_perf_{timestamp}.csv" + model_df.write_csv(csv_path) + typer.echo(f"Saved {model} results to {csv_path}") + if plot: + plots_subdir.mkdir(parents=True, exist_ok=True) + for model in models: + model_df = df.filter(pl.col("model") == model) + _plot_performance(model_df, model, plots_subdir, timestamp) + typer.echo(f"Saved {model} plots under {plots_subdir}") + + typer.echo( + f"Unified benchmark outputs written under {timestamp_dir} (CSV files) and {plots_subdir} (plots)" + ) + + +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..4998fbc2 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,20 @@ 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__", +} + +autosummary_generate_overwrite = True + # -- GitHub link and user guide settings ------------------------------------- github_root = "https://github.com/projectmesa/mesa-frames" @@ -64,7 +92,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 +101,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..4576a204 100644 --- a/docs/api/reference/agents/index.rst +++ b/docs/api/reference/agents/index.rst @@ -3,15 +3,187 @@ 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.model + AgentSet.random + AgentSet.space + 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: diff --git a/docs/api/reference/datacollector.rst b/docs/api/reference/datacollector.rst index bdf38cfd..e95e2672 100644 --- a/docs/api/reference/datacollector.rst +++ b/docs/api/reference/datacollector.rst @@ -1,10 +1,68 @@ 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: diff --git a/docs/api/reference/model.rst b/docs/api/reference/model.rst index 099e601b..0fb12b55 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 +----------- + +The ``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: diff --git a/docs/api/reference/space/index.rst b/docs/api/reference/space/index.rst index 8741b6b6..aa8692f0 100644 --- a/docs/api/reference/space/index.rst +++ b/docs/api/reference/space/index.rst @@ -4,8 +4,98 @@ 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__ + Grid.copy + + .. rubric:: Placement & Movement + + .. autosummary:: + :nosignatures: + :toctree: + + Grid.place_agents + Grid.move_agents + Grid.place_to_empty + Grid.place_to_available + Grid.move_to_empty + Grid.move_to_available + + .. rubric:: Sampling & Queries + + .. autosummary:: + :nosignatures: + :toctree: + + Grid.get_neighbors + Grid.get_directions + Grid.get_distances + Grid.sample_cells + Grid.random_pos + Grid.is_empty + Grid.is_available + Grid.is_full + + .. rubric:: Accessors & Metadata + + .. autosummary:: + :nosignatures: + :toctree: + + Grid.dimensions + Grid.neighborhood_type + Grid.torus + Grid.remaining_capacity + Grid.agents + Grid.model + Grid.random + + .. tab-item:: Full API + + .. autoclass:: Grid + :autosummary: + :autosummary-nosignatures: diff --git a/docs/general/changelog.md b/docs/general/changelog.md new file mode 100644 index 00000000..33ece5d4 --- /dev/null +++ b/docs/general/changelog.md @@ -0,0 +1 @@ +{% include-markdown "../../CHANGELOG.md" %} diff --git a/docs/general/development/index.md b/docs/general/development/index.md deleted file mode 100644 index 8dbfd993..00000000 --- a/docs/general/development/index.md +++ /dev/null @@ -1,170 +0,0 @@ -# Development Guidelines - -## Runtime Type Checking ๐Ÿ” - -mesa-frames includes optional runtime type checking using [beartype](https://github.com/beartype/beartype) for development and debugging purposes. This feature helps catch type-related errors early during development and testing. - -!!! tip "Automatically Enabled" - Runtime type checking is **automatically enabled** in the following scenarios: - - - **Hatch development environment** (`hatch shell dev`) โ€” via `pyproject.toml` configuration - - **VS Code debugging** โ€” when using the debugger (`F5` or "Python Debugger: Current File") - - **VS Code testing** โ€” when running tests through VS Code's testing interface - - No manual setup required in these environments! - -### Development Environment Setup - -#### Option 1: Hatch Development Environment (Recommended) - -The easiest way to enable runtime type checking is to use Hatch's development environment: - -```bash -# Enter the development environment (auto-enables runtime type checking) -hatch shell dev - -# Verify it's enabled -python -c "import os; print('Runtime type checking:', os.getenv('MESA_FRAMES_RUNTIME_TYPECHECKING'))" -# โ†’ Runtime type checking: true -``` - -#### Option 2: Manual Environment Variable - -For other development setups, you can manually enable runtime type checking: - -Runtime type checking can be enabled by setting the `MESA_FRAMES_RUNTIME_TYPECHECKING` environment variable: - -```bash -export MESA_FRAMES_RUNTIME_TYPECHECKING=1 -# or -export MESA_FRAMES_RUNTIME_TYPECHECKING=true -# or -export MESA_FRAMES_RUNTIME_TYPECHECKING=yes -``` - -### Usage Examples - -!!! info "Automatic Activation" - If you're using **Hatch dev environment**, **VS Code debugging**, or **VS Code testing**, runtime type checking is already enabled automatically. The examples below are for manual activation in other scenarios. - -#### For Development and Testing - -```bash -# Enable runtime type checking for testing -MESA_FRAMES_RUNTIME_TYPECHECKING=1 uv run pytest - -# Enable runtime type checking for running scripts -MESA_FRAMES_RUNTIME_TYPECHECKING=1 uv run python your_script.py -``` - -#### In Your IDE or Development Environment - -**VS Code** (Already Configured): - -- **Debugging**: Runtime type checking is automatically enabled when using VS Code's debugger -- **Testing**: Automatically enabled when running tests through VS Code's testing interface -- **Manual override**: You can also add it manually in `.vscode/settings.json`: - - ```json - { - "python.env": { - "MESA_FRAMES_RUNTIME_TYPECHECKING": "1" - } - } - ``` - -**PyCharm**: -In your run configuration, add the environment variable: - -```bash -MESA_FRAMES_RUNTIME_TYPECHECKING=1 -``` - -### How It Works - -When enabled, the runtime type checking system: - -1. **Automatically instruments** all mesa-frames packages with beartype decorators -2. **Validates function arguments** and return values at runtime -3. **Provides detailed error messages** when type mismatches occur -4. **Helps catch type-related bugs** during development - -### Requirements - -Runtime type checking requires the optional `beartype` dependency: - -```bash -# Install beartype for runtime type checking -uv add beartype -# or -pip install beartype -``` - -!!! note "Optional Dependency" - If `beartype` is not installed and runtime type checking is enabled, mesa-frames will issue a warning and continue without type checking. - -### Performance Considerations - -!!! warning "Development Only" - Runtime type checking adds significant overhead and should **only be used during development and testing**. Do not enable it in production environments. - -The overhead includes: - -- Function call interception and validation -- Type checking computations at runtime -- Memory usage for type checking infrastructure - -### When to Use Runtime Type Checking - -โœ… **Automatically enabled (recommended):** - -- Hatch development environment (`hatch shell dev`) -- VS Code debugging sessions -- VS Code test execution -- Contributing to mesa-frames development - -โœ… **Manual activation (when needed):** - -- Development and debugging in other IDEs -- Writing new features outside VS Code -- Running unit tests from command line -- Troubleshooting type-related issues - -โŒ **Not recommended for:** - -- Production deployments -- Performance benchmarking -- Large-scale simulations -- Final model runs - -### Troubleshooting - -If you encounter issues with runtime type checking: - -1. **Check beartype installation:** - - ```bash - uv run python -c "import beartype; print(beartype.__version__)" - ``` - -2. **Verify environment variable:** - - ```bash - echo $MESA_FRAMES_RUNTIME_TYPECHECKING - ``` - -3. **For automatic configurations:** - - **Hatch dev**: Ensure you're in the dev environment (`hatch shell dev`) - - **VS Code debugging**: Check that the debugger configuration in `.vscode/launch.json` includes the environment variable - - **VS Code testing**: Verify that `.env.test` file exists and contains `MESA_FRAMES_RUNTIME_TYPECHECKING=true` - -4. **Check for warnings** in your application logs - -5. **Disable temporarily** if needed: - - ```bash - unset MESA_FRAMES_RUNTIME_TYPECHECKING - ``` - -!!! tip "Pro Tip" - Runtime type checking is particularly useful when developing custom AgentSet implementations or working with complex DataFrame operations where type safety is crucial. 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/plots/boltzmann.svg b/docs/general/plots/boltzmann.svg new file mode 100644 index 00000000..f21ca936 --- /dev/null +++ b/docs/general/plots/boltzmann.svg @@ -0,0 +1,1083 @@ + + + + + + + + 2025-10-16T19:57:07.933517 + image/svg+xml + + + Matplotlib v3.10.5, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/general/plots/sugarscape.svg b/docs/general/plots/sugarscape.svg new file mode 100644 index 00000000..679002c9 --- /dev/null +++ b/docs/general/plots/sugarscape.svg @@ -0,0 +1,1091 @@ + + + + + + + + 2025-10-16T19:57:08.355947 + image/svg+xml + + + Matplotlib v3.10.5, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/general/tutorials/2_introductory_tutorial.py b/docs/general/tutorials/2_introductory_tutorial.py new file mode 100644 index 00000000..92f6f1f9 --- /dev/null +++ b/docs/general/tutorials/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/tutorials/3_advanced_tutorial.py b/docs/general/tutorials/3_advanced_tutorial.py new file mode 100644 index 00000000..bf501ec3 --- /dev/null +++ b/docs/general/tutorials/3_advanced_tutorial.py @@ -0,0 +1,1648 @@ +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={ + "sugar": "sugar", + "metabolism": "metabolism", + "vision": "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 = 20 +GRID_HEIGHT = 20 +NUM_AGENTS = 100 +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/tutorials/4_datacollector.py b/docs/general/tutorials/4_datacollector.py new file mode 100644 index 00000000..16d9837b --- /dev/null +++ b/docs/general/tutorials/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/0_getting-started.md b/docs/general/user-guide/0_getting-started.md index 1edc1587..caebc1c9 100644 --- a/docs/general/user-guide/0_getting-started.md +++ b/docs/general/user-guide/0_getting-started.md @@ -15,20 +15,20 @@ Objects can be easily subclassed to respect mesa's object-oriented philosophy. ### Vectorized Operations โšก -mesa-frames leverages the power of vectorized operations provided by DataFrame libraries: +`mesa-frames` leverages **Polars** to replace Python loops with **column-wise expressions** executed in native Rust. +This allows you to update all agents simultaneously, the main source of `mesa-frames`' performance advantage. -- Operations are performed on entire columns of data at once -- This approach is significantly faster than iterating over individual agents -- Complex behaviors can be expressed in fewer lines of code +Unlike traditional `mesa` models, where the **activation order** of agents can affect results (see [Comer, 2014](http://mars.gmu.edu/bitstream/handle/1920/9070/Comer_gmu_0883E_10539.pdf)), +`mesa-frames` processes all agents **in parallel by default**. +This removes order-dependent effects, though you should handle conflicts explicitly when sequential logic is required. -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). +!!! tip "Best practice" + Always start by expressing agent logic in a vectorized form. + Fall back to loops only when ordering or conflict resolution is essential. -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. +For a deeper understanding of vectorization and why it accelerates computation, see: -Check out these resources to understand vectorization and why it speeds up the code: - -- [What is vectorization?](https://stackoverflow.com/a/1422181) -- [Vectorization Explained, Step by Step](https://machinelearningcompass.com/machine_learning_math/vectorization/) +- [How vectorization speeds up your Python code โ€” PythonSpeed](https://pythonspeed.com/articles/vectorization-python) Here's a comparison between mesa-frames and mesa: @@ -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.df.sample( n=len(self.active_agents), with_replacement=True ) @@ -92,10 +92,10 @@ 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.df.sample(n=len(self.active_agents), with_replacement=True) self[givers, "wealth"] -= 1 - new_wealth = receivers.groupby("unique_id").count() - self[new_wealth["unique_id"], "wealth"] += new_wealth["count"] + new_wealth = receivers.group_by("unique_id").len() + self[new_wealth["unique_id"], "wealth"] += new_wealth["len"] ``` === "mesa" @@ -146,12 +146,85 @@ If you're familiar with mesa, this guide will help you understand the key differ self.schedule.step() ``` -### Transition Tips ๐Ÿ’ก +### From Imperative Code to Behavioral Rules ๐Ÿ’ญ + +When scientists describe an ABM-like process they typically write a **system of state-transition functions**: + +$$ +x_i(t+1) = f_i\big(x_i(t),\; \mathcal{N}(i,t),\; E(t)\big) +$$ + +Here, $x_i(t)$ is the agentโ€™s state, $\mathcal{N}(i,t)$ its neighborhood or local environment, and $E(t)$ a global environment; $f_i$ is the behavioral law. + +In classic `mesa`, agent behavior is implemented through explicit loops: each agent individually gathers information from its neighbors, computes its next state, and often stores this in a buffer to ensure synchronous updates. The behavioral law $f_i$ is distributed across multiple steps: neighbor iteration, temporary buffers, and scheduling logic, resulting in procedural, step-by-step control flow. + +In `mesa-frames`, these stages are unified into a single vectorized transformation. Agent interactions, state transitions, and updates are expressed as DataFrame operations (such as joins, group-bys, and column expressions) allowing all agents to process perceptions and commit actions simultaneously. This approach centralizes the behavioral law $f_i$ into concise, declarative rules, improving clarity and performance. + +#### Example: Network contagion (Linear Threshold) + +Behavioral rule: a node activates if the number of active neighbors โ‰ฅ its threshold. + +=== "mesa-frames" + + Single vectorized transformation. A join brings in source activity, a group-by aggregates exposures per destination, and a column expression applies the activation equation and commits in one pass, no explicit loops or staging structure needed. + + ```python + class Nodes(AgentSet): + # self.df columns: agent_id, active (bool), theta (int) + # self.model.space.edges: DataFrame[src, dst] + def step(self): + E = self.model.space.edges # [src, dst] + # Exposure: active neighbors per dst (vectorized join + groupby) + exposures = ( + E.join( + self.df.select(pl.col("agent_id").alias("src"), + pl.col("active").alias("src_active")), + on="src", how="left" + ) + .with_columns(pl.col("src_active").fill_null(False)) + .group_by("dst") + .agg(pl.col("src_active").sum().alias("k_active")) + ) + # Behavioral equation applied to all agents, committed in-place + self.df = ( + self.df + .join(exposures, left_on="agent_id", right_on="dst", how="left") + .with_columns(pl.col("k_active").fill_null(0)) + .with_columns( + (pl.col("active") | (pl.col("k_active") >= pl.col("theta"))) + .alias("active") + ) + .drop(["k_active", "dst"]) + ) + ``` + +=== "mesa" + + Two-phase imperative procedure. Each agent loops over its neighbors to count active ones (exposure), stores a provisional next state to avoid premature mutation, then a separate pass commits all buffered states for synchronicity. + + ```python + class Node(mesa.Agent): + def step(self): + # (1) Gather exposure: count active neighbors right now + k_active = sum( + 1 for j in self.model.G.neighbors(self.unique_id) + if self.model.id2agent[j].active + ) + # (2) Compute next state (don't mutate yet to stay synchronous) + self.next_active = self.active or (k_active >= self.theta) + + # Second pass (outside the agent method) performs the commit: + for a in model.agents: + a.active = a.next_active + ``` -1. **Think in Sets ๐ŸŽญ**: Instead of individual agents, think about operations on groups of agents. -2. **Leverage DataFrame Operations ๐Ÿ› ๏ธ**: Familiarize yourself with Polars operations for efficient agent manipulation. -3. **Vectorize Logic ๐Ÿš…**: Convert loops and conditionals to vectorized operations where possible. -4. **Use AgentSets ๐Ÿ“ฆ**: Group similar agents into AgentSets instead of creating many individual agent classes. +!!! tip "Transition tips โ€” quick summary" + 1. Think in sets: operate on AgentSets/DataFrames, not per-agent objects. + 2. Write transitions as Polars column expressions; avoid Python loops. + 3. Use joins + group-bys to compute interactions/exposure across relations. + 4. Commit state synchronously in one vectorized pass. + 5. Group similar agents into one AgentSet with typed columns. + 6. Use UDFs or staged/iterative patterns only for true race/conflict cases. ### Handling Race Conditions ๐Ÿ @@ -163,4 +236,4 @@ When simultaneous activation is not possible, you need to handle race conditions 2. **Looping Mechanism ๐Ÿ”**: Implement a looping mechanism on vectorized operations. -For a more detailed implementation of handling race conditions, please refer to the `examples/sugarscape-ig` in the mesa-frames repository. This example demonstrates how to implement the Sugarscape model with instantaneous growback, which requires careful handling of sequential agent actions. +For a more detailed implementation of handling race conditions, see the [Advanced Tutorial](../tutorials/3_advanced_tutorial.ipynb). It walks through the Sugarscape model with instantaneous growback and shows practical patterns for staged vectorization and conflict resolution. 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/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/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/5_benchmarks.md b/docs/general/user-guide/5_benchmarks.md deleted file mode 100644 index 61fca87b..00000000 --- a/docs/general/user-guide/5_benchmarks.md +++ /dev/null @@ -1,21 +0,0 @@ -# Performance Boost ๐ŸŽ๏ธ๐Ÿ’จ - -mesa-frames offers significant performance improvements over the original mesa framework. Here are some benchmark results for different models: - -## Boltzmann Wealth Model ๐Ÿ’ฐ - -[View the benchmark script](https://github.com/projectmesa/mesa-frames/blob/main/examples/boltzmann_wealth/performance_plot.py) - -### Comparison with mesa - -![Performance Graph BW](https://github.com/projectmesa/mesa-frames/raw/main/examples/boltzmann_wealth/boltzmann_no_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) - -## SugarScape with Instantaneous Growback ๐Ÿฌ - -[View the benchmark script](https://github.com/projectmesa/mesa-frames/blob/main/examples/sugarscape_ig/performance_comparison.py) - -![Performance Graph SS IG](https://github.com/projectmesa/mesa-frames/raw/main/examples/sugarscape_ig/mesa_comparison.png) diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 00000000..359bbaf8 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,106 @@ +# Examples + +This directory contains runnable example models and shared plotting/utilities +used in the tutorials and benchmarks. Each example provides **two backends**: + +- `mesa` (classic Mesa, object-per-agent) +- `frames` (Mesa Frames, vectorised agent sets / dataframe-centric) + +They expose a consistent Typer CLI so you can compare outputs and timings. + +## Contents + +```text +examples/ + boltzmann_wealth/ + backend_mesa.py # Mesa implementation + CLI (simulate() + run) + backend_frames.py # Frames implementation + CLI (simulate() + run) + sugarscape_ig/ + backend_mesa/ # Mesa Sugarscape (agents + model + CLI) + backend_frames/ # Frames Sugarscape (agents + model + CLI) + plotting.py # Shared plotting helpers (Seaborn + dark theme) + utils.py # Small dataclasses for simulation results +``` + +## Quick start + +Always run via `uv` from the project root. The simplest way to run an example +backend is to execute the module: + +```bash +uv run examples/boltzmann_wealth/backend_frames.py +``` + +Each command will: + +1. Print a short banner with configuration. +2. Run the simulation and show elapsed time. +3. Emit a tail of the collected metrics (e.g. last 5 Gini values). +4. Save CSV metrics and optional plots in a timestamped directory under that + example's `results/` folder (unless overridden by `--results-dir`). + +## CLI symmetry + +Both backends accept similar options: + +- `--agents` (population size) +- `--steps` (number of simulated steps) +- `--seed` (optional RNG seed; Mesa backend resets model RNG) +- `--plot / --no-plot` (toggle plot generation) +- `--save-results / --no-save-results` (persist CSV outputs) +- `--results-dir` (override auto-created timestamped folder) + +The Frames Boltzmann backend stores model metrics in a Polars DataFrame via +`mesa_frames.DataCollector`; the Mesa backend uses the standard `mesa.DataCollector` +returning pandas DataFrames, then converts to Polars only for plotting so plots +look identical. + +## Data and metrics + +The saved CSV layout (Frames) places `model.csv` in the results directory with +columns like: `step, gini, `. +The Mesa implementations write +compatible CSVs. + +## Plotting helpers + +`examples/plotting.py` provides: + +- `plot_model_metrics(df, output_dir, stem, title, subtitle, agents, steps)` + Produces dark theme line plots of model-level metrics (currently Gini) and + stores PNG files under `output_dir` with names like `gini__dark.png`. +- `plot_performance(df, output_dir, stem, title)` used by `benchmarks/cli.py` to + generate runtime scaling plots. + +The dark theme matches the styling used in the documentation for visual +consistency. + +## Interacting programmatically + +Instead of using the CLIs you can import the simulation entry points directly: + +```python +from examples.boltzmann_wealth import backend_frames as bw_frames +result = bw_frames.simulate(agents=2000, steps=100, seed=123) +polars_df = result.datacollector.data["model"] # Polars DataFrame of metrics +``` + +Each `simulate()` returns a small dataclass (`FramesSimulationResult` or +`MesaSimulationResult`) holding the respective `DataCollector` instance so you +can further analyse the collected data. + +## Tips + +- To compare backends fairly, disable runtime type checking when measuring performance: + set environment variable `MESA_FRAMES_RUNTIME_TYPECHECKING=0`. +- Use the same `--seed` across runs for reproducible trajectories (given the + stochastic nature of agent interactions). +- Larger Sugarscape grids (width/height) increase memory and runtime; choose + sizes proportional to the square root of agent count for balanced density. + +## Adding Examples + +You can adapt these scripts to prototype new models: copy a backend pair, +rename the module, and implement your agent rules while keeping the API +surface (`simulate`, `run`) consistent so tooling and documentation patterns +continue to apply. 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/README.md b/examples/boltzmann_wealth/README.md new file mode 100644 index 00000000..9999fe27 --- /dev/null +++ b/examples/boltzmann_wealth/README.md @@ -0,0 +1,96 @@ +# Boltzmann Wealth Exchange Model + +## Overview + +This example implements a simple wealth exchange ("Boltzmann money") model in two +backends: + +- `backend_frames.py` (Mesa Frames / vectorised `AgentSet`) +- `backend_mesa.py` (classic Mesa / object-per-agent) + +Both expose a Typer CLI with symmetric options so you can compare correctness +and performance directly. + +## Concept + +Each agent starts with 1 unit of wealth. At every step: + +1. Frames backend: all agents with strictly positive wealth become potential donors. + Each donor gives 1 unit of wealth, and a recipient is drawn (with replacement) + for every donating agent. A single vectorised update applies donor losses and + recipient gains. +2. Mesa backend: agents are shuffled and iterate sequentially; each agent with + positive wealth transfers 1 unit to a randomly selected peer. + +The stochastic exchange process leads to an emergent, increasingly unequal +wealth distribution and rising Gini coefficient, typically approaching a stable +level below 1 (due to conservation and continued mixing). + +## Reported Metrics + +The model records per-step population Gini (`gini`). You can extend reporters by +adding lambdas to `model_reporters` in either backend's constructor. + +Notes on interpretation: + +- Early steps: Gini ~ 0 (uniform initial wealth). +- Mid phase: Increasing Gini as random exchanges concentrate wealth. +- Late phase: Fluctuating plateau (a stochastic steady state) โ€” exact level + varies with agent count and RNG seed. + +## Running + +Always run examples from the project root using `uv`: + +```bash +uv run examples/boltzmann_wealth/backend_frames.py --agents 5000 --steps 200 --seed 123 --plot --save-results +uv run examples/boltzmann_wealth/backend_mesa.py --agents 5000 --steps 200 --seed 123 --plot --save-results +``` + +## CLI options + +- `--agents` Number of agents (default 5000) +- `--steps` Simulation steps (default 100) +- `--seed` Optional RNG seed for reproducibility +- `--plot / --no-plot` Generate line plot(s) of Gini +- `--save-results / --no-save-results` Persist CSV metrics +- `--results-dir` Override the auto timestamped directory under `results/` + +Frames backend additionally warns if runtime type checking is enabled because it +slows vectorised operations: set `MESA_FRAMES_RUNTIME_TYPECHECKING=0` for fair +performance comparisons. + +## Outputs + +Each run creates (or uses) a results directory like: + +```text +examples/boltzmann_wealth/results/20251016_173702/ + model.csv # step,gini + gini__dark.png (and possibly other theme variants) +``` + +Tail metrics are printed to console for quick inspection: + +```text +Metrics in the final 5 steps: shape: (5, 2) +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ step โ”† gini โ”‚ +โ”‚ --- โ”† --- โ”‚ +โ”‚ i64 โ”† f64 โ”‚ +โ”œโ”€โ”€โ”€โ”€โ”€โ”€โ”ผโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ค +โ”‚ ... โ”† ... โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +``` + +## Performance & Benchmarking + +Use the shared benchmarking CLI to compare scaling, checkout `benchmarks/README.md`. + +## Programmatic Use + +```python +from examples.boltzmann_wealth import backend_frames as bw_frames +result = bw_frames.simulate(agents=10000, steps=250, seed=42) +metrics = result.datacollector.data["model"] # Polars DataFrame +``` diff --git a/examples/sugarscape_ig/ss_polars/__init__.py b/examples/boltzmann_wealth/__init__.py similarity index 100% rename from examples/sugarscape_ig/ss_polars/__init__.py rename to examples/boltzmann_wealth/__init__.py diff --git a/examples/boltzmann_wealth/backend_frames.py b/examples/boltzmann_wealth/backend_frames.py new file mode 100644 index 00000000..da26dba9 --- /dev/null +++ b/examples/boltzmann_wealth/backend_frames.py @@ -0,0 +1,190 @@ +"""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 os +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), + ) + # Combine donor loss (1 per active agent) and recipient gains in a single adjustment. + gains = recipients.group_by("unique_id").len() + self.df = ( + self.df.join(gains, on="unique_id", how="left") + .with_columns( + ( + pl.col("wealth") + # each active agent loses 1 unit of wealth + + pl.when(pl.col("wealth") > 0).then(-1).otherwise(0) + # each agent gains 1 unit of wealth for each time they were selected as a recipient + + pl.col("len").fill_null(0) + ).alias("wealth") + ) + .drop("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) + # For benchmarks we frequently call simulate() without providing a results_dir. + # Persisting to disk would add unnecessary IO overhead and a missing storage_uri + # currently raises in DataCollector validation. Fallback to in-memory collection + # when no results_dir is supplied; otherwise write CSV files under results_dir. + if results_dir is None: + storage = "memory" + storage_uri = None + else: + storage = "csv" + storage_uri = str(results_dir) + self.datacollector = DataCollector( + model=self, + model_reporters={ + "gini": lambda m: gini(m.sets[0].df), + }, + storage=storage, + 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: + runtime_typechecking = os.environ.get("MESA_FRAMES_RUNTIME_TYPECHECKING", "") + if runtime_typechecking and runtime_typechecking.lower() not in {"0", "false"}: + typer.secho( + "Warning: MESA_FRAMES_RUNTIME_TYPECHECKING is enabled; this run will be slower.", + fg=typer.colors.YELLOW, + ) + 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..17075451 --- /dev/null +++ b/examples/plotting.py @@ -0,0 +1,290 @@ +# 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)") + leg = ax.get_legend() + if leg is not None: + # Remove redundant legend title (backend) for both themes โ€“ label colors already distinguish. + leg.set_title(None) + frame = leg.get_frame() + if theme == "dark": + frame.set_alpha(0.8) + else: # light theme: subtle boxed legend for readability on white grid + frame.set_alpha(0.9) + frame.set_edgecolor("#d0d0d0") + frame.set_linewidth(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/README.md b/examples/sugarscape_ig/README.md new file mode 100644 index 00000000..f33970d5 --- /dev/null +++ b/examples/sugarscape_ig/README.md @@ -0,0 +1,105 @@ +# Sugarscape IG (Instant Growback) + +## Overview + +This directory contains a minimal Instant Growback Sugarscape implementation in +both backends: + +- `backend_frames/` parallel (vectorised) movement variant using Mesa Frames +- `backend_mesa/` sequential (asynchronous) movement variant using classic Mesa + +The Instant Growback (IG) rule sequence is: move -> eat -> regrow -> collect. +Agents harvest sugar, pay metabolism costs, possibly die (starve), and empty +cells instantly regrow to their `max_sugar` value. + +## Concept + +Each agent has integer traits: + +- `sugar` (current stores) +- `metabolism` (per-step consumption) +- `vision` (how far the agent can see in cardinal directions) + +Movement policy (both backends conceptually): + +1. Sense visible cells along N/E/S/W up to `vision` steps (including origin). +2. Rank candidate cells by: (a) sugar (desc), (b) distance (asc), (c) coordinates + as deterministic tie-breaker. +3. Choose highest-ranked empty cell; fall back to origin if none available. + +The Frames parallel variant resolves conflicts by iterative lottery rounds using +rank promotion; the sequential Mesa variant inherently orders moves by shuffled +agent iteration. + +After moving, agents harvest sugar on their cell, pay metabolism, and starved +agents are removed. Empty cells regrow to their `max_sugar` value immediately. + +## Reported Metrics + +Both backends record population-level reporters each step: + +- `mean_sugar` Average sugar per surviving agent. +- `total_sugar` Aggregate sugar held by living agents. +- `agents_alive` Population size (declines as agents starve). +- `gini` Inequality in sugar holdings (0 = equal, higher = more unequal). +- `corr_sugar_metabolism` Pearson correlation (do high-metabolism agents retain sugar?). +- `corr_sugar_vision` Pearson correlation (does greater vision correlate with sugar?). + +Notes on interpretation: + +- `agents_alive` typically decreases until a quasi steady state (metabolism vs regrowth) or total collapse. +- `mean_sugar` and `total_sugar` may stabilise if regrowth balances metabolism. +- Rising `gini` indicates emerging inequality; sustained high values suggest strong positional advantages. +- Correlations near 0 imply weak linear relationships; positive `corr_sugar_vision` suggests high vision aids resource gathering. Negative `corr_sugar_metabolism` can emerge if high metabolism accelerates starvation. + +## Running + +From project root using `uv`: + +```bash +uv run examples/sugarscape_ig/backend_frames/model.py --agents 400 --width 40 --height 40 --steps 60 --seed 123 --plot --save-results +uv run examples/sugarscape_ig/backend_mesa/model.py --agents 400 --width 40 --height 40 --steps 60 --seed 123 --plot --save-results +``` + +## CLI options + +- `--agents` Number of agents (default 400) +- `--width`, `--height` Grid dimensions (default 40x40) +- `--steps` Max steps (default 60) +- `--max-sugar` Initial/regrowth max sugar per cell (default 4) +- `--seed` Optional RNG seed +- `--plot / --no-plot` Generate per-metric plots +- `--save-results / --no-save-results` Persist CSV outputs +- `--results-dir` Override auto timestamped directory under `results/` + +Frames backend warns if `MESA_FRAMES_RUNTIME_TYPECHECKING` is enabled (disable for benchmarks). + +## Outputs + +Example output directory (frames): + +```text +examples/sugarscape_ig/backend_frames/results/20251016_173702/ + model.csv + plots/ + gini__dark.png + agents_alive__dark.png + mean_sugar__dark.png + ... +``` + +`model.csv` columns include: `step`, `mean_sugar`, `total_sugar`, `agents_alive`, +`gini`, `corr_sugar_metabolism`, `corr_sugar_vision`, plus backend-specific bookkeeping. +Mesa backend normalises to the same layout (excluding internal columns). + +## Performance & Benchmarking + +Use the shared benchmarking CLI to compare scaling, checkout `benchmarks/README.md`. + +## Programmatic Use + +```python +from examples.sugarscape_ig.backend_frames import model as sg_frames +res = sg_frames.simulate(agents=500, steps=80, width=50, height=50, seed=42) +metrics = res.datacollector.data["model"] # Polars DataFrame +``` 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_frames/agents.py b/examples/sugarscape_ig/backend_frames/agents.py new file mode 100644 index 00000000..e619df00 --- /dev/null +++ b/examples/sugarscape_ig/backend_frames/agents.py @@ -0,0 +1,626 @@ +"""Agent implementations for the Sugarscape IG example (mesa-frames). + +This module provides the parallel (synchronous) movement variant as in the +advanced tutorial. The code and comments mirror +docs/general/user-guide/3_advanced_tutorial.py. +""" + +from __future__ import annotations + +import numpy as np +import polars as pl + +from mesa_frames import AgentSet, Model + + +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) + + +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 โ”‚ + # โ”‚ --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† --- โ”† f64 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + winners = ( + best_candidates.sort( + ["dim_0_candidate", "dim_1_candidate", "radius", "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 + + +__all__ = [ + "AntsBase", + "AntsParallel", +] diff --git a/examples/sugarscape_ig/backend_frames/model.py b/examples/sugarscape_ig/backend_frames/model.py new file mode 100644 index 00000000..1a5d336b --- /dev/null +++ b/examples/sugarscape_ig/backend_frames/model.py @@ -0,0 +1,499 @@ +"""Mesa-frames implementation of Sugarscape IG with Typer CLI. + +This mirrors the advanced tutorial in docs/general/user-guide/3_advanced_tutorial.py +and exposes a simple CLI to run the parallel update variant, save CSVs, and plot +the Gini trajectory. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +import os +from pathlib import Path +from typing import Annotated +from time import perf_counter + +import numpy as np +import polars as pl +import typer + +from mesa_frames import DataCollector, Grid, Model +from examples.utils import FramesSimulationResult +from examples.plotting import plot_model_metrics + +from examples.sugarscape_ig.backend_frames.agents import AntsBase, AntsParallel + + +# 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. + results_dir : Path | None, optional + Optional directory where CSV/plot outputs will be written. If ``None`` + the model runs without persisting CSVs to disk (in-memory storage). + + 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, + results_dir: Path | 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 + # Benchmarks may run without providing a results_dir; in that case avoid forcing + # a CSV storage backend (which requires a storage_uri) and keep data in memory. + if results_dir is None: + storage = "memory" + storage_uri = None + else: + storage = "csv" + storage_uri = str(results_dir) + 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={ + "sugar": "sugar", + "metabolism": "metabolism", + "vision": "vision", + }, + storage=storage, + storage_uri=storage_uri, + ) + 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}) + + +def simulate( + *, + agents: int, + steps: int, + width: int, + height: int, + max_sugar: int = 4, + seed: int | None = None, + results_dir: Path | None = None, +) -> FramesSimulationResult: + model = Sugarscape( + agent_type=AntsParallel, + n_agents=agents, + width=width, + height=height, + max_sugar=max_sugar, + seed=seed, + results_dir=results_dir, + ) + model.run(steps) + 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.")] = 400, + width: Annotated[int, typer.Option(help="Grid width (columns).")] = 40, + height: Annotated[int, typer.Option(help="Grid height (rows).")] = 40, + steps: Annotated[int, typer.Option(help="Number of model steps to run.")] = 60, + max_sugar: Annotated[int, typer.Option(help="Maximum sugar per cell.")] = 4, + 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 Sugarscape IG (mesa-frames, parallel) with {agents} agents on {width}x{height} for {steps} steps" + ) + runtime_typechecking = os.environ.get("MESA_FRAMES_RUNTIME_TYPECHECKING", "") + if runtime_typechecking and runtime_typechecking.lower() not in {"0", "false"}: + typer.secho( + "Warning: MESA_FRAMES_RUNTIME_TYPECHECKING is enabled; this run will be slower.", + fg=typer.colors.YELLOW, + ) + 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, + width=width, + height=height, + max_sugar=max_sugar, + seed=seed, + results_dir=results_dir, + ) + typer.echo(f"Simulation complete in {perf_counter() - start_time:.2f} seconds") + + model_metrics = result.datacollector.data["model"].drop(["seed", "batch"]) + typer.echo(f"Metrics in the final 5 steps: {model_metrics.tail(5)}") + + if save_results: + result.datacollector.flush() + + if plot: + # Create a subdirectory for per-metric plots under the timestamped + # results directory. For each column in the model metrics (except + # the step index) create a single-metric DataFrame and call the + # shared plotting helper to export light/dark PNG+SVG variants. + plots_dir = results_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + # Determine which columns to plot (preserve 'step' if present). + value_cols = [c for c in model_metrics.columns if c != "step"] + for col in value_cols: + stem = f"{col}_{timestamp}" + single = ( + model_metrics.select(["step", col]) + if "step" in model_metrics.columns + else model_metrics.select([col]) + ) + plot_model_metrics( + single, + plots_dir, + stem, + title=f"Sugarscape IG โ€” {col.capitalize()}", + subtitle=f"mesa-frames backend; seed={result.datacollector.seed}", + agents=agents, + steps=steps, + ) + + typer.echo(f"Saved plots under {plots_dir}") + + typer.echo(f"Saved CSV results under {results_dir}") + + +if __name__ == "__main__": + app() 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/backend_mesa/agents.py b/examples/sugarscape_ig/backend_mesa/agents.py new file mode 100644 index 00000000..657d8d07 --- /dev/null +++ b/examples/sugarscape_ig/backend_mesa/agents.py @@ -0,0 +1,81 @@ +"""Mesa agents for the Sugarscape IG example (sequential/asynchronous update). + +Implements the movement rule (sense along cardinal axes up to `vision`, choose +highest-sugar cell with tie-breakers by distance and coordinates). Eating, +starvation, and regrowth are orchestrated by the model to preserve the order +move -> eat -> regrow -> collect, mirroring the tutorial schedule. +""" + +from __future__ import annotations + +from typing import Tuple + +import mesa + + +class AntAgent(mesa.Agent): + """Sugarscape ant with sugar/metabolism/vision traits and movement.""" + + def __init__( + self, + model: Sugarscape, + *, + sugar: int, + metabolism: int, + vision: int, + ) -> None: + super().__init__(model) + self.sugar = int(sugar) + self.metabolism = int(metabolism) + self.vision = int(vision) + + # --- Movement helpers (sequential/asynchronous) --- + + def _visible_cells(self, origin: tuple[int, int]) -> list[tuple[int, int]]: + x0, y0 = origin + width, height = self.model.width, self.model.height + cells: list[tuple[int, int]] = [origin] + for step in range(1, self.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]) -> tuple[int, int]: + # Highest sugar; tie-break by Manhattan distance from origin; then coords. + best_cell = origin + best_sugar = int(self.model.sugar_current[origin[0], origin[1]]) + best_distance = 0 + ox, oy = origin + for cx, cy in self._visible_cells(origin): + # Block occupied cells except the origin (own cell allowed as fallback). + if (cx, cy) != origin and not self.model.grid.is_cell_empty((cx, cy)): + continue + sugar_here = int(self.model.sugar_current[cx, cy]) + distance = abs(cx - ox) + abs(cy - oy) + better = False + if sugar_here > best_sugar: + better = True + elif sugar_here == best_sugar: + if distance < best_distance: + better = True + elif distance == best_distance and (cx, cy) < best_cell: + better = True + if better: + best_cell = (cx, cy) + best_sugar = sugar_here + best_distance = distance + return best_cell + + def move(self) -> None: + best = self._choose_best_cell(self.pos) + if best != self.pos: + self.model.grid.move_agent(self, best) + + +__all__ = ["AntAgent"] diff --git a/examples/sugarscape_ig/backend_mesa/model.py b/examples/sugarscape_ig/backend_mesa/model.py new file mode 100644 index 00000000..6e62137a --- /dev/null +++ b/examples/sugarscape_ig/backend_mesa/model.py @@ -0,0 +1,315 @@ +"""Mesa implementation of Sugarscape IG with Typer CLI (sequential update). + +Follows the same structure as the Boltzmann Mesa example: `simulate()` and a +`run` CLI command that saves CSV results and plots the Gini trajectory. The +model updates in the order move -> eat -> regrow -> collect, matching the +tutorial schedule. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path +from typing import Annotated +from collections.abc import Iterable +from time import perf_counter + +import mesa +from mesa.datacollection import DataCollector +from mesa.space import SingleGrid +import numpy as np +import pandas as pd +import polars as pl +import typer + +from examples.utils import MesaSimulationResult +from examples.plotting import plot_model_metrics + +from examples.sugarscape_ig.backend_mesa.agents import AntAgent + + +def _safe_corr(x: np.ndarray, y: np.ndarray) -> float: + """Safely compute Pearson correlation between two 1-D arrays. + + Mirrors the Frames helper: returns nan for degenerate inputs. + """ + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + 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]) + + +def corr_sugar_metabolism(model: Sugarscape) -> float: + sugars = np.fromiter((a.sugar for a in model.agent_list), dtype=float) + mets = np.fromiter((a.metabolism for a in model.agent_list), dtype=float) + return _safe_corr(sugars, mets) + + +def corr_sugar_vision(model: Sugarscape) -> float: + sugars = np.fromiter((a.sugar for a in model.agent_list), dtype=float) + vision = np.fromiter((a.vision for a in model.agent_list), dtype=float) + return _safe_corr(sugars, vision) + + +def gini(values: Iterable[float]) -> float: + 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 Sugarscape(mesa.Model): + def __init__( + self, + agents: int, + *, + width: int, + height: int, + max_sugar: int = 4, + 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.width = int(width) + self.height = int(height) + + # Sugar field (current and max) as 2D arrays shaped (width, height) + numpy_rng = np.random.default_rng(seed) + self.sugar_max = numpy_rng.integers( + 0, max_sugar + 1, size=(width, height), dtype=np.int64 + ) + self.sugar_current = self.sugar_max.copy() + + # Grid with capacity 1 per cell + self.grid = SingleGrid(width, height, torus=False) + + # Agents (Python list, manually shuffled/iterated for speed) + self.agent_list: list[AntAgent] = [] + # Place all agents on empty cells; also draw initial traits from model RNG + placed = 0 + while placed < agents: + x = int(self.random.randrange(0, width)) + y = int(self.random.randrange(0, height)) + if self.grid.is_cell_empty((x, y)): + a = AntAgent( + self, + sugar=int(self.random.randint(6, 25)), + metabolism=int(self.random.randint(2, 5)), + vision=int(self.random.randint(1, 6)), + ) + self.grid.place_agent(a, (x, y)) + self.agent_list.append(a) + placed += 1 + + # Model-level reporters mirroring the Frames implementation so CSVs + # are comparable across backends. + self.datacollector = DataCollector( + model_reporters={ + "mean_sugar": lambda m: float(np.mean([a.sugar for a in m.agent_list])) + if m.agent_list + else 0.0, + "total_sugar": lambda m: float(sum(a.sugar for a in m.agent_list)) + if m.agent_list + else 0.0, + "agents_alive": lambda m: float(len(m.agent_list)), + "gini": lambda m: gini(a.sugar for a in m.agent_list), + "corr_sugar_metabolism": lambda m: corr_sugar_metabolism(m), + "corr_sugar_vision": lambda m: corr_sugar_vision(m), + "seed": lambda m: seed, + }, + agent_reporters={ + "traits": lambda a: { + "sugar": a.sugar, + "metabolism": a.metabolism, + "vision": a.vision, + } + }, + ) + self.datacollector.collect(self) + + # --- Scheduling --- + + def _harvest_and_survive(self) -> None: + survivors: list[AntAgent] = [] + for a in self.agent_list: + x, y = a.pos + a.sugar += int(self.sugar_current[x, y]) + a.sugar -= a.metabolism + # Harvested cells are emptied now; they will be refilled if empty. + self.sugar_current[x, y] = 0 + if a.sugar > 0: + survivors.append(a) + else: + # Remove dead agent from grid + self.grid.remove_agent(a) + self.agent_list = survivors + + def _regrow(self) -> None: + # Empty cells regrow to max; occupied cells set to 0 (already zeroed on harvest) + for x in range(self.width): + for y in range(self.height): + if self.grid.is_cell_empty((x, y)): + self.sugar_current[x, y] = self.sugar_max[x, y] + else: + self.sugar_current[x, y] = 0 + + def step(self) -> None: + # Randomise order, move sequentially, then eat/starve, regrow, collect + self.random.shuffle(self.agent_list) + for a in self.agent_list: + a.move() + self._harvest_and_survive() + self._regrow() + self.datacollector.collect(self) + if not self.agent_list: + self.running = False + + def run(self, steps: int) -> None: + for _ in range(steps): + if not getattr(self, "running", True): + break + self.step() + + +def simulate( + *, + agents: int, + steps: int, + width: int, + height: int, + max_sugar: int = 4, + seed: int | None = None, +) -> MesaSimulationResult: + model = Sugarscape( + agents, width=width, height=height, max_sugar=max_sugar, 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.")] = 400, + width: Annotated[int, typer.Option(help="Grid width (columns).")] = 40, + height: Annotated[int, typer.Option(help="Grid height (rows).")] = 40, + steps: Annotated[int, typer.Option(help="Number of model steps to run.")] = 60, + max_sugar: Annotated[int, typer.Option(help="Maximum sugar per cell.")] = 4, + 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: + typer.echo( + f"Running Sugarscape IG (mesa, sequential) with {agents} agents on {width}x{height} 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() + result = simulate( + agents=agents, + steps=steps, + width=width, + height=height, + max_sugar=max_sugar, + seed=seed, + ) + typer.echo(f"Simulation completed in {perf_counter() - start_time:.3f} seconds") + dc = result.datacollector + + # Extract metrics using DataCollector API + model_pd = ( + dc.get_model_vars_dataframe().reset_index().rename(columns={"index": "step"}) + ) + # Keep the full model metrics (step + any model reporters) + seed_val = None + if "seed" in model_pd.columns and not model_pd.empty: + seed_val = model_pd["seed"].iloc[0] + + # Show tail for quick inspection (exclude seed column from display) + display_pd = ( + model_pd.drop(columns=["seed"]) if "seed" in model_pd.columns else model_pd + ) + typer.echo( + f"Metrics in the final 5 steps:\n{display_pd.tail(5).to_string(index=False)}" + ) + + # Save CSV (full model metrics) + if save_results: + csv_path = results_dir / "model.csv" + model_pd.to_csv(csv_path, index=False) + + # Plot per-metric similar to the backend_frames example: create a + # `plots/` subdirectory and generate one figure per model metric column + if plot and not model_pd.empty: + plots_dir = results_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + # Determine which columns to plot (preserve 'step' if present). + # Exclude 'seed' from plots so we don't create a chart for a constant + # model reporter; keep 'seed' in the CSV/dataframe for reproducibility. + value_cols = [c for c in model_pd.columns if c not in {"step", "seed"}] + for col in value_cols: + stem = f"{col}_{timestamp}" + single = ( + model_pd[["step", col]] + if "step" in model_pd.columns + else model_pd[[col]] + ) + # Convert the single-column pandas DataFrame to Polars for the + # shared plotting helper. + single_pl = pl.from_pandas(single) + # Omit seed from subtitle/plot metadata to avoid leaking a constant + # value into the figure (it remains in the saved CSV). If you want + # to include the seed in filenames or external metadata, prefer + # annotating the output folder or README instead. + plot_model_metrics( + single_pl, + plots_dir, + stem, + title=f"Sugarscape IG - {col.capitalize()}", + subtitle="mesa backend", + agents=agents, + steps=steps, + ) + + typer.echo(f"Saved plots under {plots_dir}") + + typer.echo(f"Saved CSV results under {results_dir}") + + +if __name__ == "__main__": + app() 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/mesa_frames/abstract/agentset.py b/mesa_frames/abstract/agentset.py index ae5db2db..ad693bb6 100644 --- a/mesa_frames/abstract/agentset.py +++ b/mesa_frames/abstract/agentset.py @@ -23,6 +23,8 @@ from contextlib import suppress from typing import Any, Literal, Self, overload +from collections.abc import Callable, Sequence + from numpy.random import Generator from mesa_frames.abstract.mixin import CopyMixin, DataFrameMixin @@ -222,7 +224,25 @@ def get( self, attr_names: str | Collection[str] | None = None, mask: AgentMask | None = None, - ) -> Series | DataFrame: ... + ) -> Series | DataFrame: + """Retrieve agent attributes as a Series or DataFrame. + + Parameters + ---------- + attr_names : str | Collection[str] | None, optional + Column name or collection of names to fetch. When ``None``, return + all agent attributes (excluding any internal identifiers). + mask : AgentMask | None, optional + Subset selector limiting which agents are included. ``None`` means + operate on the full set. + + Returns + ------- + Series | DataFrame + A Series when selecting a single attribute, otherwise a DataFrame + containing the requested columns. + """ + ... @abstractmethod def step(self) -> None: @@ -406,6 +426,13 @@ def __reversed__(self) -> Iterator: @property def df(self) -> DataFrame: + """Return the full backing DataFrame for this agent set. + + Returns + ------- + DataFrame + Table containing every agent, including inactive records. + """ return self._df @df.setter @@ -421,18 +448,54 @@ def df(self, agents: DataFrame) -> None: @property @abstractmethod - def active_agents(self) -> DataFrame: ... + def active_agents(self) -> DataFrame: + """Return the subset of agents currently marked as active. + + Returns + ------- + DataFrame + DataFrame view containing only active agents. + """ + ... @property @abstractmethod - def inactive_agents(self) -> DataFrame: ... + def inactive_agents(self) -> DataFrame: + """Return the subset of agents currently marked as inactive. + + Returns + ------- + DataFrame + DataFrame view containing only inactive agents. + """ + ... @property @abstractmethod - def index(self) -> Index: ... + def index(self) -> Index: + """Return the unique identifier index for agents in this set. + + Returns + ------- + Index + Collection of unique agent identifiers. + """ + ... @property def pos(self) -> DataFrame: + """Return positional data for agents from the attached space. + + Returns + ------- + DataFrame + Position records aligned with each agent's ``unique_id``. + + Raises + ------ + AttributeError + If the model has no space attached. + """ if self.space is None: raise AttributeError( "Attempted to access `pos`, but the model has no space attached." @@ -458,14 +521,35 @@ def name(self) -> str: @property def model(self) -> mesa_frames.concrete.model.Model: + """Return the parent model for this agent set. + + Returns + ------- + mesa_frames.concrete.model.Model + The model instance that owns this agent set. + """ return self._model @property def random(self) -> Generator: + """Return the random number generator shared with the model. + + Returns + ------- + Generator + Generator used for stochastic operations. + """ return self.model.random @property def space(self) -> mesa_frames.abstract.space.Space | None: + """Return the space attached to the parent model, if any. + + Returns + ------- + mesa_frames.abstract.space.Space | None + Spatial structure registered on the model, or ``None`` when absent. + """ return self.model.space @abstractmethod @@ -521,6 +605,81 @@ def set( """ ... + @abstractmethod + def select( + self, + mask: AgentMask | None = None, + filter_func: Callable[[Self], BoolSeries] | None = None, + n: int | None = None, + negate: bool = False, + inplace: bool = True, + ) -> Self: + """Update the active-agent mask using selection criteria. + + Parameters + ---------- + mask : AgentMask | None, optional + Pre-computed mask identifying agents to activate. + filter_func : Callable[[Self], BoolSeries] | None, optional + Callable evaluated on the agent set to produce an additional mask. + n : int | None, optional + Randomly sample ``n`` agents from the selected mask when provided. + negate : bool, optional + Invert the effective mask, by default False. + inplace : bool, optional + Whether to mutate in place or return an updated copy, by default True. + + Returns + ------- + Self + The updated AgentSet (or a modified copy when ``inplace=False``). + """ + ... + + @abstractmethod + def shuffle(self, inplace: bool = True) -> Self: + """Randomly permute agent order. + + Parameters + ---------- + inplace : bool, optional + Whether to mutate in place or return a shuffled copy, by default True. + + Returns + ------- + Self + The shuffled AgentSet (or a shuffled copy when ``inplace=False``). + """ + ... + + @abstractmethod + def sort( + self, + by: str | Sequence[str], + ascending: bool | Sequence[bool] = True, + inplace: bool = True, + **kwargs: Any, + ) -> Self: + """Sort agents by one or more columns. + + Parameters + ---------- + by : str | Sequence[str] + Column name(s) to sort on. + ascending : bool | Sequence[bool], optional + Sort order per column, by default True. + inplace : bool, optional + Whether to mutate in place or return a sorted copy, by default True. + **kwargs : Any + Backend-specific keyword arguments forwarded to the concrete sorter. + + Returns + ------- + Self + The sorted AgentSet (or a sorted copy when ``inplace=False``). + """ + ... + def __setitem__( self, key: str diff --git a/mesa_frames/abstract/agentsetregistry.py b/mesa_frames/abstract/agentsetregistry.py index abb0ef69..44a52eb4 100644 --- a/mesa_frames/abstract/agentsetregistry.py +++ b/mesa_frames/abstract/agentsetregistry.py @@ -350,7 +350,6 @@ def remove( """ ... - # select() intentionally removed from the abstract API. @abstractmethod def replace( diff --git a/mesa_frames/abstract/datacollector.py b/mesa_frames/abstract/datacollector.py index 6505408f..dd5cce57 100644 --- a/mesa_frames/abstract/datacollector.py +++ b/mesa_frames/abstract/datacollector.py @@ -121,8 +121,8 @@ def collect(self) -> None: This method calls _collect() to perform actual data collection. - Example - ------- + Examples + -------- >>> datacollector.collect() """ self._collect() @@ -133,8 +133,8 @@ def conditional_collect(self) -> None: This method calls _collect() to perform actual data collection only if trigger returns True - Example - ------- + Examples + -------- >>> datacollector.conditional_collect() """ if self._should_collect(): @@ -166,8 +166,8 @@ def data(self) -> Any: """ Returns collected data currently in memory as a dataframe. - Example: - ------- + Examples + -------- >>> df = datacollector.data >>> print(df) """ @@ -183,8 +183,8 @@ def flush(self) -> None: use this method to save collected data. - Example - ------- + Examples + -------- >>> datacollector.flush() >>> # Data is saved externally and in-memory buffers are cleared if configured """ @@ -219,7 +219,7 @@ def seed(self) -> int: """ Function to get the model seed. - Example: + Examples -------- >>> seed = datacollector.seed """ diff --git a/mesa_frames/concrete/agentset.py b/mesa_frames/concrete/agentset.py index 2a9b1a55..7e1b5e0d 100644 --- a/mesa_frames/concrete/agentset.py +++ b/mesa_frames/concrete/agentset.py @@ -132,19 +132,12 @@ def rename(self, new_name: str, inplace: bool = True) -> Self: # Check if we have a model and can find the AgentSetRegistry that contains this set try: if self in self.model.sets: - # Save index to locate the copy on non-inplace path - try: - idx = list(self.model.sets).index(self) # type: ignore[arg-type] - except Exception: - idx = None reg = self.model.sets.rename(self, new_name, inplace=inplace) if inplace: return self - if idx is not None: - return reg[idx] - return reg.get(new_name) # type: ignore[return-value] - except Exception: - # Fall back to local rename if delegation fails + return reg[new_name] + except KeyError: + # Fall back to local rename if isn't found in a an AgentSetRegistry obj._name = new_name return obj diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py index b91db207..00b1cf25 100644 --- a/mesa_frames/concrete/model.py +++ b/mesa_frames/concrete/model.py @@ -113,6 +113,28 @@ def reset_randomizer(self, seed: int | Sequence[int] | None) -> None: self._seed = seed self.random = np.random.default_rng(seed=self._seed) + @property + def seed(self) -> int | Sequence[int]: + """Return the current seed used by the model's random generator. + + Returns + ------- + int | Sequence[int] + The seed that initialized the underlying RNG. + """ + return self._seed + + @seed.setter + def seed(self, seed: int | Sequence[int] | None) -> None: + """Reset the model random generator using a new seed. + + Parameters + ---------- + seed : int | Sequence[int] | None + A new seed value; falls back to system entropy when ``None``. + """ + self.reset_randomizer(seed) + def run_model(self) -> None: """Run the model until the end condition is reached. diff --git a/mkdocs.yml b/mkdocs.yml index 8a462881..40afb584 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,8 +1,9 @@ # 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 +edit_uri: edit/main/docs/general/ docs_dir: docs/general # Theme configuration @@ -40,12 +41,17 @@ 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: - search - mkdocs-jupyter: - execute: true # Ensures the notebooks run and generate output + execute: false # Ensures the notebooks run and generate output + include: ["*.ipynb"] # Restrict processing to notebooks only (avoid executing raw .py tutorial files) - git-revision-date-localized: enable_creation_date: true - minify: @@ -92,10 +98,9 @@ markdown_extensions: # Extra JavaScript and CSS for rendering extra_javascript: - - javascripts/mathjax.js - - 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) # Customization extra: social: @@ -110,12 +115,12 @@ nav: - User Guide: - Getting Started: user-guide/0_getting-started.md - 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 + - Tutorials: + - Introductory Tutorial: tutorials/2_introductory_tutorial.ipynb + - Advanced Tutorial: tutorials/3_advanced_tutorial.ipynb + - Data Collector Tutorial: tutorials/4_datacollector.ipynb - API Reference: api/index.html - Contributing: - Contribution Guide: contributing.md - - Development Guidelines: development/index.md - Roadmap: roadmap.md + - Changelog: changelog.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..dfd4cc46 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", 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", 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"