From 69e407a06c06a64bbd6ae7592b0db425edbeb490 Mon Sep 17 00:00:00 2001 From: Adam Amer <136176500+adamamer20@users.noreply.github.com> Date: Mon, 8 Dec 2025 23:57:29 +0100 Subject: [PATCH 1/2] Revamp documentation content and navigation --- .gitignore | 10 +- CHANGELOG.md | 40 + CONTRIBUTING.md | 74 +- README.md | 213 ++-- ROADMAP.md | 82 +- docs/api/_static/brand-core.css | 0 docs/api/_static/brand-pydata.css | 0 docs/api/conf.py | 40 +- docs/api/index.rst | 57 +- docs/api/reference/agents/index.rst | 194 ++- docs/api/reference/datacollector.rst | 70 +- docs/api/reference/model.rst | 69 +- docs/api/reference/space/index.rst | 100 +- docs/general/changelog.md | 1 + docs/general/index.md | 90 +- docs/general/plots/boltzmann.svg | 1083 ++++++++++++++++ docs/general/plots/sugarscape.svg | 1091 +++++++++++++++++ .../tutorials/2_introductory_tutorial.py | 277 +++++ docs/general/tutorials/4_datacollector.py | 229 ++++ docs/general/user-guide/0_getting-started.md | 113 +- mkdocs.yml | 23 +- pyproject.toml | 2 +- uv.lock | 49 +- 23 files changed, 3484 insertions(+), 423 deletions(-) create mode 100644 CHANGELOG.md create mode 100644 docs/api/_static/brand-core.css create mode 100644 docs/api/_static/brand-pydata.css create mode 100644 docs/general/changelog.md create mode 100644 docs/general/plots/boltzmann.svg create mode 100644 docs/general/plots/sugarscape.svg create mode 100644 docs/general/tutorials/2_introductory_tutorial.py create mode 100644 docs/general/tutorials/4_datacollector.py diff --git a/.gitignore b/.gitignore index 98823b67..e545a0a4 100644 --- a/.gitignore +++ b/.gitignore @@ -156,8 +156,10 @@ llm_rules.md .python-version benchmarks/results/* -docs/api/_build/* -docs/api/reference/* +docs/api/_build +docs/api/reference/**/mesa_frames.*.rst examples/**/results/* -docs/general/**/data_* -docs/site/* \ No newline at end of file +docs/general/tutorials/data_csv +docs/general/tutorials/data_parquet +docs/general/tutorials/*.ipynb +docs/site \ 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/docs/api/_static/brand-core.css b/docs/api/_static/brand-core.css new file mode 100644 index 00000000..e69de29b diff --git a/docs/api/_static/brand-pydata.css b/docs/api/_static/brand-pydata.css new file mode 100644 index 00000000..e69de29b 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/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/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/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/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 5bca2fcf..7f63a94c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ test = [ docs = [ { include-group = "typechecking" }, + "jupytext>=1.17.3", "mkdocs-material>=9.6.14", "mkdocs-jupyter>=0.25.1", "mkdocs-git-revision-date-localized-plugin>=1.4.7", @@ -77,7 +78,6 @@ 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", diff --git a/uv.lock b/uv.lock index ee2f031c..4db6f107 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" @@ -1235,6 +1222,7 @@ dependencies = [ dev = [ { name = "autodocsumm" }, { name = "beartype" }, + { name = "jupytext" }, { name = "mesa" }, { name = "mkdocs-git-revision-date-localized-plugin" }, { name = "mkdocs-include-markdown-plugin" }, @@ -1243,7 +1231,6 @@ dev = [ { name = "mkdocs-minify-plugin" }, { name = "numba" }, { name = "numpydoc" }, - { name = "perfplot" }, { name = "pre-commit" }, { name = "pydata-sphinx-theme" }, { name = "pytest" }, @@ -1259,6 +1246,7 @@ dev = [ docs = [ { name = "autodocsumm" }, { name = "beartype" }, + { name = "jupytext" }, { name = "mesa" }, { name = "mkdocs-git-revision-date-localized-plugin" }, { name = "mkdocs-include-markdown-plugin" }, @@ -1266,7 +1254,6 @@ docs = [ { name = "mkdocs-material" }, { name = "mkdocs-minify-plugin" }, { name = "numpydoc" }, - { name = "perfplot" }, { name = "pydata-sphinx-theme" }, { name = "seaborn" }, { name = "sphinx" }, @@ -1298,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" }, @@ -1306,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" }, @@ -1322,6 +1309,7 @@ dev = [ 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" }, @@ -1329,7 +1317,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" }, @@ -1730,21 +1717,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" @@ -2275,19 +2247,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, ] -[[package]] -name = "rich" -version = "14.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "markdown-it-py" }, - { name = "pygments" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fe/75/af448d8e52bf1d8fa6a9d089ca6c07ff4453d86c65c145d0a300bb073b9b/rich-14.1.0.tar.gz", hash = "sha256:e497a48b844b0320d45007cdebfeaeed8db2a4f4bcf49f15e455cfc4af11eaa8", size = 224441, upload-time = "2025-07-25T07:32:58.125Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e3/30/3c4d035596d3cf444529e0b2953ad0466f6049528a879d27534700580395/rich-14.1.0-py3-none-any.whl", hash = "sha256:536f5f1785986d6dbdea3c75205c473f970777b4a0d6c6dd1b696aa05a3fa04f", size = 243368, upload-time = "2025-07-25T07:32:56.73Z" }, -] - [[package]] name = "roman-numerals-py" version = "3.1.0" From ae64f6eade417cb4319f5bee5b5a5d38d368c0ee Mon Sep 17 00:00:00 2001 From: Adam Amer <136176500+adamamer20@users.noreply.github.com> Date: Tue, 16 Dec 2025 20:25:55 +0100 Subject: [PATCH 2/2] Merge branch 'main' of https://github.com/projectmesa/mesa-frames into split/docs-and-toptexts --- .github/ISSUE_TEMPLATE/config.yml | 2 +- .github/workflows/build.yml | 105 +- .github/workflows/docs-gh-pages.yml | 85 +- .github/workflows/docs-preview-dispatch.yml | 97 + .github/workflows/publish.yml | 75 +- .github/workflows/sync_coderabbit.yml | 6 +- .gitignore | 2 +- CONTRIBUTING.md | 4 +- README.md | 14 +- benchmarks/README.md | 84 + benchmarks/cli.py | 285 +++ docs/api/conf.py | 4 +- docs/general/tutorials/3_advanced_tutorial.py | 1749 +++++++++++++++++ .../user-guide/2_introductory-tutorial.ipynb | 4 +- .../general/user-guide/3_advanced-tutorial.md | 4 - docs/general/user-guide/4_datacollector.ipynb | 4 +- docs/general/user-guide/5_benchmarks.md | 10 +- examples/README.md | 106 + examples/__init__.py | 6 + examples/boltzmann_wealth/README.md | 96 + examples/boltzmann_wealth/__init__.py | 0 examples/boltzmann_wealth/backend_frames.py | 188 ++ examples/boltzmann_wealth/backend_mesa.py | 178 ++ examples/plotting.py | 306 +++ examples/sugarscape_ig/README.md | 105 + .../sugarscape_ig/backend_frames/__init__.py | 1 + .../sugarscape_ig/backend_frames/agents.py | 645 ++++++ .../sugarscape_ig/backend_frames/model.py | 497 +++++ .../sugarscape_ig/backend_mesa/__init__.py | 1 + examples/sugarscape_ig/backend_mesa/agents.py | 81 + examples/sugarscape_ig/backend_mesa/model.py | 312 +++ examples/utils.py | 25 + mesa_frames/__init__.py | 2 +- mesa_frames/abstract/space.py | 2 +- mkdocs.yml | 14 +- pyproject.toml | 9 +- tests/test_benchmarks_cli.py | 112 ++ tests/test_sugarscape_mesa_cli.py | 26 + uv.lock | 102 +- 39 files changed, 5161 insertions(+), 187 deletions(-) create mode 100644 .github/workflows/docs-preview-dispatch.yml create mode 100644 benchmarks/README.md create mode 100644 benchmarks/cli.py create mode 100644 docs/general/tutorials/3_advanced_tutorial.py delete mode 100644 docs/general/user-guide/3_advanced-tutorial.md create mode 100644 examples/README.md create mode 100644 examples/__init__.py create mode 100644 examples/boltzmann_wealth/README.md create mode 100644 examples/boltzmann_wealth/__init__.py create mode 100644 examples/boltzmann_wealth/backend_frames.py create mode 100644 examples/boltzmann_wealth/backend_mesa.py create mode 100644 examples/plotting.py create mode 100644 examples/sugarscape_ig/README.md create mode 100644 examples/sugarscape_ig/backend_frames/__init__.py create mode 100644 examples/sugarscape_ig/backend_frames/agents.py create mode 100644 examples/sugarscape_ig/backend_frames/model.py create mode 100644 examples/sugarscape_ig/backend_mesa/__init__.py create mode 100644 examples/sugarscape_ig/backend_mesa/agents.py create mode 100644 examples/sugarscape_ig/backend_mesa/model.py create mode 100644 examples/utils.py create mode 100644 tests/test_benchmarks_cli.py create mode 100644 tests/test_sugarscape_mesa_cli.py diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index c26704fd..aae48c78 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,7 +1,7 @@ blank_issues_enabled: true contact_links: - name: Discussions/Questions etc - url: https://github.com/projectmesa/mesa-frames/discussions + url: https://github.com/mesa/mesa-frames/discussions about: Discuss Mesa-Frames, ask questions, do discussions, share ideas, and showcase your projects diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 193f450c..a28b8290 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,24 +19,31 @@ concurrency: cancel-in-progress: true jobs: - test_on_ubuntu: - name: build (ubuntu, py${{ matrix.python-version }}) - runs-on: ubuntu-latest - timeout-minutes: 6 - + build: + runs-on: ${{ matrix.os }}-latest + timeout-minutes: 10 strategy: fail-fast: false matrix: - python-version: ["3.11", "3.12", "3.13"] + os: [windows, ubuntu, macos] + python-version: ["3.13"] + pip-pre: [""] + include: + - os: ubuntu + python-version: "3.12" + - os: ubuntu + python-version: "3.11" + - os: ubuntu + python-version: "3.13" + pip-pre: "--pre" env: MESA_FRAMES_RUNTIME_TYPECHECKING: "true" POSTGRES_URI: postgresql://user:password@localhost:5432/testdb - SKIP_PG_TESTS: "false" services: postgres: - image: postgres:15 + image: ${{ matrix.os == 'ubuntu' && 'postgres:15' || '' }} ports: - 5432:5432 env: @@ -50,94 +57,48 @@ jobs: --health-retries=5 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} + allow-prereleases: true + cache: "pip" - name: Install uv via GitHub Action - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v7 with: - cache: true + enable-cache: true - - name: Install mesa-frames + dev dependencies + - name: Install mesa-frames and dependencies run: | - uv pip install --system . - uv pip install --group dev --system - - - name: Test with pytest - run: pytest --durations=10 --cov=mesa_frames tests/ --cov-report=xml - - - name: Upload coverage to Codecov - uses: codecov/codecov-action@v4 - with: - token: ${{ secrets.CODECOV_TOKEN }} - - test_on_mac_windows: - name: build (${{ matrix.os }}, py3.13) - runs-on: ${{ matrix.os }}-latest - timeout-minutes: 10 - - strategy: - matrix: - include: - - os: windows - skip_pg: true - - os: macos - skip_pg: false - - env: - MESA_FRAMES_RUNTIME_TYPECHECKING: "true" - POSTGRES_URI: postgresql://user:password@localhost:5432/testdb - SKIP_PG_TESTS: ${{ matrix.skip_pg }} - - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.13" - - - name: Install uv via GitHub Action - uses: astral-sh/setup-uv@v6 - with: - cache: true - - - name: Install mesa-frames + dev dependencies - run: | - uv pip install --system . + uv pip install --system . ${{ matrix.pip-pre }} uv pip install --group dev --system - name: Install and Start PostgreSQL (macOS) if: matrix.os == 'macos' run: | brew install postgresql@15 - export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH" export PGDATA="/opt/homebrew/var/postgresql@15" - - # Ensure a clean database directory rm -rf $PGDATA mkdir -p $PGDATA - initdb --username=user --auth=trust --encoding=UTF8 $PGDATA - pg_ctl -D $PGDATA -l logfile start - sleep 5 - createdb testdb -U user - shell: bash + + - name: Set SKIP_PG_TESTS for Windows + if: matrix.os == 'windows' + run: echo "SKIP_PG_TESTS=true" >> $env:GITHUB_ENV - name: Test with pytest run: pytest --durations=10 --cov=mesa_frames tests/ --cov-report=xml - build: - name: build - runs-on: ubuntu-latest - needs: [test_on_ubuntu, test_on_mac_windows] - steps: - - run: echo "All matrix jobs completed successfully." + - if: matrix.os == 'ubuntu' + name: Codecov + uses: codecov/codecov-action@v5 + with: + fail_ci_if_error: true + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/.github/workflows/docs-gh-pages.yml b/.github/workflows/docs-gh-pages.yml index 435af957..ab7157e4 100644 --- a/.github/workflows/docs-gh-pages.yml +++ b/.github/workflows/docs-gh-pages.yml @@ -1,40 +1,89 @@ -name: Build and Deploy Documentation +name: Docs โ€” Build & Deploy to Main +permissions: + contents: read 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 + permissions: + contents: read + actions: write steps: - uses: actions/checkout@v4 - with: - fetch-depth: 0 # Fetch all history for .git-restore-mtime to work correctly + with: { fetch-depth: 0 } - - name: Set up Python - uses: actions/setup-python@v5 + - name: Set up Python 3.13 + uses: actions/setup-python@v6 with: - python-version: '3.x' + python-version: "3.13" + allow-prereleases: true + cache: "pip" - - name: Install uv via GitHub Action - uses: astral-sh/setup-uv@v6 + - uses: astral-sh/setup-uv@v7 - - 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 documentation) - run: sphinx-build -b html docs/api site/api + - name: Build Sphinx docs (API) + run: uv run sphinx-build -b html docs/api site/api - - name: Deploy to GitHub Pages + - name: Upload site artifact + uses: actions/upload-artifact@v4 + with: + name: site + path: site + + deploy-main: + needs: build + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + permissions: + contents: write + pages: write + steps: + - uses: actions/download-artifact@v4 + with: { name: site, path: site } + - name: Deploy to GitHub Pages (main) uses: peaceiris/actions-gh-pages@v4 with: github_token: ${{ secrets.GITHUB_TOKEN }} + publish_branch: gh-pages publish_dir: ./site - force_orphan: true \ No newline at end of file + force_orphan: true diff --git a/.github/workflows/docs-preview-dispatch.yml b/.github/workflows/docs-preview-dispatch.yml new file mode 100644 index 00000000..24ca12e2 --- /dev/null +++ b/.github/workflows/docs-preview-dispatch.yml @@ -0,0 +1,97 @@ +# Deploys preview artifacts produced by the Docs โ€” Build & Preview workflow for PRs. +# The build workflow runs with read-only permissions on PRs (including forks) and +# uploads the built site as an artifact. This workflow runs in the base repo with +# the permissions needed to publish the preview to the gh-pages branch, but it +# never checks out or executes untrusted PR code. +name: Docs โ€” Preview Deploy + +on: + workflow_run: + workflows: ["Docs โ€” Build & Deploy to Main"] + types: [completed] + +permissions: + contents: write + actions: read + pages: write + issues: write + pull-requests: write + +jobs: + deploy-preview: + if: > + github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.conclusion == 'success' && + (github.event.workflow_run.pull_requests || null) + runs-on: ubuntu-latest + steps: + - name: Download built site artifact + uses: actions/download-artifact@v4 + with: + name: site + path: site + run-id: ${{ github.event.workflow_run.id }} + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Derive preview metadata + id: meta + env: + HEAD_SHA: ${{ github.event.workflow_run.head_sha }} + HEAD_BRANCH: ${{ github.event.workflow_run.head_branch }} + PR_NUMBER: ${{ github.event.workflow_run.pull_requests[0].number }} + run: | + set -euo pipefail + short_sha="$(printf '%s' "$HEAD_SHA" | cut -c1-7)" + { + echo "short_sha=$short_sha" + echo "head_branch=$HEAD_BRANCH" + echo "pr_number=$PR_NUMBER" + } >> "$GITHUB_OUTPUT" + + - name: Deploy preview under subfolder + uses: peaceiris/actions-gh-pages@v4 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_branch: gh-pages + publish_dir: ./site + destination_dir: preview/${{ steps.meta.outputs.head_branch }}/${{ steps.meta.outputs.short_sha }} + keep_files: true + + - name: Print preview URL + env: + PREVIEW_OWNER: ${{ github.repository_owner }} + PREVIEW_REPO: ${{ github.repository }} + HEAD_BRANCH: ${{ steps.meta.outputs.head_branch }} + SHORT_SHA: ${{ steps.meta.outputs.short_sha }} + PR_NUMBER: ${{ steps.meta.outputs.pr_number }} + run: | + echo "Preview for PR #${PR_NUMBER}:" + echo "https://${PREVIEW_OWNER}.github.io/$(basename "$PREVIEW_REPO")/preview/${HEAD_BRANCH}/${SHORT_SHA}/" + + - name: Comment preview URL on PR + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const prNumber = Number('${{ steps.meta.outputs.pr_number }}'); + if (!Number.isInteger(prNumber)) { + throw new Error(`Invalid PR number: ${{ steps.meta.outputs.pr_number }}`); + } + + const owner = context.repo.owner; + const repo = context.repo.repo; + const headBranch = '${{ steps.meta.outputs.head_branch }}'; + const shortSha = '${{ steps.meta.outputs.short_sha }}'; + const previewUrl = `https://${owner}.github.io/${repo}/preview/${headBranch}/${shortSha}/`; + const body = [ + `Docs preview ready for PR #${prNumber}.`, + '', + previewUrl, + ].join('\n'); + + await github.rest.issues.createComment({ + owner, + repo, + issue_number: prNumber, + body, + }); diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 105824de..a2070903 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,52 @@ 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 + # Ensure we are on the main branch before committing so the push lands on main. + git checkout main + 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 +109,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/.github/workflows/sync_coderabbit.yml b/.github/workflows/sync_coderabbit.yml index 49213652..741e64ed 100644 --- a/.github/workflows/sync_coderabbit.yml +++ b/.github/workflows/sync_coderabbit.yml @@ -19,7 +19,7 @@ jobs: pull-requests: write # open PR & enable autoโ€‘merge env: - UPSTREAM_RAW_URL: "https://raw.githubusercontent.com/projectmesa/mesa/main/.coderabbit.yaml" + UPSTREAM_RAW_URL: "https://raw.githubusercontent.com/mesa/mesa/main/.coderabbit.yaml" SYNC_BRANCH: "sync-coderabbit-from-mesa" steps: @@ -51,11 +51,11 @@ jobs: with: token: ${{ secrets.GITHUB_TOKEN }} branch: ${{ env.SYNC_BRANCH }} - commit-message: "chore: sync .coderabbit.yaml from projectmesa/mesa@main" + commit-message: "chore: sync .coderabbit.yaml from mesa/mesa@main" title: "๐Ÿ”„ Sync .coderabbit.yaml from mesa:main" body: | Automated update of **.coderabbit.yaml** from \ - [projectmesa/mesa:main](${{ env.UPSTREAM_RAW_URL }}). + [mesa/mesa:main](${{ env.UPSTREAM_RAW_URL }}). This PR will merge automatically once all required checks succeed. labels: automated, coderabbit-sync diff --git a/.gitignore b/.gitignore index e545a0a4..9f4686a4 100644 --- a/.gitignore +++ b/.gitignore @@ -162,4 +162,4 @@ examples/**/results/* docs/general/tutorials/data_csv docs/general/tutorials/data_parquet docs/general/tutorials/*.ipynb -docs/site \ No newline at end of file +docs/site diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 82a340c9..d865ab46 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,11 +2,11 @@ Thank you for taking the time to contribute to **mesa-frames**! Since the project is still in its early stages, we warmly welcome contributions that will help shape its development. ๐ŸŽ‰ -For a more general and comprehensive guide, please refer to [mesa's main contribution guidelines](https://github.com/projectmesa/mesa/blob/main/CONTRIBUTING.md). ๐Ÿ“œ +For a more general and comprehensive guide, please refer to [mesa's main contribution guidelines](https://github.com/mesa/mesa/blob/main/CONTRIBUTING.md). ๐Ÿ“œ ## Project Roadmap ๐Ÿ—บ๏ธ -Before contributing, we recommend reviewing our [roadmap](https://projectmesa.github.io/mesa-frames/roadmap/) file to understand the project's current priorities, upcoming features, and long-term vision. This will help ensure your contributions align with the project's direction. +Before contributing, we recommend reviewing our [roadmap](https://mesa.github.io/mesa-frames/roadmap/) file to understand the project's current priorities, upcoming features, and long-term vision. This will help ensure your contributions align with the project's direction. ## How to Contribute ๐Ÿ’ก diff --git a/README.md b/README.md index b46acc8e..2f380037 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@

- Mesa logo + Mesa logo

mesa-frames

@@ -8,7 +8,7 @@ | | | | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| 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) | +| CI/CD | [![CI Checks](https://github.com/mesa/mesa-frames/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/mesa/mesa-frames/actions/workflows/build.yml) [![codecov](https://codecov.io/gh/mesa/mesa-frames/branch/main/graph/badge.svg)](https://app.codecov.io/gh/mesa/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) | @@ -17,7 +17,7 @@ ## Scale Mesa beyond its limits -Classic [Mesa](https://github.com/projectmesa/mesa) stores each agent as a Python object, which quickly becomes a bottleneck at scale. +Classic [Mesa](https://github.com/mesa/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. You keep the Mesa-style `Model` / `AgentSet` structure, but updates are vectorized and memory-efficient. @@ -56,7 +56,7 @@ mesa-frames currently uses **Polars** as its backend. ## Benchmarks -[![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) +[![Reproduce Benchmarks](https://img.shields.io/badge/Reproduce%20Benchmarks-๐Ÿ“Š-orange?style=for-the-badge)](https://github.com/mesa/mesa-frames/blob/main/benchmarks/README.md) **mesa-frames consistently outperforms classic Mesa across both toy and canonical ABMs.** @@ -85,7 +85,7 @@ We still have room to optimize performance further (see [Roadmap](#roadmap)). Or for development: ```bash -git clone https://github.com/projectmesa/mesa-frames.git +git clone https://github.com/mesa/mesa-frames.git cd mesa-frames uv sync --all-extras ``` @@ -135,6 +135,6 @@ uv sync --all-extras ## License -Copyright ยฉ 2025 Adam Amer, Project Mesa team and contributors +Copyright ยฉ 2025 Adam Amer, Mesa team and contributors -Licensed under the [Apache License, Version 2.0](https://raw.githubusercontent.com/projectmesa/mesa-frames/refs/heads/main/LICENSE). +Licensed under the [Apache License, Version 2.0](https://raw.githubusercontent.com/mesa/mesa-frames/refs/heads/main/LICENSE). diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..42da5535 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,84 @@ +# 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. diff --git a/benchmarks/cli.py b/benchmarks/cli.py new file mode 100644 index 00000000..f97e99c4 --- /dev/null +++ b/benchmarks/cli.py @@ -0,0 +1,285 @@ +"""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") + # We keep start = 0 to benchmark initialization time + 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 | list[str], + typer.Option( + help="Models to benchmark: boltzmann, sugarscape, or all", + callback=_parse_models, + ), + ] = "all", + agents: Annotated[ + str | list[int], + typer.Option( + help="Agent count or range (start:stop:step)", callback=_parse_agents + ), + ] = "1000:5000:1000", + steps: Annotated[ + int, + typer.Option( + min=0, + help="Number of steps per run.", + ), + ] = 100, + repeats: Annotated[int, typer.Option(help="Repeats per configuration.", min=1)] = 1, + seed: Annotated[int, typer.Option(help="Optional RNG seed.")] = 42, + save: Annotated[bool, typer.Option(help="Persist benchmark CSV results.")] = True, + plot: Annotated[bool, typer.Option(help="Render performance plots.")] = True, + results_dir: Annotated[ + Path | None, + 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. Defaults to the module's results directory." + ), + ), + ] = None, +) -> None: + """Run performance benchmarks for the selected models.""" + # Support both CLI (via callbacks) and direct function calls + if isinstance(models, str): + models = _parse_models(models) + if isinstance(agents, str): + agents = _parse_agents(agents) + # Ensure module-relative default is computed at call time (avoids import-time side effects) + if results_dir is None: + results_dir = Path(__file__).resolve().parent / "results" + + 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}") + + destinations: list[str] = [] + if save: + destinations.append(f"CSVs under {timestamp_dir}") + if plot: + destinations.append(f"plots under {plots_subdir}") + + if destinations: + typer.echo("Unified benchmark outputs written: " + "; ".join(destinations)) + else: + typer.echo( + "Benchmark run completed (save=False, plot=False; no files written)." + ) + + +if __name__ == "__main__": + app() diff --git a/docs/api/conf.py b/docs/api/conf.py index 4998fbc2..10326190 100644 --- a/docs/api/conf.py +++ b/docs/api/conf.py @@ -85,8 +85,8 @@ # -- GitHub link and user guide settings ------------------------------------- -github_root = "https://github.com/projectmesa/mesa-frames" -web_root = "https://projectmesa.github.io/mesa-frames" +github_root = "https://github.com/mesa/mesa-frames" +web_root = "https://mesa.github.io/mesa-frames" html_theme_options = { "external_links": [ diff --git a/docs/general/tutorials/3_advanced_tutorial.py b/docs/general/tutorials/3_advanced_tutorial.py new file mode 100644 index 00000000..caa5c0a0 --- /dev/null +++ b/docs/general/tutorials/3_advanced_tutorial.py @@ -0,0 +1,1749 @@ +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). In the instant-growback variant used here, any cell + that was occupied at the end of the previous step has sugar 0 (it was harvested + and did not regrow). +* **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):** ants rank candidate destinations using the + start-of-step sugar field; conflicts are resolved by a random lottery in + rounds (losers fall back to their next choice). If an ant wins a destination + other than its origin, its origin becomes available to other ants in later + rounds of the same step. + +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 compare these update rules and show how far a synchronous, +dataframe-friendly rule can go as a **performance-oriented relaxation** of the +classic sequential schedule. Some macroscopic summaries (like total sugar or the +Gini coefficient of wealth) often remain qualitatively similar, while more +fine-grained statistics (like wealthโ€“trait correlations) can drift noticeably for +individual seeds because conflict resolution changes which traits win contested +cells. +""" + +# %% [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 + + +# Simple display helper to render HTML in notebooks while keeping stdout +# output for scripts/CI runs. +def _in_ipython() -> bool: + """Return True when running inside an IPython/Jupyter session.""" + try: + from IPython import get_ipython + except Exception: + return False + return get_ipython() is not None + + +def show_output( + obj: object, + *, + title: str | None = None, + max_rows: int = 12, + collapsible: bool = False, + open_by_default: bool = False, +) -> None: + """Display rich HTML when available, otherwise fall back to prints.""" + rich_env = _in_ipython() + + if isinstance(obj, pl.DataFrame): + df = obj.head(max_rows) if max_rows else obj + if rich_env: + from IPython.display import HTML, display + + if collapsible: + open_attr = " open" if open_by_default else "" + summary = title or "Table" + html = df.to_pandas().to_html(index=False) + display( + HTML( + f""" + + {summary} +
{html}
+ +""" + ) + ) + else: + if title: + display(HTML(f"

{title}

")) + display(df.to_pandas()) + else: + if title: + print(f"\n=== {title} ===") + print(df) + return + + if title: + if rich_env: + from IPython.display import HTML, display + + display(HTML(f"

{title}

")) + else: + print(f"\n=== {title} ===") + + if rich_env: + from IPython.display import display + + display(obj) + else: + print(obj) + + +# %% [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 + 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() + + 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. Occupied cells have already been harvested in + :meth:`AntsBase.eat`, so we only need to refresh empty cells here. + 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"]}) + + +# %% [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.4 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", + ] + ) + .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 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + # Treat all currently occupied cells (origins) as taken from the start. + # Each agent may still target its own origin; we handle that exception + # when filtering candidate pools. + taken = origins.select( + [ + pl.col("dim_0").alias("dim_0_candidate"), + pl.col("dim_1").alias("dim_1_candidate"), + ] + ) + origins_for_filter = origins.rename( + {"dim_0": "dim_0_origin", "dim_1": "dim_1_origin"} + ) + + # 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") + ) + candidate_pool = ( + candidate_pool.join(origins_for_filter, on="agent_id", how="left") + .join( + taken.with_columns(pl.lit(True).alias("is_taken")), + on=["dim_0_candidate", "dim_1_candidate"], + how="left", + ) + .filter( + pl.col("is_taken").is_null() + | ( + (pl.col("dim_0_candidate") == pl.col("dim_0_origin")) + & (pl.col("dim_1_candidate") == pl.col("dim_1_origin")) + ) + ) + .drop(["dim_0_origin", "dim_1_origin", "is_taken"]) + ) + + 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", + ) + 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", + ) + # Origins of agents that move away become available to others in + # subsequent rounds. Keep origins for agents that stayed put. + vacated = ( + winners.join(origins_for_filter, on="agent_id", how="left") + .filter( + (pl.col("dim_0_candidate") != pl.col("dim_0_origin")) + | (pl.col("dim_1_candidate") != pl.col("dim_1_origin")) + ) + .select( + pl.col("dim_0_origin").alias("dim_0_candidate"), + pl.col("dim_1_origin").alias("dim_1_candidate"), + ) + ) + if not vacated.is_empty(): + taken = taken.join( + vacated, + on=["dim_0_candidate", "dim_1_candidate"], + how="anti", + ) + + 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 + + show_output( + frames[variant_name] + .select(["step", "mean_sugar", "total_sugar", "agents_alive"]) + .tail(5), + title=f"{variant_name} aggregate trajectory (last 5 steps)", + max_rows=5, + collapsible=True, + ) + show_output(f"{variant_name} runtime: {runtime:.3f} s") + if not _in_ipython(): + 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) +) + +show_output( + runtime_table, + title="Runtime comparison (fastest first)", + collapsible=True, + open_by_default=True, +) + +# 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, with vacated origins available to others), the final-step Gini differs by โ‰ˆ0.045, and wealth-trait correlations diverge by a few 1e-2 to 1e-1. +These gaps vary by seed and grid size. In practice the parallel rule is best seen as a *performance-oriented relaxation* of the sequential schedule: it can preserve broad macro trends, but it is not a drop-in replacement when you care about seed-level microstructure.""" + +# %% +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"), +) +show_output( + comparison.select(["step", "mean_diff", "total_diff", "count_diff"]).head(10), + title="Step-level absolute differences (first 10 steps)", + max_rows=10, + collapsible=True, +) + + +# 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() +) + +show_output( + 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"), + ] + ), + title="Steady-state inequality metrics", + collapsible=True, + open_by_default=True, +) + +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] + show_output( + f"Absolute Gini gap (numba vs parallel): {abs(numba_gini - par_gini):.4f}" + ) + +# %% [markdown] +""" +## 6. Takeaways and Next Steps + +Some final notes: +- mesa-frames should preferably be used when you have many agents and operations can be vectorized. +- If your model is not easily vectorizable, consider using Numba or reducing your microscopic rule to a vectorizable form. As we saw, the macroscopic behavior can remain consistent (and be more similar to real-world systems). + + +Currently, the Polars implementation spends most of the time in join operations. + +**Polars + LazyFrames roadmap** - future mesa-frames releases will expose + LazyFrame-powered sets and spaces (which can also use a GPU cuda accelerated backend which greatly accelerates joins), so the same Polars + code you wrote here will scale even further without touching Numba. +""" diff --git a/docs/general/user-guide/2_introductory-tutorial.ipynb b/docs/general/user-guide/2_introductory-tutorial.ipynb index 11391f9d..52233265 100644 --- a/docs/general/user-guide/2_introductory-tutorial.ipynb +++ b/docs/general/user-guide/2_introductory-tutorial.ipynb @@ -5,7 +5,7 @@ "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)" + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mesa/mesa-frames/blob/main/docs/general/user-guide/2_introductory-tutorial.ipynb)" ] }, { @@ -25,7 +25,7 @@ "metadata": {}, "outputs": [], "source": [ - "# !pip install git+https://github.com/projectmesa/mesa-frames mesa" + "# !pip install git+https://github.com/mesa/mesa-frames mesa" ] }, { 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 index 0809caa2..4319050f 100644 --- a/docs/general/user-guide/4_datacollector.ipynb +++ b/docs/general/user-guide/4_datacollector.ipynb @@ -7,7 +7,7 @@ "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", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mesa/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", @@ -33,7 +33,7 @@ }, "outputs": [], "source": [ - "# !pip install git+https://github.com/projectmesa/mesa-frames mesa" + "# !pip install git+https://github.com/mesa/mesa-frames mesa" ] }, { diff --git a/docs/general/user-guide/5_benchmarks.md b/docs/general/user-guide/5_benchmarks.md index 61fca87b..416b1bb8 100644 --- a/docs/general/user-guide/5_benchmarks.md +++ b/docs/general/user-guide/5_benchmarks.md @@ -4,18 +4,18 @@ mesa-frames offers significant performance improvements over the original mesa f ## Boltzmann Wealth Model ๐Ÿ’ฐ -[View the benchmark script](https://github.com/projectmesa/mesa-frames/blob/main/examples/boltzmann_wealth/performance_plot.py) +[View the benchmark script](https://github.com/mesa/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) +![Performance Graph BW](https://github.com/mesa/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) +![Performance Graph BW without Mesa](https://github.com/mesa/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) +[View the benchmark script](https://github.com/mesa/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) +![Performance Graph SS IG](https://github.com/mesa/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..51b3e260 --- /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, check out `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/boltzmann_wealth/__init__.py b/examples/boltzmann_wealth/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/boltzmann_wealth/backend_frames.py b/examples/boltzmann_wealth/backend_frames.py new file mode 100644 index 00000000..a4403aa8 --- /dev/null +++ b/examples/boltzmann_wealth/backend_frames.py @@ -0,0 +1,188 @@ +"""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() + typer.echo(f"Saved CSV results under {results_dir}") + + 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}") + + +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..dda875ac --- /dev/null +++ b/examples/boltzmann_wealth/backend_mesa.py @@ -0,0 +1,178 @@ +"""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 mesa +from mesa.datacollection import DataCollector +import numpy as np +import polars as pl +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"}) + extracted_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={extracted_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/plotting.py b/examples/plotting.py new file mode 100644 index 00000000..b2a0473d --- /dev/null +++ b/examples/plotting.py @@ -0,0 +1,306 @@ +# 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, + "legend.facecolor": "#ffffff", + "legend.edgecolor": "#d0d0d0", + }, + ), + "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 leg is not None: + # Remove redundant legend title and ensure a readable + # boxed background for the light theme (subtle) while + # keeping a slightly transparent frame for dark theme. + leg.set_title(None) + frame = leg.get_frame() + if theme == "dark": + frame.set_alpha(0.8) + else: + frame.set_alpha(0.9) + frame.set_edgecolor("#d0d0d0") + frame.set_linewidth(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") + + leg = ax.get_legend() + if leg is not None: + leg.set_title(None) + frame = leg.get_frame() + if theme == "dark": + frame.set_alpha(0.8) + else: + frame.set_alpha(0.9) + frame.set_edgecolor("#d0d0d0") + frame.set_linewidth(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_agent_metrics", + "plot_model_metrics", + "plot_performance", +] diff --git a/examples/sugarscape_ig/README.md b/examples/sugarscape_ig/README.md new file mode 100644 index 00000000..d3734b21 --- /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 agents aid resource gathering. Negative `corr_sugar_metabolism` can emerge if high-metabolism agents accelerate 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, check out `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..3a9c4932 --- /dev/null +++ b/examples/sugarscape_ig/backend_frames/agents.py @@ -0,0 +1,645 @@ +"""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/tutorials/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", + ] + ) + .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 โ”‚ + # โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก + # Treat all currently occupied cells (origins) as taken from the start. + # Each agent may still target its own origin; we handle that exception + # when filtering candidate pools. + taken = origins.select( + [ + pl.col("dim_0").alias("dim_0_candidate"), + pl.col("dim_1").alias("dim_1_candidate"), + ] + ) + origins_for_filter = origins.rename( + {"dim_0": "dim_0_origin", "dim_1": "dim_1_origin"} + ) + + # 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") + ) + candidate_pool = ( + candidate_pool.join(origins_for_filter, on="agent_id", how="left") + .join( + taken.with_columns(pl.lit(True).alias("is_taken")), + on=["dim_0_candidate", "dim_1_candidate"], + how="left", + ) + .filter( + pl.col("is_taken").is_null() + | ( + (pl.col("dim_0_candidate") == pl.col("dim_0_origin")) + & (pl.col("dim_1_candidate") == pl.col("dim_1_origin")) + ) + ) + .drop(["dim_0_origin", "dim_1_origin", "is_taken"]) + ) + + 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", + ) + 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", "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", + ) + + # Origins of agents that move away become available to others in + # subsequent rounds. Keep origins for agents that stayed put. + vacated = ( + winners.join(origins_for_filter, on="agent_id", how="left") + .filter( + (pl.col("dim_0_candidate") != pl.col("dim_0_origin")) + | (pl.col("dim_1_candidate") != pl.col("dim_1_origin")) + ) + .select( + pl.col("dim_0_origin").alias("dim_0_candidate"), + pl.col("dim_1_origin").alias("dim_1_candidate"), + ) + ) + if not vacated.is_empty(): + taken = taken.join( + vacated, + on=["dim_0_candidate", "dim_1_candidate"], + how="anti", + ) + + 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..c51122fe --- /dev/null +++ b/examples/sugarscape_ig/backend_frames/model.py @@ -0,0 +1,497 @@ +"""Mesa-frames implementation of Sugarscape IG with Typer CLI. + +This mirrors the advanced tutorial in docs/general/tutorials/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) + sorted_vals = np.sort(sugar) + 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() + typer.echo(f"Saved CSV results under {results_dir}") + + 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}") + + # Skip CSV-saved confirmation when results are not persisted + + +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..b8626d20 --- /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: mesa.Model, + *, + 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..5b7bbf90 --- /dev/null +++ b/examples/sugarscape_ig/backend_mesa/model.py @@ -0,0 +1,312 @@ +"""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 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) + + # 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) + typer.echo(f"Saved CSV results under {results_dir}") + + # 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}") + + # Skip CSV-saved confirmation when results are not persisted + + +if __name__ == "__main__": + app() 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/__init__.py b/mesa_frames/__init__.py index 20fcbeef..946974a7 100644 --- a/mesa_frames/__init__.py +++ b/mesa_frames/__init__.py @@ -38,7 +38,7 @@ def __init__(self, width, height): Developed by: Project Mesa, Adam Amer License: MIT -GitHub: https://github.com/projectmesa/mesa-frames +GitHub: https://github.com/mesa/mesa-frames """ from __future__ import annotations diff --git a/mesa_frames/abstract/space.py b/mesa_frames/abstract/space.py index 39abe6bd..ca1e8161 100644 --- a/mesa_frames/abstract/space.py +++ b/mesa_frames/abstract/space.py @@ -1236,7 +1236,7 @@ class AbstractGrid(AbstractDiscreteSpace): The [0, 0] coordinate is at the bottom-left corner of the grid. .. _np.genfromtxt: https://numpy.org/doc/stable/reference/generated/numpy.genfromtxt.html - .. _mesa-examples Sugarscape model: https://github.com/projectmesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 + .. _mesa-examples Sugarscape model: https://github.com/mesa/mesa-examples/blob/e137a60e4e2f2546901bec497e79c4a7b0cc69bb/examples/sugarscape_g1mt/sugarscape_g1mt/model.py#L93-L94 """ _cells_capacity: ( diff --git a/mkdocs.yml b/mkdocs.yml index 40afb584..aa427850 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,8 +1,8 @@ # Project information 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 +site_url: https://mesa.github.io/mesa-frames +repo_url: https://github.com/mesa/mesa-frames +repo_name: mesa/mesa-frames edit_uri: edit/main/docs/general/ docs_dir: docs/general @@ -42,9 +42,9 @@ theme: icon: repo: fontawesome/brands/github # Logo (PNG) - logo: https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.png + logo: https://raw.githubusercontent.com/mesa/mesa/main/docs/images/mesa_logo.png # Favicon (ICO) - favicon: https://raw.githubusercontent.com/projectmesa/mesa/main/docs/images/mesa_logo.ico + favicon: https://raw.githubusercontent.com/mesa/mesa/main/docs/images/mesa_logo.ico # Plugins plugins: @@ -82,7 +82,7 @@ markdown_extensions: - pymdownx.keys - pymdownx.magiclink: repo_url_shorthand: true - user: projectmesa + user: mesa repo: mesa-frames - pymdownx.mark - pymdownx.smartsymbols @@ -105,7 +105,7 @@ extra_javascript: extra: social: - icon: fontawesome/brands/github - link: https://github.com/projectmesa + link: https://github.com/mesa - icon: fontawesome/brands/python link: https://pypi.org/project/mesa-frames/ diff --git a/pyproject.toml b/pyproject.toml index 7f63a94c..ab83b905 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,15 +36,15 @@ dependencies = [ "pyarrow>=20.0.0", # polars._typing added in 1.0.0 "polars>=1.30.0", - "psycopg2-binary==2.9.10", + "psycopg2-binary>=2.9.10", "boto3>=1.35.91", "typing-extensions>=4.15.0", ] dynamic = ["version"] [project.urls] -Documentation = "https://projectmesa.github.io/mesa-frames" -Repository = "https://github.com/projectmesa/mesa-frames.git" +Documentation = "https://mesa.github.io/mesa-frames" +Repository = "https://github.com/mesa/mesa-frames.git" # โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -66,6 +66,7 @@ test = [ docs = [ { include-group = "typechecking" }, "jupytext>=1.17.3", + "typer>=0.9.0", "mkdocs-material>=9.6.14", "mkdocs-jupyter>=0.25.1", "mkdocs-git-revision-date-localized-plugin>=1.4.7", @@ -87,7 +88,7 @@ docs = [ dev = [ { include-group = "test" }, { include-group = "docs" }, - "numba>=0.60.0", + "numba>=0.62", "ruff>=0.11.12", "pre-commit>=4.2.0", ] diff --git a/tests/test_benchmarks_cli.py b/tests/test_benchmarks_cli.py new file mode 100644 index 00000000..d63de8a2 --- /dev/null +++ b/tests/test_benchmarks_cli.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from pathlib import Path + +import polars as pl +import pytest + +from benchmarks import cli + + +def _register_dummy_model(monkeypatch: pytest.MonkeyPatch) -> None: + # Small, fast no-op backends to avoid running real simulations in tests + b1 = cli.Backend(name="mesa", runner=lambda agents, steps, seed=None: None) + b2 = cli.Backend(name="frames", runner=lambda agents, steps, seed=None: None) + monkeypatch.setitem( + cli.MODELS, "dummy", cli.ModelConfig(name="dummy", backends=[b1, b2]) + ) + + +def _pick_timestamp_dir(base: Path) -> Path: + # After a run that writes, the base results dir will contain a single timestamp subdir + subs = [p for p in base.iterdir() if p.is_dir()] + assert len(subs) <= 1 + return subs[0] if subs else base + + +def test_summary_save_and_plot( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture +): + _register_dummy_model(monkeypatch) + cli.run( + models=["dummy"], + agents=[1], + steps=1, + repeats=1, + seed=1, + save=True, + plot=True, + results_dir=tmp_path, + ) + out = capsys.readouterr().out + assert "Unified benchmark outputs written:" in out + assert "CSVs under" in out and "plots under" in out + ts = _pick_timestamp_dir(tmp_path) + # CSV should be present + assert any(ts.glob("*_perf_*.csv")) + + +def test_summary_plot_only( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture +): + _register_dummy_model(monkeypatch) + cli.run( + models=["dummy"], + agents=[1], + steps=1, + repeats=1, + seed=1, + save=False, + plot=True, + results_dir=tmp_path, + ) + out = capsys.readouterr().out + assert "Unified benchmark outputs written:" in out + assert "plots under" in out and "CSVs under" not in out + ts = _pick_timestamp_dir(tmp_path) + # No CSVs should be present, but plots subdir should exist + assert not any(ts.glob("*_perf_*.csv")) + assert (ts / "plots").exists() + + +def test_summary_save_only( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture +): + _register_dummy_model(monkeypatch) + cli.run( + models=["dummy"], + agents=[1], + steps=1, + repeats=1, + seed=1, + save=True, + plot=False, + results_dir=tmp_path, + ) + out = capsys.readouterr().out + assert "Unified benchmark outputs written:" in out + assert "CSVs under" in out and "plots under" not in out + ts = _pick_timestamp_dir(tmp_path) + # CSV should be present, plots subdir should not exist + assert any(ts.glob("*_perf_*.csv")) + assert not (ts / "plots").exists() + + +def test_summary_neither( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture +): + _register_dummy_model(monkeypatch) + cli.run( + models=["dummy"], + agents=[1], + steps=1, + repeats=1, + seed=1, + save=False, + plot=False, + results_dir=tmp_path, + ) + out = capsys.readouterr().out + assert "Benchmark run completed (save=False, plot=False; no files written)." in out + # No directories created + assert not any(p for p in tmp_path.iterdir()) diff --git a/tests/test_sugarscape_mesa_cli.py b/tests/test_sugarscape_mesa_cli.py new file mode 100644 index 00000000..651143c7 --- /dev/null +++ b/tests/test_sugarscape_mesa_cli.py @@ -0,0 +1,26 @@ +from pathlib import Path + + +def test_saved_csv_echo_is_guarded_by_save_results_flag() -> None: + """Ensure the CSV-saved confirmation is printed only when results are saved. + + This is a light-weight static check that avoids importing heavy runtime + dependencies (e.g., `mesa`) in CI environments where they may be absent. + """ + path = Path("examples/sugarscape_ig/backend_mesa/model.py") + src = path.read_text() + + needle = 'typer.echo(f"Saved CSV results under {results_dir}")' + # Ensure the message exists in the file (we just moved it into the if block) + assert needle in src + + # For each occurrence, ensure it is preceded by an `if save_results:` within + # the previous 6 lines (simple heuristic that is robust enough for this fix). + lines = src.splitlines() + for idx, line in enumerate(lines): + if needle in line: + window_start = max(0, idx - 6) + window = "\n".join(lines[window_start:idx]) + assert "if save_results:" in window, ( + "Found CSV-saved confirmation not guarded by `if save_results:`" + ) diff --git a/uv.lock b/uv.lock index 4db6f107..61d824e6 100644 --- a/uv.lock +++ b/uv.lock @@ -995,25 +995,26 @@ wheels = [ [[package]] name = "llvmlite" -version = "0.44.0" +version = "0.46.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/89/6a/95a3d3610d5c75293d5dbbb2a76480d5d4eeba641557b69fe90af6c5b84e/llvmlite-0.44.0.tar.gz", hash = "sha256:07667d66a5d150abed9157ab6c0b9393c9356f229784a4385c02f99e94fc94d4", size = 171880, upload-time = "2025-01-20T11:14:41.342Z" } +sdist = { url = "https://files.pythonhosted.org/packages/74/cd/08ae687ba099c7e3d21fe2ea536500563ef1943c5105bf6ab4ee3829f68e/llvmlite-0.46.0.tar.gz", hash = "sha256:227c9fd6d09dce2783c18b754b7cd9d9b3b3515210c46acc2d3c5badd9870ceb", size = 193456, upload-time = "2025-12-08T18:15:36.295Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/e2/86b245397052386595ad726f9742e5223d7aea999b18c518a50e96c3aca4/llvmlite-0.44.0-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:eed7d5f29136bda63b6d7804c279e2b72e08c952b7c5df61f45db408e0ee52f3", size = 28132305, upload-time = "2025-01-20T11:12:53.936Z" }, - { url = "https://files.pythonhosted.org/packages/ff/ec/506902dc6870249fbe2466d9cf66d531265d0f3a1157213c8f986250c033/llvmlite-0.44.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ace564d9fa44bb91eb6e6d8e7754977783c68e90a471ea7ce913bff30bd62427", size = 26201090, upload-time = "2025-01-20T11:12:59.847Z" }, - { url = "https://files.pythonhosted.org/packages/99/fe/d030f1849ebb1f394bb3f7adad5e729b634fb100515594aca25c354ffc62/llvmlite-0.44.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c5d22c3bfc842668168a786af4205ec8e3ad29fb1bc03fd11fd48460d0df64c1", size = 42361858, upload-time = "2025-01-20T11:13:07.623Z" }, - { url = "https://files.pythonhosted.org/packages/d7/7a/ce6174664b9077fc673d172e4c888cb0b128e707e306bc33fff8c2035f0d/llvmlite-0.44.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f01a394e9c9b7b1d4e63c327b096d10f6f0ed149ef53d38a09b3749dcf8c9610", size = 41184200, upload-time = "2025-01-20T11:13:20.058Z" }, - { url = "https://files.pythonhosted.org/packages/5f/c6/258801143975a6d09a373f2641237992496e15567b907a4d401839d671b8/llvmlite-0.44.0-cp311-cp311-win_amd64.whl", hash = "sha256:d8489634d43c20cd0ad71330dde1d5bc7b9966937a263ff1ec1cebb90dc50955", size = 30331193, upload-time = "2025-01-20T11:13:26.976Z" }, - { url = "https://files.pythonhosted.org/packages/15/86/e3c3195b92e6e492458f16d233e58a1a812aa2bfbef9bdd0fbafcec85c60/llvmlite-0.44.0-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:1d671a56acf725bf1b531d5ef76b86660a5ab8ef19bb6a46064a705c6ca80aad", size = 28132297, upload-time = "2025-01-20T11:13:32.57Z" }, - { url = "https://files.pythonhosted.org/packages/d6/53/373b6b8be67b9221d12b24125fd0ec56b1078b660eeae266ec388a6ac9a0/llvmlite-0.44.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5f79a728e0435493611c9f405168682bb75ffd1fbe6fc360733b850c80a026db", size = 26201105, upload-time = "2025-01-20T11:13:38.744Z" }, - { url = "https://files.pythonhosted.org/packages/cb/da/8341fd3056419441286c8e26bf436923021005ece0bff5f41906476ae514/llvmlite-0.44.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0143a5ef336da14deaa8ec26c5449ad5b6a2b564df82fcef4be040b9cacfea9", size = 42361901, upload-time = "2025-01-20T11:13:46.711Z" }, - { url = "https://files.pythonhosted.org/packages/53/ad/d79349dc07b8a395a99153d7ce8b01d6fcdc9f8231355a5df55ded649b61/llvmlite-0.44.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d752f89e31b66db6f8da06df8b39f9b91e78c5feea1bf9e8c1fba1d1c24c065d", size = 41184247, upload-time = "2025-01-20T11:13:56.159Z" }, - { url = "https://files.pythonhosted.org/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, - { url = "https://files.pythonhosted.org/packages/89/24/4c0ca705a717514c2092b18476e7a12c74d34d875e05e4d742618ebbf449/llvmlite-0.44.0-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:319bddd44e5f71ae2689859b7203080716448a3cd1128fb144fe5c055219d516", size = 28132306, upload-time = "2025-01-20T11:14:09.035Z" }, - { url = "https://files.pythonhosted.org/packages/01/cf/1dd5a60ba6aee7122ab9243fd614abcf22f36b0437cbbe1ccf1e3391461c/llvmlite-0.44.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9c58867118bad04a0bb22a2e0068c693719658105e40009ffe95c7000fcde88e", size = 26201090, upload-time = "2025-01-20T11:14:15.401Z" }, - { url = "https://files.pythonhosted.org/packages/d2/1b/656f5a357de7135a3777bd735cc7c9b8f23b4d37465505bd0eaf4be9befe/llvmlite-0.44.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:46224058b13c96af1365290bdfebe9a6264ae62fb79b2b55693deed11657a8bf", size = 42361904, upload-time = "2025-01-20T11:14:22.949Z" }, - { url = "https://files.pythonhosted.org/packages/d8/e1/12c5f20cb9168fb3464a34310411d5ad86e4163c8ff2d14a2b57e5cc6bac/llvmlite-0.44.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0097052c32bf721a4efc03bd109d335dfa57d9bffb3d4c24cc680711b8b4fc", size = 41184245, upload-time = "2025-01-20T11:14:31.731Z" }, - { url = "https://files.pythonhosted.org/packages/d0/81/e66fc86539293282fd9cb7c9417438e897f369e79ffb62e1ae5e5154d4dd/llvmlite-0.44.0-cp313-cp313-win_amd64.whl", hash = "sha256:2fb7c4f2fb86cbae6dca3db9ab203eeea0e22d73b99bc2341cdf9de93612e930", size = 30331193, upload-time = "2025-01-20T11:14:38.578Z" }, + { url = "https://files.pythonhosted.org/packages/7a/a1/2ad4b2367915faeebe8447f0a057861f646dbf5fbbb3561db42c65659cf3/llvmlite-0.46.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:82f3d39b16f19aa1a56d5fe625883a6ab600d5cc9ea8906cca70ce94cabba067", size = 37232766, upload-time = "2025-12-08T18:14:48.836Z" }, + { url = "https://files.pythonhosted.org/packages/12/b5/99cf8772fdd846c07da4fd70f07812a3c8fd17ea2409522c946bb0f2b277/llvmlite-0.46.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a3df43900119803bbc52720e758c76f316a9a0f34612a886862dfe0a5591a17e", size = 56275175, upload-time = "2025-12-08T18:14:51.604Z" }, + { url = "https://files.pythonhosted.org/packages/38/f2/ed806f9c003563732da156139c45d970ee435bd0bfa5ed8de87ba972b452/llvmlite-0.46.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de183fefc8022d21b0aa37fc3e90410bc3524aed8617f0ff76732fc6c3af5361", size = 55128630, upload-time = "2025-12-08T18:14:55.107Z" }, + { url = "https://files.pythonhosted.org/packages/19/0c/8f5a37a65fc9b7b17408508145edd5f86263ad69c19d3574e818f533a0eb/llvmlite-0.46.0-cp311-cp311-win_amd64.whl", hash = "sha256:e8b10bc585c58bdffec9e0c309bb7d51be1f2f15e169a4b4d42f2389e431eb93", size = 38138652, upload-time = "2025-12-08T18:14:58.171Z" }, + { url = "https://files.pythonhosted.org/packages/2b/f8/4db016a5e547d4e054ff2f3b99203d63a497465f81ab78ec8eb2ff7b2304/llvmlite-0.46.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6b9588ad4c63b4f0175a3984b85494f0c927c6b001e3a246a3a7fb3920d9a137", size = 37232767, upload-time = "2025-12-08T18:15:00.737Z" }, + { url = "https://files.pythonhosted.org/packages/aa/85/4890a7c14b4fa54400945cb52ac3cd88545bbdb973c440f98ca41591cdc5/llvmlite-0.46.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3535bd2bb6a2d7ae4012681ac228e5132cdb75fefb1bcb24e33f2f3e0c865ed4", size = 56275176, upload-time = "2025-12-08T18:15:03.936Z" }, + { url = "https://files.pythonhosted.org/packages/6a/07/3d31d39c1a1a08cd5337e78299fca77e6aebc07c059fbd0033e3edfab45c/llvmlite-0.46.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4cbfd366e60ff87ea6cc62f50bc4cd800ebb13ed4c149466f50cf2163a473d1e", size = 55128630, upload-time = "2025-12-08T18:15:07.196Z" }, + { url = "https://files.pythonhosted.org/packages/2a/6b/d139535d7590a1bba1ceb68751bef22fadaa5b815bbdf0e858e3875726b2/llvmlite-0.46.0-cp312-cp312-win_amd64.whl", hash = "sha256:398b39db462c39563a97b912d4f2866cd37cba60537975a09679b28fbbc0fb38", size = 38138940, upload-time = "2025-12-08T18:15:10.162Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ff/3eba7eb0aed4b6fca37125387cd417e8c458e750621fce56d2c541f67fa8/llvmlite-0.46.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:30b60892d034bc560e0ec6654737aaa74e5ca327bd8114d82136aa071d611172", size = 37232767, upload-time = "2025-12-08T18:15:13.22Z" }, + { url = "https://files.pythonhosted.org/packages/0e/54/737755c0a91558364b9200702c3c9c15d70ed63f9b98a2c32f1c2aa1f3ba/llvmlite-0.46.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6cc19b051753368a9c9f31dc041299059ee91aceec81bd57b0e385e5d5bf1a54", size = 56275176, upload-time = "2025-12-08T18:15:16.339Z" }, + { url = "https://files.pythonhosted.org/packages/e6/91/14f32e1d70905c1c0aa4e6609ab5d705c3183116ca02ac6df2091868413a/llvmlite-0.46.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bca185892908f9ede48c0acd547fe4dc1bafefb8a4967d47db6cf664f9332d12", size = 55128629, upload-time = "2025-12-08T18:15:19.493Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a7/d526ae86708cea531935ae777b6dbcabe7db52718e6401e0fb9c5edea80e/llvmlite-0.46.0-cp313-cp313-win_amd64.whl", hash = "sha256:67438fd30e12349ebb054d86a5a1a57fd5e87d264d2451bcfafbbbaa25b82a35", size = 38138941, upload-time = "2025-12-08T18:15:22.536Z" }, + { url = "https://files.pythonhosted.org/packages/95/ae/af0ffb724814cc2ea64445acad05f71cff5f799bb7efb22e47ee99340dbc/llvmlite-0.46.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:d252edfb9f4ac1fcf20652258e3f102b26b03eef738dc8a6ffdab7d7d341d547", size = 37232768, upload-time = "2025-12-08T18:15:25.055Z" }, + { url = "https://files.pythonhosted.org/packages/c9/19/5018e5352019be753b7b07f7759cdabb69ca5779fea2494be8839270df4c/llvmlite-0.46.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:379fdd1c59badeff8982cb47e4694a6143bec3bb49aa10a466e095410522064d", size = 56275173, upload-time = "2025-12-08T18:15:28.109Z" }, + { url = "https://files.pythonhosted.org/packages/9f/c9/d57877759d707e84c082163c543853245f91b70c804115a5010532890f18/llvmlite-0.46.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e8cbfff7f6db0fa2c771ad24154e2a7e457c2444d7673e6de06b8b698c3b269", size = 55128628, upload-time = "2025-12-08T18:15:31.098Z" }, + { url = "https://files.pythonhosted.org/packages/30/a8/e61a8c2b3cc7a597073d9cde1fcbb567e9d827f1db30c93cf80422eac70d/llvmlite-0.46.0-cp314-cp314-win_amd64.whl", hash = "sha256:7821eda3ec1f18050f981819756631d60b6d7ab1a6cf806d9efefbe3f4082d61", size = 39153056, upload-time = "2025-12-08T18:15:33.938Z" }, ] [[package]] @@ -1242,6 +1243,7 @@ dev = [ { name = "sphinx-copybutton" }, { name = "sphinx-design" }, { name = "sphinx-rtd-theme" }, + { name = "typer" }, ] docs = [ { name = "autodocsumm" }, @@ -1261,6 +1263,7 @@ docs = [ { name = "sphinx-copybutton" }, { name = "sphinx-design" }, { name = "sphinx-rtd-theme" }, + { name = "typer" }, ] test = [ { name = "beartype" }, @@ -1276,7 +1279,7 @@ requires-dist = [ { name = "boto3", specifier = ">=1.35.91" }, { name = "numpy", specifier = ">=2.0.2" }, { name = "polars", specifier = ">=1.30.0" }, - { name = "psycopg2-binary", specifier = "==2.9.10" }, + { name = "psycopg2-binary", specifier = ">=2.9.10" }, { name = "pyarrow", specifier = ">=20.0.0" }, { name = "typing-extensions", specifier = ">=4.15.0" }, ] @@ -1292,7 +1295,7 @@ dev = [ { name = "mkdocs-jupyter", specifier = ">=0.25.1" }, { name = "mkdocs-material", specifier = ">=9.6.14" }, { name = "mkdocs-minify-plugin", specifier = ">=0.8.0" }, - { name = "numba", specifier = ">=0.60.0" }, + { name = "numba", specifier = ">=0.62" }, { name = "numpydoc", specifier = ">=1.8.0" }, { name = "pre-commit", specifier = ">=4.2.0" }, { name = "pydata-sphinx-theme", specifier = ">=0.16.1" }, @@ -1305,6 +1308,7 @@ 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" }, @@ -1324,6 +1328,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" }, @@ -1546,29 +1551,30 @@ wheels = [ [[package]] name = "numba" -version = "0.61.2" +version = "0.63.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "llvmlite" }, { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1c/a0/e21f57604304aa03ebb8e098429222722ad99176a4f979d34af1d1ee80da/numba-0.61.2.tar.gz", hash = "sha256:8750ee147940a6637b80ecf7f95062185ad8726c8c28a2295b8ec1160a196f7d", size = 2820615, upload-time = "2025-04-09T02:58:07.659Z" } +sdist = { url = "https://files.pythonhosted.org/packages/dc/60/0145d479b2209bd8fdae5f44201eceb8ce5a23e0ed54c71f57db24618665/numba-0.63.1.tar.gz", hash = "sha256:b320aa675d0e3b17b40364935ea52a7b1c670c9037c39cf92c49502a75902f4b", size = 2761666, upload-time = "2025-12-10T02:57:39.002Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/97/c99d1056aed767503c228f7099dc11c402906b42a4757fec2819329abb98/numba-0.61.2-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:efd3db391df53aaa5cfbee189b6c910a5b471488749fd6606c3f33fc984c2ae2", size = 2775825, upload-time = "2025-04-09T02:57:43.442Z" }, - { url = "https://files.pythonhosted.org/packages/95/9e/63c549f37136e892f006260c3e2613d09d5120672378191f2dc387ba65a2/numba-0.61.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:49c980e4171948ffebf6b9a2520ea81feed113c1f4890747ba7f59e74be84b1b", size = 2778695, upload-time = "2025-04-09T02:57:44.968Z" }, - { url = "https://files.pythonhosted.org/packages/97/c8/8740616c8436c86c1b9a62e72cb891177d2c34c2d24ddcde4c390371bf4c/numba-0.61.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3945615cd73c2c7eba2a85ccc9c1730c21cd3958bfcf5a44302abae0fb07bb60", size = 3829227, upload-time = "2025-04-09T02:57:46.63Z" }, - { url = "https://files.pythonhosted.org/packages/fc/06/66e99ae06507c31d15ff3ecd1f108f2f59e18b6e08662cd5f8a5853fbd18/numba-0.61.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:bbfdf4eca202cebade0b7d43896978e146f39398909a42941c9303f82f403a18", size = 3523422, upload-time = "2025-04-09T02:57:48.222Z" }, - { url = "https://files.pythonhosted.org/packages/0f/a4/2b309a6a9f6d4d8cfba583401c7c2f9ff887adb5d54d8e2e130274c0973f/numba-0.61.2-cp311-cp311-win_amd64.whl", hash = "sha256:76bcec9f46259cedf888041b9886e257ae101c6268261b19fda8cfbc52bec9d1", size = 2831505, upload-time = "2025-04-09T02:57:50.108Z" }, - { url = "https://files.pythonhosted.org/packages/b4/a0/c6b7b9c615cfa3b98c4c63f4316e3f6b3bbe2387740277006551784218cd/numba-0.61.2-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:34fba9406078bac7ab052efbf0d13939426c753ad72946baaa5bf9ae0ebb8dd2", size = 2776626, upload-time = "2025-04-09T02:57:51.857Z" }, - { url = "https://files.pythonhosted.org/packages/92/4a/fe4e3c2ecad72d88f5f8cd04e7f7cff49e718398a2fac02d2947480a00ca/numba-0.61.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4ddce10009bc097b080fc96876d14c051cc0c7679e99de3e0af59014dab7dfe8", size = 2779287, upload-time = "2025-04-09T02:57:53.658Z" }, - { url = "https://files.pythonhosted.org/packages/9a/2d/e518df036feab381c23a624dac47f8445ac55686ec7f11083655eb707da3/numba-0.61.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b1bb509d01f23d70325d3a5a0e237cbc9544dd50e50588bc581ba860c213546", size = 3885928, upload-time = "2025-04-09T02:57:55.206Z" }, - { url = "https://files.pythonhosted.org/packages/10/0f/23cced68ead67b75d77cfcca3df4991d1855c897ee0ff3fe25a56ed82108/numba-0.61.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:48a53a3de8f8793526cbe330f2a39fe9a6638efcbf11bd63f3d2f9757ae345cd", size = 3577115, upload-time = "2025-04-09T02:57:56.818Z" }, - { url = "https://files.pythonhosted.org/packages/68/1d/ddb3e704c5a8fb90142bf9dc195c27db02a08a99f037395503bfbc1d14b3/numba-0.61.2-cp312-cp312-win_amd64.whl", hash = "sha256:97cf4f12c728cf77c9c1d7c23707e4d8fb4632b46275f8f3397de33e5877af18", size = 2831929, upload-time = "2025-04-09T02:57:58.45Z" }, - { url = "https://files.pythonhosted.org/packages/0b/f3/0fe4c1b1f2569e8a18ad90c159298d862f96c3964392a20d74fc628aee44/numba-0.61.2-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:3a10a8fc9afac40b1eac55717cece1b8b1ac0b946f5065c89e00bde646b5b154", size = 2771785, upload-time = "2025-04-09T02:57:59.96Z" }, - { url = "https://files.pythonhosted.org/packages/e9/71/91b277d712e46bd5059f8a5866862ed1116091a7cb03bd2704ba8ebe015f/numba-0.61.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7d3bcada3c9afba3bed413fba45845f2fb9cd0d2b27dd58a1be90257e293d140", size = 2773289, upload-time = "2025-04-09T02:58:01.435Z" }, - { url = "https://files.pythonhosted.org/packages/0d/e0/5ea04e7ad2c39288c0f0f9e8d47638ad70f28e275d092733b5817cf243c9/numba-0.61.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bdbca73ad81fa196bd53dc12e3aaf1564ae036e0c125f237c7644fe64a4928ab", size = 3893918, upload-time = "2025-04-09T02:58:02.933Z" }, - { url = "https://files.pythonhosted.org/packages/17/58/064f4dcb7d7e9412f16ecf80ed753f92297e39f399c905389688cf950b81/numba-0.61.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f154aaea625fb32cfbe3b80c5456d514d416fcdf79733dd69c0df3a11348e9e", size = 3584056, upload-time = "2025-04-09T02:58:04.538Z" }, - { url = "https://files.pythonhosted.org/packages/af/a4/6d3a0f2d3989e62a18749e1e9913d5fa4910bbb3e3311a035baea6caf26d/numba-0.61.2-cp313-cp313-win_amd64.whl", hash = "sha256:59321215e2e0ac5fa928a8020ab00b8e57cda8a97384963ac0dfa4d4e6aa54e7", size = 2831846, upload-time = "2025-04-09T02:58:06.125Z" }, + { url = "https://files.pythonhosted.org/packages/70/90/5f8614c165d2e256fbc6c57028519db6f32e4982475a372bbe550ea0454c/numba-0.63.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b33db00f18ccc790ee9911ce03fcdfe9d5124637d1ecc266f5ae0df06e02fec3", size = 2680501, upload-time = "2025-12-10T02:57:09.797Z" }, + { url = "https://files.pythonhosted.org/packages/dc/9d/d0afc4cf915edd8eadd9b2ab5b696242886ee4f97720d9322650d66a88c6/numba-0.63.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7d31ea186a78a7c0f6b1b2a3fe68057fdb291b045c52d86232b5383b6cf4fc25", size = 3744945, upload-time = "2025-12-10T02:57:11.697Z" }, + { url = "https://files.pythonhosted.org/packages/05/a9/d82f38f2ab73f3be6f838a826b545b80339762ee8969c16a8bf1d39395a8/numba-0.63.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ed3bb2fbdb651d6aac394388130a7001aab6f4541837123a4b4ab8b02716530c", size = 3450827, upload-time = "2025-12-10T02:57:13.709Z" }, + { url = "https://files.pythonhosted.org/packages/18/3f/a9b106e93c5bd7434e65f044bae0d204e20aa7f7f85d72ceb872c7c04216/numba-0.63.1-cp311-cp311-win_amd64.whl", hash = "sha256:1ecbff7688f044b1601be70113e2fb1835367ee0b28ffa8f3adf3a05418c5c87", size = 2747262, upload-time = "2025-12-10T02:57:15.664Z" }, + { url = "https://files.pythonhosted.org/packages/14/9c/c0974cd3d00ff70d30e8ff90522ba5fbb2bcee168a867d2321d8d0457676/numba-0.63.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2819cd52afa5d8d04e057bdfd54367575105f8829350d8fb5e4066fb7591cc71", size = 2680981, upload-time = "2025-12-10T02:57:17.579Z" }, + { url = "https://files.pythonhosted.org/packages/cb/70/ea2bc45205f206b7a24ee68a159f5097c9ca7e6466806e7c213587e0c2b1/numba-0.63.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5cfd45dbd3d409e713b1ccfdc2ee72ca82006860254429f4ef01867fdba5845f", size = 3801656, upload-time = "2025-12-10T02:57:19.106Z" }, + { url = "https://files.pythonhosted.org/packages/0d/82/4f4ba4fd0f99825cbf3cdefd682ca3678be1702b63362011de6e5f71f831/numba-0.63.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:69a599df6976c03b7ecf15d05302696f79f7e6d10d620367407517943355bcb0", size = 3501857, upload-time = "2025-12-10T02:57:20.721Z" }, + { url = "https://files.pythonhosted.org/packages/af/fd/6540456efa90b5f6604a86ff50dabefb187e43557e9081adcad3be44f048/numba-0.63.1-cp312-cp312-win_amd64.whl", hash = "sha256:bbad8c63e4fc7eb3cdb2c2da52178e180419f7969f9a685f283b313a70b92af3", size = 2750282, upload-time = "2025-12-10T02:57:22.474Z" }, + { url = "https://files.pythonhosted.org/packages/57/f7/e19e6eff445bec52dde5bed1ebb162925a8e6f988164f1ae4b3475a73680/numba-0.63.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:0bd4fd820ef7442dcc07da184c3f54bb41d2bdb7b35bacf3448e73d081f730dc", size = 2680954, upload-time = "2025-12-10T02:57:24.145Z" }, + { url = "https://files.pythonhosted.org/packages/e9/6c/1e222edba1e20e6b113912caa9b1665b5809433cbcb042dfd133c6f1fd38/numba-0.63.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:53de693abe4be3bd4dee38e1c55f01c55ff644a6a3696a3670589e6e4c39cde2", size = 3809736, upload-time = "2025-12-10T02:57:25.836Z" }, + { url = "https://files.pythonhosted.org/packages/76/0a/590bad11a8b3feeac30a24d01198d46bdb76ad15c70d3a530691ce3cae58/numba-0.63.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:81227821a72a763c3d4ac290abbb4371d855b59fdf85d5af22a47c0e86bf8c7e", size = 3508854, upload-time = "2025-12-10T02:57:27.438Z" }, + { url = "https://files.pythonhosted.org/packages/4e/f5/3800384a24eed1e4d524669cdbc0b9b8a628800bb1e90d7bd676e5f22581/numba-0.63.1-cp313-cp313-win_amd64.whl", hash = "sha256:eb227b07c2ac37b09432a9bda5142047a2d1055646e089d4a240a2643e508102", size = 2750228, upload-time = "2025-12-10T02:57:30.36Z" }, + { url = "https://files.pythonhosted.org/packages/36/2f/53be2aa8a55ee2608ebe1231789cbb217f6ece7f5e1c685d2f0752e95a5b/numba-0.63.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f180883e5508940cc83de8a8bea37fc6dd20fbe4e5558d4659b8b9bef5ff4731", size = 2681153, upload-time = "2025-12-10T02:57:32.016Z" }, + { url = "https://files.pythonhosted.org/packages/13/91/53e59c86759a0648282368d42ba732c29524a745fd555ed1fb1df83febbe/numba-0.63.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f0938764afa82a47c0e895637a6c55547a42c9e1d35cac42285b1fa60a8b02bb", size = 3778718, upload-time = "2025-12-10T02:57:33.764Z" }, + { url = "https://files.pythonhosted.org/packages/6c/0c/2be19eba50b0b7636f6d1f69dfb2825530537708a234ba1ff34afc640138/numba-0.63.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f90a929fa5094e062d4e0368ede1f4497d5e40f800e80aa5222c4734236a2894", size = 3478712, upload-time = "2025-12-10T02:57:35.518Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5f/4d0c9e756732577a52211f31da13a3d943d185f7fb90723f56d79c696caa/numba-0.63.1-cp314-cp314-win_amd64.whl", hash = "sha256:8d6d5ce85f572ed4e1a135dbb8c0114538f9dd0e3657eeb0bb64ab204cbe2a8f", size = 2752161, upload-time = "2025-12-10T02:57:37.12Z" }, ] [[package]] @@ -2481,6 +2487,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" @@ -2793,6 +2808,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"