Adds numba documentation and examples for issue#133

docs/general/user-guide/notebooks/numba_acceleration.ipynb

@@ -0,0 +1,334 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Numba Acceleration in mesa-frames\n\n",
+    "## Introduction\n\n",
+    "This guide explains how to use Numba to accelerate agent-based models in mesa-frames. [Numba](https://numba.pydata.org/) is a just-in-time (JIT) compiler for Python that can significantly improve performance of numerical Python code by compiling it to optimized machine code at runtime.\n\n",
+    "Mesa-frames already offers substantial performance improvements over standard mesa by using DataFrame-based storage (especially with Polars), but for computationally intensive simulations, Numba can provide additional acceleration."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## When to Use Numba\n\n",
+    "Consider using Numba acceleration in the following scenarios:\n\n",
+    "1. **Large agent populations**: When your simulation involves thousands or millions of agents\n",
+    "2. **Computationally intensive agent methods**: When agents perform complex calculations or numerical operations\n",
+    "3. **Spatial operations**: For optimizing neighbor search and spatial movement calculations\n",
+    "4. **Performance bottlenecks**: When profiling reveals specific methods as performance bottlenecks"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Numba Integration Options\n\n",
+    "Mesa-frames supports several Numba integration approaches:\n\n",
+    "1. **CPU acceleration**: Standard Numba acceleration on a single CPU core\n",
+    "2. **Parallel CPU acceleration**: Utilizing multiple CPU cores for parallel processing\n",
+    "3. **GPU acceleration**: Leveraging NVIDIA GPUs through CUDA (requires a compatible GPU and CUDA installation)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Basic Implementation Pattern\n\n",
+    "The recommended pattern for implementing Numba acceleration in mesa-frames follows these steps:\n\n",
+    "1. Identify the performance-critical method in your agent class\n",
+    "2. Extract the numerical computation into a separate function\n",
+    "3. Decorate this function with Numba's `@jit`, `@vectorize`, or `@guvectorize` decorators\n",
+    "4. Call this accelerated function from your agent class method"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example: Basic Numba Acceleration\n\n",
+    "Here's a simple example of using Numba to accelerate an agent method:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import polars as pl\n",
+    "from numba import jit\n",
+    "from mesa_frames import AgentSetPolars, ModelDF\n",
+    "\n",
+    "\n",
+    "class MyAgentSet(AgentSetPolars):\n",
+    "    def __init__(self, model: ModelDF, n_agents: int):\n",
+    "        super().__init__(model)\n",
+    "        # Initialize agents\n",
+    "        self += pl.DataFrame(\n",
+    "            {\n",
+    "                \"unique_id\": pl.arange(n_agents, eager=True),\n",
+    "                \"value\": pl.ones(n_agents, eager=True),\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    def complex_calculation(self):\n",
+    "        # Extract data to numpy arrays for Numba processing\n",
+    "        values = self.agents[\"value\"].to_numpy()\n",
+    "\n",
+    "        # Call the Numba-accelerated function\n",
+    "        results = self._calculate_with_numba(values)\n",
+    "\n",
+    "        # Update the agent values\n",
+    "        self[\"value\"] = results\n",
+    "\n",
+    "    @staticmethod\n",
+    "    @jit(nopython=True)\n",
+    "    def _calculate_with_numba(values):\n",
+    "        # This function will be compiled by Numba\n",
+    "        result = np.empty_like(values)\n",
+    "        for i in range(len(values)):\n",
+    "            # Complex calculation that benefits from Numba\n",
+    "            result[i] = values[i] ** 2 + np.sin(values[i])\n",
+    "        return result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Advanced Implementation: Vectorized Operations\n\n",
+    "For even better performance, you can use Numba's vectorization capabilities:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import polars as pl\n",
+    "from numba import vectorize, float64\n",
+    "from mesa_frames import AgentSetPolars, ModelDF\n",
+    "\n",
+    "\n",
+    "class MyVectorizedAgentSet(AgentSetPolars):\n",
+    "    def __init__(self, model: ModelDF, n_agents: int):\n",
+    "        super().__init__(model)\n",
+    "        # Initialize agents\n",
+    "        self += pl.DataFrame(\n",
+    "            {\n",
+    "                \"unique_id\": pl.arange(n_agents, eager=True),\n",
+    "                \"value\": pl.ones(n_agents, eager=True),\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    def complex_calculation(self):\n",
+    "        # Extract data to numpy arrays\n",
+    "        values = self.agents[\"value\"].to_numpy()\n",
+    "\n",
+    "        # Call the vectorized function\n",
+    "        results = self._vectorized_calculation(values)\n",
+    "\n",
+    "        # Update the agent values\n",
+    "        self[\"value\"] = results\n",
+    "\n",
+    "    @staticmethod\n",
+    "    @vectorize([float64(float64)], nopython=True)\n",
+    "    def _vectorized_calculation(x):\n",
+    "        # This function will be applied to each element\n",
+    "        return x**2 + np.sin(x)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## GPU Acceleration with CUDA\n\n",
+    "If you have a compatible NVIDIA GPU, you can use Numba's CUDA capabilities for massive parallelization:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import polars as pl\n",
+    "import math\n",
+    "from numba import cuda\n",
+    "from mesa_frames import AgentSetPolars, ModelDF\n",
+    "\n",
+    "\n",
+    "class MyCudaAgentSet(AgentSetPolars):\n",
+    "    def __init__(self, model: ModelDF, n_agents: int):\n",
+    "        super().__init__(model)\n",
+    "        # Initialize agents\n",
+    "        self += pl.DataFrame(\n",
+    "            {\n",
+    "                \"unique_id\": pl.arange(n_agents, eager=True),\n",
+    "                \"value\": pl.ones(n_agents, eager=True),\n",
+    "            }\n",
+    "        )\n",
+    "\n",
+    "    def complex_calculation(self):\n",
+    "        # Extract data to numpy arrays\n",
+    "        values = self.agents[\"value\"].to_numpy()\n",
+    "\n",
+    "        # Prepare output array\n",
+    "        results = np.empty_like(values)\n",
+    "\n",
+    "        # Call the CUDA kernel\n",
+    "        threads_per_block = 256\n",
+    "        blocks_per_grid = (len(values) + threads_per_block - 1) // threads_per_block\n",
+    "        self._cuda_calculation[blocks_per_grid, threads_per_block](values, results)\n",
+    "\n",
+    "        # Update the agent values\n",
+    "        self[\"value\"] = results\n",
+    "\n",
+    "    @staticmethod\n",
+    "    @cuda.jit\n",
+    "    def _cuda_calculation(values, results):\n",
+    "        # Calculate thread index\n",
+    "        i = cuda.grid(1)\n",
+    "\n",
+    "        # Check array bounds\n",
+    "        if i < values.size:\n",
+    "            # Complex calculation\n",
+    "            results[i] = values[i] ** 2 + math.sin(values[i])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## General Usage Pattern with guvectorize\n\n",
+    "The Sugarscape example in mesa-frames demonstrates a more advanced pattern using `guvectorize`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from numba import guvectorize, int32\n",
+    "from mesa_frames import AgentSetPolars\n",
+    "\n",
+    "\n",
+    "class AgentSetWithNumba(AgentSetPolars):\n",
+    "    numba_target = \"cpu\"  # Can be \"cpu\", \"parallel\", or \"cuda\"\n",
+    "\n",
+    "    def _get_accelerated_function(self):\n",
+    "        @guvectorize(\n",
+    "            [\n",
+    "                (\n",
+    "                    int32[:],  # input array 1\n",
+    "                    int32[:],  # input array 2\n",
+    "                    # ... more input arrays\n",
+    "                    int32[:],  # output array\n",
+    "                )\n",
+    "            ],\n",
+    "            \"(n), (m), ... -> (p)\",  # Signature defining array shapes\n",
+    "            nopython=True,\n",
+    "            target=self.numba_target,\n",
+    "        )\n",
+    "        def vectorized_function(input1, input2, output):\n",
+    "            # Function implementation\n",
+    "            # This will be compiled for the specified target\n",
+    "            # (CPU, parallel, or CUDA)\n",
+    "\n",
+    "            # Perform calculations and populate output array\n",
+    "            pass\n",
+    "\n",
+    "        return vectorized_function"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Real-World Example: Sugarscape Implementation\n\n",
+    "The mesa-frames repository includes a complete example of Numba acceleration in the Sugarscape model.\n",
+    "The implementation includes three variants:\n\n",
+    "1. **AntPolarsNumbaCPU**: Single-core CPU acceleration\n",
+    "2. **AntPolarsNumbaParallel**: Multi-core CPU acceleration\n",
+    "3. **AntPolarsNumbaGPU**: GPU acceleration using CUDA\n\n",
+    "You can find this implementation in the `examples/sugarscape_ig/ss_polars/agents.py` file."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Performance Considerations\n\n",
+    "When using Numba with mesa-frames, keep the following in mind:\n\n",
+    "1. **Compilation overhead**: The first call to a Numba function includes compilation time\n",
+    "2. **Data transfer overhead**: Moving data between DataFrame and NumPy arrays has a cost\n",
+    "3. **Function complexity**: Numba benefits most for computationally intensive functions\n",
+    "4. **Best practices**: Follow [Numba's best practices](https://numba.pydata.org/numba-doc/latest/user/performance-tips.html) for maximum performance"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Installation\n\n",
+    "To use Numba with mesa-frames, install it as an optional dependency:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# In your terminal, run one of these commands:\n",
+    "# pip install mesa-frames[numba]\n",
+    "#\n",
+    "# Or if you're installing from source:\n",
+    "# pip install -e \".[numba]\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Conclusion\n\n",
+    "Numba acceleration provides a powerful way to optimize performance-critical parts of your mesa-frames models. By selectively applying Numba to computationally intensive methods, you can achieve significant performance improvements while maintaining the overall structure and readability of your model code.\n\n",
+    "The mesa-frames repository includes complete examples of Numba acceleration, including:\n\n",
+    "1. The diffusion example in `examples/numba_example`\n",
+    "2. The Sugarscape implementation with Numba variants in `examples/sugarscape_ig/ss_polars/agents.py`\n\n",
+    "These examples demonstrate how to effectively integrate Numba with mesa-frames in real-world models."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

examples/numba_example/__init__.py

@@ -0,0 +1,6 @@
+"""
+Numba acceleration example for mesa-frames.
+
+This example demonstrates how to use Numba to accelerate agent-based models in mesa-frames.
+It compares standard Polars implementations with several Numba acceleration approaches.
+"""