From 13f3d6023fdcf879c29445fecf9e730639e46b3a Mon Sep 17 00:00:00 2001
From: Suryansh Garg <133861447+suryanshgargbpgc@users.noreply.github.com>
Date: Tue, 25 Mar 2025 17:27:19 +0530
Subject: [PATCH 1/4] Adds numba doc and examples for issue#133

1.adds doc on numba usage.
2. include example with performance comparision.
3.adds tests for numba acceleration functions.
4. makes numba available as an optional dependency.
---
 .../user-guide/5_numba_acceleration.md        | 224 ++++++++++++++
 examples/numba_example/__init__.py            |   6 +
 examples/numba_example/model.py               | 289 ++++++++++++++++++
 mkdocs.yml                                    |   1 +
 pyproject.toml                                |   7 +-
 tests/test_numba_acceleration.py              |  86 ++++++
 6 files changed, 611 insertions(+), 2 deletions(-)
 create mode 100644 docs/general/user-guide/5_numba_acceleration.md
 create mode 100644 examples/numba_example/__init__.py
 create mode 100644 examples/numba_example/model.py
 create mode 100644 tests/test_numba_acceleration.py

diff --git a/docs/general/user-guide/5_numba_acceleration.md b/docs/general/user-guide/5_numba_acceleration.md
new file mode 100644
index 00000000..92b2055d
--- /dev/null
+++ b/docs/general/user-guide/5_numba_acceleration.md
@@ -0,0 +1,224 @@
+# Numba Acceleration in mesa-frames
+
+## Introduction
+
+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.
+
+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.
+
+## When to Use Numba
+
+Consider using Numba acceleration in the following scenarios:
+
+1. **Large agent populations**: When your simulation involves thousands or millions of agents
+2. **Computationally intensive agent methods**: When agents perform complex calculations or numerical operations
+3. **Spatial operations**: For optimizing neighbor search and spatial movement calculations
+4. **Performance bottlenecks**: When profiling reveals specific methods as performance bottlenecks
+
+## Numba Integration Options
+
+Mesa-frames supports several Numba integration approaches:
+
+1. **CPU acceleration**: Standard Numba acceleration on a single CPU core
+2. **Parallel CPU acceleration**: Utilizing multiple CPU cores for parallel processing
+3. **GPU acceleration**: Leveraging NVIDIA GPUs through CUDA (requires a compatible GPU and CUDA installation)
+
+## Basic Implementation Pattern
+
+The recommended pattern for implementing Numba acceleration in mesa-frames follows these steps:
+
+1. Identify the performance-critical method in your agent class
+2. Extract the numerical computation into a separate function
+3. Decorate this function with Numba's `@jit`, `@vectorize`, or `@guvectorize` decorators
+4. Call this accelerated function from your agent class method
+
+## Example: Basic Numba Acceleration
+
+Here's a simple example of using Numba to accelerate an agent method:
+
+```python
+import numpy as np
+from numba import jit
+from mesa_frames import AgentSetPolars, ModelDF
+
+class MyAgentSet(AgentSetPolars):
+    def __init__(self, model: ModelDF, n_agents: int):
+        super().__init__(model)
+        # Initialize agents
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": pl.ones(n_agents, eager=True)
+        })
+    
+    def complex_calculation(self):
+        # Extract data to numpy arrays for Numba processing
+        values = self.agents["value"].to_numpy()
+        
+        # Call the Numba-accelerated function
+        results = self._calculate_with_numba(values)
+        
+        # Update the agent values
+        self["value"] = results
+    
+    @staticmethod
+    @jit(nopython=True)
+    def _calculate_with_numba(values):
+        # This function will be compiled by Numba
+        result = np.empty_like(values)
+        for i in range(len(values)):
+            # Complex calculation that benefits from Numba
+            result[i] = values[i] ** 2 + np.sin(values[i])
+        return result
+```
+
+## Advanced Implementation: Vectorized Operations
+
+For even better performance, you can use Numba's vectorization capabilities:
+
+```python
+import numpy as np
+from numba import vectorize, float64
+from mesa_frames import AgentSetPolars, ModelDF
+
+class MyVectorizedAgentSet(AgentSetPolars):
+    def __init__(self, model: ModelDF, n_agents: int):
+        super().__init__(model)
+        # Initialize agents
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": pl.ones(n_agents, eager=True)
+        })
+    
+    def complex_calculation(self):
+        # Extract data to numpy arrays
+        values = self.agents["value"].to_numpy()
+        
+        # Call the vectorized function
+        results = self._vectorized_calculation(values)
+        
+        # Update the agent values
+        self["value"] = results
+    
+    @staticmethod
+    @vectorize([float64(float64)], nopython=True)
+    def _vectorized_calculation(x):
+        # This function will be applied to each element
+        return x ** 2 + np.sin(x)
+```
+
+## GPU Acceleration with CUDA
+
+If you have a compatible NVIDIA GPU, you can use Numba's CUDA capabilities for massive parallelization:
+
+```python
+import numpy as np
+from numba import cuda
+from mesa_frames import AgentSetPolars, ModelDF
+
+class MyCudaAgentSet(AgentSetPolars):
+    def __init__(self, model: ModelDF, n_agents: int):
+        super().__init__(model)
+        # Initialize agents
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": pl.ones(n_agents, eager=True)
+        })
+    
+    def complex_calculation(self):
+        # Extract data to numpy arrays
+        values = self.agents["value"].to_numpy()
+        
+        # Prepare output array
+        results = np.empty_like(values)
+        
+        # Call the CUDA kernel
+        threads_per_block = 256
+        blocks_per_grid = (len(values) + threads_per_block - 1) // threads_per_block
+        self._cuda_calculation[blocks_per_grid, threads_per_block](values, results)
+        
+        # Update the agent values
+        self["value"] = results
+    
+    @staticmethod
+    @cuda.jit
+    def _cuda_calculation(values, results):
+        # Calculate thread index
+        i = cuda.grid(1)
+        
+        # Check array bounds
+        if i < values.size:
+            # Complex calculation
+            results[i] = values[i] ** 2 + math.sin(values[i])
+```
+
+## General Usage Pattern with guvectorize
+
+The Sugarscape example in mesa-frames demonstrates a more advanced pattern using `guvectorize`:
+
+```python
+class AgentSetWithNumba(AgentSetPolars):
+    numba_target = "cpu"  # Can be "cpu", "parallel", or "cuda"
+    
+    def _get_accelerated_function(self):
+        @guvectorize(
+            [
+                (
+                    int32[:],  # input array 1
+                    int32[:],  # input array 2
+                    # ... more input arrays
+                    int32[:],  # output array
+                )
+            ],
+            "(n), (m), ... -> (p)",  # Signature defining array shapes
+            nopython=True,
+            target=self.numba_target,
+        )
+        def vectorized_function(
+            input1, input2, ..., output
+        ):
+            # Function implementation
+            # This will be compiled for the specified target
+            # (CPU, parallel, or CUDA)
+            
+            # Perform calculations and populate output array
+            
+        return vectorized_function
+```
+
+## Real-World Example: Sugarscape Implementation
+
+The mesa-frames repository includes a complete example of Numba acceleration in the Sugarscape model. 
+The implementation includes three variants:
+
+1. **AntPolarsNumbaCPU**: Single-core CPU acceleration
+2. **AntPolarsNumbaParallel**: Multi-core CPU acceleration
+3. **AntPolarsNumbaGPU**: GPU acceleration using CUDA
+
+You can find this implementation in the `examples/sugarscape_ig/ss_polars/agents.py` file.
+
+## Performance Considerations
+
+When using Numba with mesa-frames, keep the following in mind:
+
+1. **Compilation overhead**: The first call to a Numba function includes compilation time
+2. **Data transfer overhead**: Moving data between DataFrame and NumPy arrays has a cost
+3. **Function complexity**: Numba benefits most for computationally intensive functions
+4. **Best practices**: Follow [Numba's best practices](https://numba.pydata.org/numba-doc/latest/user/performance-tips.html) for maximum performance
+
+## Installation
+
+To use Numba with mesa-frames, install it as an optional dependency:
+
+```bash
+pip install mesa-frames[numba]
+```
+
+Or if you're installing from source:
+
+```bash
+pip install -e ".[numba]"
+```
+
+## Conclusion
+
+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. 
\ No newline at end of file
diff --git a/examples/numba_example/__init__.py b/examples/numba_example/__init__.py
new file mode 100644
index 00000000..7791ecf0
--- /dev/null
+++ b/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.
+""" 
\ No newline at end of file
diff --git a/examples/numba_example/model.py b/examples/numba_example/model.py
new file mode 100644
index 00000000..a7e11270
--- /dev/null
+++ b/examples/numba_example/model.py
@@ -0,0 +1,289 @@
+"""
+A simple example demonstrating Numba acceleration in mesa-frames.
+
+This example compares different implementations of the same model:
+1. A standard Polars-based implementation
+2. A basic Numba-accelerated implementation
+3. A vectorized Numba implementation 
+4. A parallel Numba implementation
+
+The model simulates agents on a 2D grid with a simple diffusion process.
+"""
+
+import numpy as np
+import polars as pl
+from numba import jit, vectorize, prange, float64, int64
+from mesa_frames import AgentSetPolars, GridPolars, ModelDF
+
+
+class DiffusionAgentStandard(AgentSetPolars):
+    """Standard implementation using Polars without Numba."""
+    
+    def __init__(self, model, n_agents):
+        super().__init__(model)
+        # Initialize agents with random values
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": model.random.random(n_agents),
+        })
+    
+    def step(self):
+        """Standard implementation using pure Polars operations."""
+        # Get neighborhood
+        neighborhood = self.space.get_neighbors(agents=self, include_center=True)
+        
+        # Group by center agent to get all neighbors for each agent
+        grouped = neighborhood.group_by("unique_id_center")
+        
+        # For each agent, calculate new value based on neighbor average
+        for group in grouped:
+            center_id = group["unique_id_center"][0]
+            neighbor_ids = group["unique_id"]
+            neighbor_values = self[neighbor_ids, "value"]
+            new_value = neighbor_values.mean()
+            self[center_id, "value"] = new_value
+
+
+class DiffusionAgentNumba(AgentSetPolars):
+    """Implementation using basic Numba acceleration."""
+    
+    def __init__(self, model, n_agents):
+        super().__init__(model)
+        # Initialize agents with random values
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": model.random.random(n_agents),
+        })
+    
+    def step(self):
+        """Numba-accelerated implementation."""
+        # Get neighborhood
+        neighborhood = self.space.get_neighbors(agents=self, include_center=True)
+        
+        # Extract arrays for Numba processing
+        center_ids = neighborhood["unique_id_center"].to_numpy()
+        neighbor_ids = neighborhood["unique_id"].to_numpy()
+        values = self.agents["value"].to_numpy()
+        
+        # Use Numba to calculate new values
+        new_values = self._calculate_new_values(center_ids, neighbor_ids, values)
+        
+        # Update agent values
+        for i, agent_id in enumerate(self.agents["unique_id"]):
+            if i < len(new_values):
+                self[agent_id, "value"] = new_values[i]
+    
+    @staticmethod
+    @jit(nopython=True)
+    def _calculate_new_values(center_ids, neighbor_ids, values):
+        """Numba-accelerated calculation of new values."""
+        # Get unique center IDs
+        unique_centers = np.unique(center_ids)
+        new_values = np.zeros_like(unique_centers, dtype=np.float64)
+        
+        # For each center, calculate average of neighbors
+        for i, center in enumerate(unique_centers):
+            # Find indices where center_ids matches this center
+            indices = np.where(center_ids == center)[0]
+            
+            # Get neighbor IDs for this center
+            neighbors = neighbor_ids[indices]
+            
+            # Calculate mean of neighbor values
+            neighbor_values = np.array([values[n] for n in neighbors])
+            new_values[i] = np.mean(neighbor_values)
+            
+        return new_values
+
+
+class DiffusionAgentVectorized(AgentSetPolars):
+    """Implementation using vectorized Numba operations."""
+    
+    def __init__(self, model, n_agents):
+        super().__init__(model)
+        # Initialize agents with random values
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": model.random.random(n_agents),
+        })
+    
+    def step(self):
+        """Implementation using vectorized operations where possible."""
+        # Get neighborhood
+        neighborhood = self.space.get_neighbors(agents=self, include_center=True)
+        
+        # Extract data for processing
+        unique_ids = self.agents["unique_id"].to_numpy()
+        values = self.agents["value"].to_numpy()
+        
+        # Create a lookup dictionary for values
+        value_dict = {uid: val for uid, val in zip(unique_ids, values)}
+        
+        # Process the neighborhoods
+        new_values = np.zeros_like(values)
+        
+        # Group by center ID and process each group
+        for center_id, group in neighborhood.group_by("unique_id_center"):
+            neighbor_ids = group["unique_id"].to_numpy()
+            neighbor_values = np.array([value_dict[nid] for nid in neighbor_ids])
+            
+            # Use vectorized functions for calculations (mean in this case)
+            idx = np.where(unique_ids == center_id)[0][0]
+            new_values[idx] = np.mean(neighbor_values)  # Use NumPy's mean directly
+        
+        # Update all values at once
+        self["value"] = new_values
+    
+    # The vectorize decorator doesn't work with arrays as input types in this context
+    # We'll use a different approach with jit instead
+    @staticmethod
+    @jit(nopython=True)
+    def _calculate_mean(values):
+        """Numba-accelerated calculation of mean."""
+        return np.mean(values)
+
+
+class DiffusionAgentParallel(AgentSetPolars):
+    """Implementation using parallel Numba operations."""
+    
+    def __init__(self, model, n_agents):
+        super().__init__(model)
+        # Initialize agents with random values
+        self += pl.DataFrame({
+            "unique_id": pl.arange(n_agents, eager=True),
+            "value": model.random.random(n_agents),
+        })
+    
+    def step(self):
+        """Implementation using parallel processing."""
+        # Get neighborhood
+        neighborhood = self.space.get_neighbors(agents=self, include_center=True)
+        
+        # Process in parallel using Numba
+        unique_ids = self.agents["unique_id"].to_numpy()
+        values = self.agents["value"].to_numpy()
+        
+        # Create arrays for center IDs and their neighbors
+        centers = []
+        neighbors_list = []
+        
+        # Group neighborhoods by center ID
+        for center_id, group in neighborhood.group_by("unique_id_center"):
+            centers.append(center_id)
+            neighbors_list.append(group["unique_id"].to_numpy())
+        
+        # Convert to arrays for Numba
+        centers = np.array(centers)
+        max_neighbors = max(len(n) for n in neighbors_list)
+        
+        # Create 2D array of neighbor IDs with padding
+        neighbors_array = np.zeros((len(centers), max_neighbors), dtype=np.int64)
+        for i, neighbors in enumerate(neighbors_list):
+            neighbors_array[i, :len(neighbors)] = neighbors
+        
+        # Calculate new values in parallel
+        new_values = self._parallel_calculate(centers, neighbors_array, unique_ids, values, max_neighbors)
+        
+        # Update agent values
+        for center_id, new_value in zip(centers, new_values):
+            self[center_id, "value"] = new_value
+    
+    @staticmethod
+    @jit(nopython=True, parallel=True)
+    def _parallel_calculate(centers, neighbors_array, unique_ids, values, max_neighbors):
+        """Parallel calculation of new values using Numba."""
+        result = np.zeros_like(centers, dtype=np.float64)
+        
+        # Process each center in parallel
+        for i in prange(len(centers)):
+            center = centers[i]
+            neighbors = neighbors_array[i]
+            
+            # Calculate mean of neighbor values
+            sum_val = 0.0
+            count = 0
+            
+            for j in range(max_neighbors):
+                neighbor = neighbors[j]
+                if neighbor == 0 and j > 0:  # Padding value
+                    break
+                    
+                # Find this neighbor's value
+                for k in range(len(unique_ids)):
+                    if unique_ids[k] == neighbor:
+                        sum_val += values[k]
+                        count += 1
+                        break
+            
+            result[i] = sum_val / count if count > 0 else 0
+            
+        return result
+
+
+class DiffusionModel(ModelDF):
+    """Model demonstrating different implementation approaches."""
+    
+    def __init__(self, width, height, n_agents, agent_class):
+        super().__init__()
+        self.grid = GridPolars(self, (width, height))
+        self.agents += agent_class(self, n_agents)
+        self.grid.place_to_empty(self.agents)
+        
+    def step(self):
+        self.agents.do("step")
+        
+    def run(self, steps):
+        for _ in range(steps):
+            self.step()
+
+
+def run_comparison(width, height, n_agents, steps):
+    """Run and compare different implementations."""
+    import time
+    
+    results = {}
+    
+    for name, agent_class in [
+        ("Standard", DiffusionAgentStandard),
+        ("Basic Numba", DiffusionAgentNumba),
+        ("Vectorized", DiffusionAgentVectorized),
+        ("Parallel", DiffusionAgentParallel)
+    ]:
+        # Initialize model
+        model = DiffusionModel(width, height, n_agents, agent_class)
+        
+        # Run with timing
+        start = time.time()
+        model.run(steps)
+        end = time.time()
+        
+        results[name] = end - start
+        print(f"{name}: {results[name]:.4f} seconds")
+    
+    # Return the results
+    return results
+
+
+if __name__ == "__main__":
+    print("Running comparison of Numba acceleration approaches")
+    results = run_comparison(
+        width=50,
+        height=50,
+        n_agents=1000,
+        steps=10
+    )
+    
+    # Plot results if matplotlib is available
+    try:
+        import matplotlib.pyplot as plt
+        
+        plt.figure(figsize=(10, 6))
+        plt.bar(results.keys(), results.values())
+        plt.ylabel("Execution time (seconds)")
+        plt.title("Comparison of Numba Acceleration Approaches")
+        plt.savefig("numba_comparison.png")
+        plt.close()
+        
+        print("Results saved to numba_comparison.png")
+    except ImportError:
+        print("Matplotlib not available for plotting") 
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index e7eb915b..e6effb54 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -112,6 +112,7 @@ nav:
     - Introductory Tutorial: user-guide/2_introductory-tutorial.md
     - Advanced Tutorial: user-guide/3_advanced-tutorial.md
     - Benchmarks: user-guide/4_benchmarks.md
+    - Numba Acceleration: user-guide/numba_acceleration.md
   - API Reference: api/index.html
   - Contributing:
     - Contribution Guide: contributing.md
diff --git a/pyproject.toml b/pyproject.toml
index c81ebe01..9b4c8d39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,6 +63,10 @@ mkdocs = [
   "mkdocs-include-markdown-plugin"
 ]
 
+numba = [
+  "numba>=0.56.0",
+]
+
 sphinx = [
   "sphinx",
   "sphinx-rtd-theme",
@@ -87,9 +91,8 @@ test = [
 ]
 
 dev = [
-  "mesa_frames[test, docs]",
+  "mesa_frames[test, docs, numba]",
   "mesa~=2.4.0",
-  "numba",
 ]
 
 [tool.hatch.envs.test]
diff --git a/tests/test_numba_acceleration.py b/tests/test_numba_acceleration.py
new file mode 100644
index 00000000..a8354c27
--- /dev/null
+++ b/tests/test_numba_acceleration.py
@@ -0,0 +1,86 @@
+"""
+Unit tests for Numba acceleration in mesa-frames.
+
+This module tests the Numba acceleration functionality in mesa-frames, focusing
+on the acceleration functions themselves rather than full model integration.
+"""
+
+import numpy as np
+import pytest
+
+from examples.numba_example.model import (
+    DiffusionAgentNumba,
+    DiffusionAgentVectorized,
+    DiffusionAgentParallel
+)
+
+
+class Test_NumbaAcceleration:
+    """Test suite for Numba acceleration functionality."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Set up common test parameters."""
+        self.seed = 42
+
+    def test_basic_numba_function(self):
+        """Test the basic Numba function directly."""
+        # Test the basic Numba function
+        center_ids = np.array([0, 0, 0, 1, 1, 2])
+        neighbor_ids = np.array([0, 1, 2, 0, 1, 2])
+        values = np.array([0.5, 0.3, 0.8])
+        
+        new_values = DiffusionAgentNumba._calculate_new_values(
+            center_ids, neighbor_ids, values
+        )
+        
+        # We should get 3 new values (one for each unique center ID)
+        assert len(new_values) == 3
+        
+        # For center ID 0, the neighbors are [0, 1, 2], so mean is (0.5 + 0.3 + 0.8)/3 = 0.53333
+        assert abs(new_values[0] - 0.53333) < 0.0001
+        
+        # For center ID 1, the neighbors are [0, 1], so mean is (0.5 + 0.3)/2 = 0.4
+        assert abs(new_values[1] - 0.4) < 0.0001
+        
+        # For center ID 2, the neighbor is [2], so mean is 0.8
+        assert abs(new_values[2] - 0.8) < 0.0001
+
+    def test_vectorized_function(self):
+        """Test the vectorized Numba function."""
+        # Call the vectorized function directly
+        values = np.array([0.5, 0.3, 0.8])
+        
+        # The function is now a standard Numba jit function
+        result = DiffusionAgentVectorized._calculate_mean(values)
+        
+        # The mean should be (0.5 + 0.3 + 0.8)/3 = 0.53333
+        assert abs(result - 0.53333) < 0.0001
+
+    def test_parallel_function(self):
+        """Test the parallel Numba function."""
+        # Set up test data
+        centers = np.array([0, 1, 2])
+        neighbors_array = np.array([
+            [0, 1, 2],  # neighbors for center 0
+            [0, 1, 0],  # neighbors for center 1
+            [2, 0, 0]   # neighbors for center 2
+        ])
+        unique_ids = np.array([0, 1, 2])
+        values = np.array([0.5, 0.3, 0.8])
+        max_neighbors = 3
+        
+        # Call the parallel function
+        results = DiffusionAgentParallel._parallel_calculate(
+            centers, neighbors_array, unique_ids, values, max_neighbors
+        )
+        
+        # Check the results
+        # For center 0, the neighbors are [0, 1, 2], so mean is (0.5 + 0.3 + 0.8)/3 = 0.53333
+        assert abs(results[0] - 0.53333) < 0.0001
+        
+        # For center 1, the neighbors are [0, 1], so mean is (0.5 + 0.3)/2 = 0.4
+        assert abs(results[1] - 0.4) < 0.0001
+        
+        # For center 2, the neighbor is [2], so mean is 0.8/1 = 0.8
+        assert abs(results[2] - 0.8) < 0.0001 
\ No newline at end of file

From 8d6b6e93b44ad016f86cdb1584a15f3f3e83c39e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 25 Mar 2025 12:01:00 +0000
Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../user-guide/5_numba_acceleration.md        |  38 ++--
 examples/numba_example/__init__.py            |   2 +-
 examples/numba_example/model.py               | 163 +++++++++---------
 tests/test_numba_acceleration.py              |  38 ++--
 4 files changed, 125 insertions(+), 116 deletions(-)

diff --git a/docs/general/user-guide/5_numba_acceleration.md b/docs/general/user-guide/5_numba_acceleration.md
index 92b2055d..f4ad7f11 100644
--- a/docs/general/user-guide/5_numba_acceleration.md
+++ b/docs/general/user-guide/5_numba_acceleration.md
@@ -49,17 +49,17 @@ class MyAgentSet(AgentSetPolars):
             "unique_id": pl.arange(n_agents, eager=True),
             "value": pl.ones(n_agents, eager=True)
         })
-    
+
     def complex_calculation(self):
         # Extract data to numpy arrays for Numba processing
         values = self.agents["value"].to_numpy()
-        
+
         # Call the Numba-accelerated function
         results = self._calculate_with_numba(values)
-        
+
         # Update the agent values
         self["value"] = results
-    
+
     @staticmethod
     @jit(nopython=True)
     def _calculate_with_numba(values):
@@ -88,17 +88,17 @@ class MyVectorizedAgentSet(AgentSetPolars):
             "unique_id": pl.arange(n_agents, eager=True),
             "value": pl.ones(n_agents, eager=True)
         })
-    
+
     def complex_calculation(self):
         # Extract data to numpy arrays
         values = self.agents["value"].to_numpy()
-        
+
         # Call the vectorized function
         results = self._vectorized_calculation(values)
-        
+
         # Update the agent values
         self["value"] = results
-    
+
     @staticmethod
     @vectorize([float64(float64)], nopython=True)
     def _vectorized_calculation(x):
@@ -123,28 +123,28 @@ class MyCudaAgentSet(AgentSetPolars):
             "unique_id": pl.arange(n_agents, eager=True),
             "value": pl.ones(n_agents, eager=True)
         })
-    
+
     def complex_calculation(self):
         # Extract data to numpy arrays
         values = self.agents["value"].to_numpy()
-        
+
         # Prepare output array
         results = np.empty_like(values)
-        
+
         # Call the CUDA kernel
         threads_per_block = 256
         blocks_per_grid = (len(values) + threads_per_block - 1) // threads_per_block
         self._cuda_calculation[blocks_per_grid, threads_per_block](values, results)
-        
+
         # Update the agent values
         self["value"] = results
-    
+
     @staticmethod
     @cuda.jit
     def _cuda_calculation(values, results):
         # Calculate thread index
         i = cuda.grid(1)
-        
+
         # Check array bounds
         if i < values.size:
             # Complex calculation
@@ -158,7 +158,7 @@ The Sugarscape example in mesa-frames demonstrates a more advanced pattern using
 ```python
 class AgentSetWithNumba(AgentSetPolars):
     numba_target = "cpu"  # Can be "cpu", "parallel", or "cuda"
-    
+
     def _get_accelerated_function(self):
         @guvectorize(
             [
@@ -179,15 +179,15 @@ class AgentSetWithNumba(AgentSetPolars):
             # Function implementation
             # This will be compiled for the specified target
             # (CPU, parallel, or CUDA)
-            
+
             # Perform calculations and populate output array
-            
+
         return vectorized_function
 ```
 
 ## Real-World Example: Sugarscape Implementation
 
-The mesa-frames repository includes a complete example of Numba acceleration in the Sugarscape model. 
+The mesa-frames repository includes a complete example of Numba acceleration in the Sugarscape model.
 The implementation includes three variants:
 
 1. **AntPolarsNumbaCPU**: Single-core CPU acceleration
@@ -221,4 +221,4 @@ pip install -e ".[numba]"
 
 ## Conclusion
 
-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. 
\ No newline at end of file
+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.
diff --git a/examples/numba_example/__init__.py b/examples/numba_example/__init__.py
index 7791ecf0..7df8352b 100644
--- a/examples/numba_example/__init__.py
+++ b/examples/numba_example/__init__.py
@@ -3,4 +3,4 @@
 
 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.
-""" 
\ No newline at end of file
+"""
diff --git a/examples/numba_example/model.py b/examples/numba_example/model.py
index a7e11270..9c2f3072 100644
--- a/examples/numba_example/model.py
+++ b/examples/numba_example/model.py
@@ -4,7 +4,7 @@
 This example compares different implementations of the same model:
 1. A standard Polars-based implementation
 2. A basic Numba-accelerated implementation
-3. A vectorized Numba implementation 
+3. A vectorized Numba implementation
 4. A parallel Numba implementation
 
 The model simulates agents on a 2D grid with a simple diffusion process.
@@ -18,23 +18,25 @@
 
 class DiffusionAgentStandard(AgentSetPolars):
     """Standard implementation using Polars without Numba."""
-    
+
     def __init__(self, model, n_agents):
         super().__init__(model)
         # Initialize agents with random values
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": model.random.random(n_agents),
-        })
-    
+        self += pl.DataFrame(
+            {
+                "unique_id": pl.arange(n_agents, eager=True),
+                "value": model.random.random(n_agents),
+            }
+        )
+
     def step(self):
         """Standard implementation using pure Polars operations."""
         # Get neighborhood
         neighborhood = self.space.get_neighbors(agents=self, include_center=True)
-        
+
         # Group by center agent to get all neighbors for each agent
         grouped = neighborhood.group_by("unique_id_center")
-        
+
         # For each agent, calculate new value based on neighbor average
         for group in grouped:
             center_id = group["unique_id_center"][0]
@@ -46,33 +48,35 @@ def step(self):
 
 class DiffusionAgentNumba(AgentSetPolars):
     """Implementation using basic Numba acceleration."""
-    
+
     def __init__(self, model, n_agents):
         super().__init__(model)
         # Initialize agents with random values
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": model.random.random(n_agents),
-        })
-    
+        self += pl.DataFrame(
+            {
+                "unique_id": pl.arange(n_agents, eager=True),
+                "value": model.random.random(n_agents),
+            }
+        )
+
     def step(self):
         """Numba-accelerated implementation."""
         # Get neighborhood
         neighborhood = self.space.get_neighbors(agents=self, include_center=True)
-        
+
         # Extract arrays for Numba processing
         center_ids = neighborhood["unique_id_center"].to_numpy()
         neighbor_ids = neighborhood["unique_id"].to_numpy()
         values = self.agents["value"].to_numpy()
-        
+
         # Use Numba to calculate new values
         new_values = self._calculate_new_values(center_ids, neighbor_ids, values)
-        
+
         # Update agent values
         for i, agent_id in enumerate(self.agents["unique_id"]):
             if i < len(new_values):
                 self[agent_id, "value"] = new_values[i]
-    
+
     @staticmethod
     @jit(nopython=True)
     def _calculate_new_values(center_ids, neighbor_ids, values):
@@ -80,60 +84,62 @@ def _calculate_new_values(center_ids, neighbor_ids, values):
         # Get unique center IDs
         unique_centers = np.unique(center_ids)
         new_values = np.zeros_like(unique_centers, dtype=np.float64)
-        
+
         # For each center, calculate average of neighbors
         for i, center in enumerate(unique_centers):
             # Find indices where center_ids matches this center
             indices = np.where(center_ids == center)[0]
-            
+
             # Get neighbor IDs for this center
             neighbors = neighbor_ids[indices]
-            
+
             # Calculate mean of neighbor values
             neighbor_values = np.array([values[n] for n in neighbors])
             new_values[i] = np.mean(neighbor_values)
-            
+
         return new_values
 
 
 class DiffusionAgentVectorized(AgentSetPolars):
     """Implementation using vectorized Numba operations."""
-    
+
     def __init__(self, model, n_agents):
         super().__init__(model)
         # Initialize agents with random values
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": model.random.random(n_agents),
-        })
-    
+        self += pl.DataFrame(
+            {
+                "unique_id": pl.arange(n_agents, eager=True),
+                "value": model.random.random(n_agents),
+            }
+        )
+
     def step(self):
         """Implementation using vectorized operations where possible."""
         # Get neighborhood
         neighborhood = self.space.get_neighbors(agents=self, include_center=True)
-        
+
         # Extract data for processing
         unique_ids = self.agents["unique_id"].to_numpy()
         values = self.agents["value"].to_numpy()
-        
+
         # Create a lookup dictionary for values
         value_dict = {uid: val for uid, val in zip(unique_ids, values)}
-        
+
         # Process the neighborhoods
         new_values = np.zeros_like(values)
-        
+
         # Group by center ID and process each group
         for center_id, group in neighborhood.group_by("unique_id_center"):
             neighbor_ids = group["unique_id"].to_numpy()
             neighbor_values = np.array([value_dict[nid] for nid in neighbor_ids])
-            
+
             # Use vectorized functions for calculations (mean in this case)
             idx = np.where(unique_ids == center_id)[0][0]
             new_values[idx] = np.mean(neighbor_values)  # Use NumPy's mean directly
-        
+
         # Update all values at once
         self["value"] = new_values
-    
+
     # The vectorize decorator doesn't work with arrays as input types in this context
     # We'll use a different approach with jit instead
     @staticmethod
@@ -145,93 +151,99 @@ def _calculate_mean(values):
 
 class DiffusionAgentParallel(AgentSetPolars):
     """Implementation using parallel Numba operations."""
-    
+
     def __init__(self, model, n_agents):
         super().__init__(model)
         # Initialize agents with random values
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": model.random.random(n_agents),
-        })
-    
+        self += pl.DataFrame(
+            {
+                "unique_id": pl.arange(n_agents, eager=True),
+                "value": model.random.random(n_agents),
+            }
+        )
+
     def step(self):
         """Implementation using parallel processing."""
         # Get neighborhood
         neighborhood = self.space.get_neighbors(agents=self, include_center=True)
-        
+
         # Process in parallel using Numba
         unique_ids = self.agents["unique_id"].to_numpy()
         values = self.agents["value"].to_numpy()
-        
+
         # Create arrays for center IDs and their neighbors
         centers = []
         neighbors_list = []
-        
+
         # Group neighborhoods by center ID
         for center_id, group in neighborhood.group_by("unique_id_center"):
             centers.append(center_id)
             neighbors_list.append(group["unique_id"].to_numpy())
-        
+
         # Convert to arrays for Numba
         centers = np.array(centers)
         max_neighbors = max(len(n) for n in neighbors_list)
-        
+
         # Create 2D array of neighbor IDs with padding
         neighbors_array = np.zeros((len(centers), max_neighbors), dtype=np.int64)
         for i, neighbors in enumerate(neighbors_list):
-            neighbors_array[i, :len(neighbors)] = neighbors
-        
+            neighbors_array[i, : len(neighbors)] = neighbors
+
         # Calculate new values in parallel
-        new_values = self._parallel_calculate(centers, neighbors_array, unique_ids, values, max_neighbors)
-        
+        new_values = self._parallel_calculate(
+            centers, neighbors_array, unique_ids, values, max_neighbors
+        )
+
         # Update agent values
         for center_id, new_value in zip(centers, new_values):
             self[center_id, "value"] = new_value
-    
+
     @staticmethod
     @jit(nopython=True, parallel=True)
-    def _parallel_calculate(centers, neighbors_array, unique_ids, values, max_neighbors):
+    def _parallel_calculate(
+        centers, neighbors_array, unique_ids, values, max_neighbors
+    ):
         """Parallel calculation of new values using Numba."""
         result = np.zeros_like(centers, dtype=np.float64)
-        
+
         # Process each center in parallel
         for i in prange(len(centers)):
             center = centers[i]
             neighbors = neighbors_array[i]
-            
+
             # Calculate mean of neighbor values
             sum_val = 0.0
             count = 0
-            
+
             for j in range(max_neighbors):
                 neighbor = neighbors[j]
                 if neighbor == 0 and j > 0:  # Padding value
                     break
-                    
+
                 # Find this neighbor's value
                 for k in range(len(unique_ids)):
                     if unique_ids[k] == neighbor:
                         sum_val += values[k]
                         count += 1
                         break
-            
+
             result[i] = sum_val / count if count > 0 else 0
-            
+
         return result
 
 
 class DiffusionModel(ModelDF):
     """Model demonstrating different implementation approaches."""
-    
+
     def __init__(self, width, height, n_agents, agent_class):
         super().__init__()
         self.grid = GridPolars(self, (width, height))
         self.agents += agent_class(self, n_agents)
         self.grid.place_to_empty(self.agents)
-        
+
     def step(self):
         self.agents.do("step")
-        
+
     def run(self, steps):
         for _ in range(steps):
             self.step()
@@ -240,50 +252,45 @@ def run(self, steps):
 def run_comparison(width, height, n_agents, steps):
     """Run and compare different implementations."""
     import time
-    
+
     results = {}
-    
+
     for name, agent_class in [
         ("Standard", DiffusionAgentStandard),
         ("Basic Numba", DiffusionAgentNumba),
         ("Vectorized", DiffusionAgentVectorized),
-        ("Parallel", DiffusionAgentParallel)
+        ("Parallel", DiffusionAgentParallel),
     ]:
         # Initialize model
         model = DiffusionModel(width, height, n_agents, agent_class)
-        
+
         # Run with timing
         start = time.time()
         model.run(steps)
         end = time.time()
-        
+
         results[name] = end - start
         print(f"{name}: {results[name]:.4f} seconds")
-    
+
     # Return the results
     return results
 
 
 if __name__ == "__main__":
     print("Running comparison of Numba acceleration approaches")
-    results = run_comparison(
-        width=50,
-        height=50,
-        n_agents=1000,
-        steps=10
-    )
-    
+    results = run_comparison(width=50, height=50, n_agents=1000, steps=10)
+
     # Plot results if matplotlib is available
     try:
         import matplotlib.pyplot as plt
-        
+
         plt.figure(figsize=(10, 6))
         plt.bar(results.keys(), results.values())
         plt.ylabel("Execution time (seconds)")
         plt.title("Comparison of Numba Acceleration Approaches")
         plt.savefig("numba_comparison.png")
         plt.close()
-        
+
         print("Results saved to numba_comparison.png")
     except ImportError:
-        print("Matplotlib not available for plotting") 
\ No newline at end of file
+        print("Matplotlib not available for plotting")
diff --git a/tests/test_numba_acceleration.py b/tests/test_numba_acceleration.py
index a8354c27..76451287 100644
--- a/tests/test_numba_acceleration.py
+++ b/tests/test_numba_acceleration.py
@@ -11,7 +11,7 @@
 from examples.numba_example.model import (
     DiffusionAgentNumba,
     DiffusionAgentVectorized,
-    DiffusionAgentParallel
+    DiffusionAgentParallel,
 )
 
 
@@ -29,20 +29,20 @@ def test_basic_numba_function(self):
         center_ids = np.array([0, 0, 0, 1, 1, 2])
         neighbor_ids = np.array([0, 1, 2, 0, 1, 2])
         values = np.array([0.5, 0.3, 0.8])
-        
+
         new_values = DiffusionAgentNumba._calculate_new_values(
             center_ids, neighbor_ids, values
         )
-        
+
         # We should get 3 new values (one for each unique center ID)
         assert len(new_values) == 3
-        
+
         # For center ID 0, the neighbors are [0, 1, 2], so mean is (0.5 + 0.3 + 0.8)/3 = 0.53333
         assert abs(new_values[0] - 0.53333) < 0.0001
-        
+
         # For center ID 1, the neighbors are [0, 1], so mean is (0.5 + 0.3)/2 = 0.4
         assert abs(new_values[1] - 0.4) < 0.0001
-        
+
         # For center ID 2, the neighbor is [2], so mean is 0.8
         assert abs(new_values[2] - 0.8) < 0.0001
 
@@ -50,10 +50,10 @@ def test_vectorized_function(self):
         """Test the vectorized Numba function."""
         # Call the vectorized function directly
         values = np.array([0.5, 0.3, 0.8])
-        
+
         # The function is now a standard Numba jit function
         result = DiffusionAgentVectorized._calculate_mean(values)
-        
+
         # The mean should be (0.5 + 0.3 + 0.8)/3 = 0.53333
         assert abs(result - 0.53333) < 0.0001
 
@@ -61,26 +61,28 @@ def test_parallel_function(self):
         """Test the parallel Numba function."""
         # Set up test data
         centers = np.array([0, 1, 2])
-        neighbors_array = np.array([
-            [0, 1, 2],  # neighbors for center 0
-            [0, 1, 0],  # neighbors for center 1
-            [2, 0, 0]   # neighbors for center 2
-        ])
+        neighbors_array = np.array(
+            [
+                [0, 1, 2],  # neighbors for center 0
+                [0, 1, 0],  # neighbors for center 1
+                [2, 0, 0],  # neighbors for center 2
+            ]
+        )
         unique_ids = np.array([0, 1, 2])
         values = np.array([0.5, 0.3, 0.8])
         max_neighbors = 3
-        
+
         # Call the parallel function
         results = DiffusionAgentParallel._parallel_calculate(
             centers, neighbors_array, unique_ids, values, max_neighbors
         )
-        
+
         # Check the results
         # For center 0, the neighbors are [0, 1, 2], so mean is (0.5 + 0.3 + 0.8)/3 = 0.53333
         assert abs(results[0] - 0.53333) < 0.0001
-        
+
         # For center 1, the neighbors are [0, 1], so mean is (0.5 + 0.3)/2 = 0.4
         assert abs(results[1] - 0.4) < 0.0001
-        
+
         # For center 2, the neighbor is [2], so mean is 0.8/1 = 0.8
-        assert abs(results[2] - 0.8) < 0.0001 
\ No newline at end of file
+        assert abs(results[2] - 0.8) < 0.0001

From f21e34144cf6c5ab2c6fabe85f91c667291042a7 Mon Sep 17 00:00:00 2001
From: Suryansh Garg <133861447+suryanshgargbpgc@users.noreply.github.com>
Date: Thu, 3 Apr 2025 04:18:13 +0530
Subject: [PATCH 3/4] Convert Numba acceleration documentation from markdown to
 Jupyter notebook as requested in PR review

---
 .../user-guide/5_numba_acceleration.md        | 224 ------------
 .../notebooks/numba_acceleration.ipynb        | 326 ++++++++++++++++++
 mkdocs.yml                                    |   2 +-
 3 files changed, 327 insertions(+), 225 deletions(-)
 delete mode 100644 docs/general/user-guide/5_numba_acceleration.md
 create mode 100644 docs/general/user-guide/notebooks/numba_acceleration.ipynb

diff --git a/docs/general/user-guide/5_numba_acceleration.md b/docs/general/user-guide/5_numba_acceleration.md
deleted file mode 100644
index f4ad7f11..00000000
--- a/docs/general/user-guide/5_numba_acceleration.md
+++ /dev/null
@@ -1,224 +0,0 @@
-# Numba Acceleration in mesa-frames
-
-## Introduction
-
-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.
-
-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.
-
-## When to Use Numba
-
-Consider using Numba acceleration in the following scenarios:
-
-1. **Large agent populations**: When your simulation involves thousands or millions of agents
-2. **Computationally intensive agent methods**: When agents perform complex calculations or numerical operations
-3. **Spatial operations**: For optimizing neighbor search and spatial movement calculations
-4. **Performance bottlenecks**: When profiling reveals specific methods as performance bottlenecks
-
-## Numba Integration Options
-
-Mesa-frames supports several Numba integration approaches:
-
-1. **CPU acceleration**: Standard Numba acceleration on a single CPU core
-2. **Parallel CPU acceleration**: Utilizing multiple CPU cores for parallel processing
-3. **GPU acceleration**: Leveraging NVIDIA GPUs through CUDA (requires a compatible GPU and CUDA installation)
-
-## Basic Implementation Pattern
-
-The recommended pattern for implementing Numba acceleration in mesa-frames follows these steps:
-
-1. Identify the performance-critical method in your agent class
-2. Extract the numerical computation into a separate function
-3. Decorate this function with Numba's `@jit`, `@vectorize`, or `@guvectorize` decorators
-4. Call this accelerated function from your agent class method
-
-## Example: Basic Numba Acceleration
-
-Here's a simple example of using Numba to accelerate an agent method:
-
-```python
-import numpy as np
-from numba import jit
-from mesa_frames import AgentSetPolars, ModelDF
-
-class MyAgentSet(AgentSetPolars):
-    def __init__(self, model: ModelDF, n_agents: int):
-        super().__init__(model)
-        # Initialize agents
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": pl.ones(n_agents, eager=True)
-        })
-
-    def complex_calculation(self):
-        # Extract data to numpy arrays for Numba processing
-        values = self.agents["value"].to_numpy()
-
-        # Call the Numba-accelerated function
-        results = self._calculate_with_numba(values)
-
-        # Update the agent values
-        self["value"] = results
-
-    @staticmethod
-    @jit(nopython=True)
-    def _calculate_with_numba(values):
-        # This function will be compiled by Numba
-        result = np.empty_like(values)
-        for i in range(len(values)):
-            # Complex calculation that benefits from Numba
-            result[i] = values[i] ** 2 + np.sin(values[i])
-        return result
-```
-
-## Advanced Implementation: Vectorized Operations
-
-For even better performance, you can use Numba's vectorization capabilities:
-
-```python
-import numpy as np
-from numba import vectorize, float64
-from mesa_frames import AgentSetPolars, ModelDF
-
-class MyVectorizedAgentSet(AgentSetPolars):
-    def __init__(self, model: ModelDF, n_agents: int):
-        super().__init__(model)
-        # Initialize agents
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": pl.ones(n_agents, eager=True)
-        })
-
-    def complex_calculation(self):
-        # Extract data to numpy arrays
-        values = self.agents["value"].to_numpy()
-
-        # Call the vectorized function
-        results = self._vectorized_calculation(values)
-
-        # Update the agent values
-        self["value"] = results
-
-    @staticmethod
-    @vectorize([float64(float64)], nopython=True)
-    def _vectorized_calculation(x):
-        # This function will be applied to each element
-        return x ** 2 + np.sin(x)
-```
-
-## GPU Acceleration with CUDA
-
-If you have a compatible NVIDIA GPU, you can use Numba's CUDA capabilities for massive parallelization:
-
-```python
-import numpy as np
-from numba import cuda
-from mesa_frames import AgentSetPolars, ModelDF
-
-class MyCudaAgentSet(AgentSetPolars):
-    def __init__(self, model: ModelDF, n_agents: int):
-        super().__init__(model)
-        # Initialize agents
-        self += pl.DataFrame({
-            "unique_id": pl.arange(n_agents, eager=True),
-            "value": pl.ones(n_agents, eager=True)
-        })
-
-    def complex_calculation(self):
-        # Extract data to numpy arrays
-        values = self.agents["value"].to_numpy()
-
-        # Prepare output array
-        results = np.empty_like(values)
-
-        # Call the CUDA kernel
-        threads_per_block = 256
-        blocks_per_grid = (len(values) + threads_per_block - 1) // threads_per_block
-        self._cuda_calculation[blocks_per_grid, threads_per_block](values, results)
-
-        # Update the agent values
-        self["value"] = results
-
-    @staticmethod
-    @cuda.jit
-    def _cuda_calculation(values, results):
-        # Calculate thread index
-        i = cuda.grid(1)
-
-        # Check array bounds
-        if i < values.size:
-            # Complex calculation
-            results[i] = values[i] ** 2 + math.sin(values[i])
-```
-
-## General Usage Pattern with guvectorize
-
-The Sugarscape example in mesa-frames demonstrates a more advanced pattern using `guvectorize`:
-
-```python
-class AgentSetWithNumba(AgentSetPolars):
-    numba_target = "cpu"  # Can be "cpu", "parallel", or "cuda"
-
-    def _get_accelerated_function(self):
-        @guvectorize(
-            [
-                (
-                    int32[:],  # input array 1
-                    int32[:],  # input array 2
-                    # ... more input arrays
-                    int32[:],  # output array
-                )
-            ],
-            "(n), (m), ... -> (p)",  # Signature defining array shapes
-            nopython=True,
-            target=self.numba_target,
-        )
-        def vectorized_function(
-            input1, input2, ..., output
-        ):
-            # Function implementation
-            # This will be compiled for the specified target
-            # (CPU, parallel, or CUDA)
-
-            # Perform calculations and populate output array
-
-        return vectorized_function
-```
-
-## Real-World Example: Sugarscape Implementation
-
-The mesa-frames repository includes a complete example of Numba acceleration in the Sugarscape model.
-The implementation includes three variants:
-
-1. **AntPolarsNumbaCPU**: Single-core CPU acceleration
-2. **AntPolarsNumbaParallel**: Multi-core CPU acceleration
-3. **AntPolarsNumbaGPU**: GPU acceleration using CUDA
-
-You can find this implementation in the `examples/sugarscape_ig/ss_polars/agents.py` file.
-
-## Performance Considerations
-
-When using Numba with mesa-frames, keep the following in mind:
-
-1. **Compilation overhead**: The first call to a Numba function includes compilation time
-2. **Data transfer overhead**: Moving data between DataFrame and NumPy arrays has a cost
-3. **Function complexity**: Numba benefits most for computationally intensive functions
-4. **Best practices**: Follow [Numba's best practices](https://numba.pydata.org/numba-doc/latest/user/performance-tips.html) for maximum performance
-
-## Installation
-
-To use Numba with mesa-frames, install it as an optional dependency:
-
-```bash
-pip install mesa-frames[numba]
-```
-
-Or if you're installing from source:
-
-```bash
-pip install -e ".[numba]"
-```
-
-## Conclusion
-
-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.
diff --git a/docs/general/user-guide/notebooks/numba_acceleration.ipynb b/docs/general/user-guide/notebooks/numba_acceleration.ipynb
new file mode 100644
index 00000000..cee91aab
--- /dev/null
+++ b/docs/general/user-guide/notebooks/numba_acceleration.ipynb
@@ -0,0 +1,326 @@
+{
+  "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",
+        "class MyAgentSet(AgentSetPolars):\n",
+        "    def __init__(self, model: ModelDF, n_agents: int):\n",
+        "        super().__init__(model)\n",
+        "        # Initialize agents\n",
+        "        self += pl.DataFrame({\n",
+        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
+        "            \"value\": pl.ones(n_agents, eager=True)\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",
+        "class MyVectorizedAgentSet(AgentSetPolars):\n",
+        "    def __init__(self, model: ModelDF, n_agents: int):\n",
+        "        super().__init__(model)\n",
+        "        # Initialize agents\n",
+        "        self += pl.DataFrame({\n",
+        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
+        "            \"value\": pl.ones(n_agents, eager=True)\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",
+        "class MyCudaAgentSet(AgentSetPolars):\n",
+        "    def __init__(self, model: ModelDF, n_agents: int):\n",
+        "        super().__init__(model)\n",
+        "        # Initialize agents\n",
+        "        self += pl.DataFrame({\n",
+        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
+        "            \"value\": pl.ones(n_agents, eager=True)\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",
+        "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(\n",
+        "            input1, input2, output\n",
+        "        ):\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
+}
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index e6effb54..192631c6 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -112,7 +112,7 @@ nav:
     - Introductory Tutorial: user-guide/2_introductory-tutorial.md
     - Advanced Tutorial: user-guide/3_advanced-tutorial.md
     - Benchmarks: user-guide/4_benchmarks.md
-    - Numba Acceleration: user-guide/numba_acceleration.md
+    - Numba Acceleration: user-guide/notebooks/numba_acceleration.ipynb
   - API Reference: api/index.html
   - Contributing:
     - Contribution Guide: contributing.md

From 8774406b5bd71ba8e529ec8bd5662971498d8870 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 2 Apr 2025 22:58:42 +0000
Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../notebooks/numba_acceleration.ipynb        | 654 +++++++++---------
 1 file changed, 331 insertions(+), 323 deletions(-)

diff --git a/docs/general/user-guide/notebooks/numba_acceleration.ipynb b/docs/general/user-guide/notebooks/numba_acceleration.ipynb
index cee91aab..08d78b64 100644
--- a/docs/general/user-guide/notebooks/numba_acceleration.ipynb
+++ b/docs/general/user-guide/notebooks/numba_acceleration.ipynb
@@ -1,326 +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",
-        "class MyAgentSet(AgentSetPolars):\n",
-        "    def __init__(self, model: ModelDF, n_agents: int):\n",
-        "        super().__init__(model)\n",
-        "        # Initialize agents\n",
-        "        self += pl.DataFrame({\n",
-        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
-        "            \"value\": pl.ones(n_agents, eager=True)\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",
-        "class MyVectorizedAgentSet(AgentSetPolars):\n",
-        "    def __init__(self, model: ModelDF, n_agents: int):\n",
-        "        super().__init__(model)\n",
-        "        # Initialize agents\n",
-        "        self += pl.DataFrame({\n",
-        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
-        "            \"value\": pl.ones(n_agents, eager=True)\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",
-        "class MyCudaAgentSet(AgentSetPolars):\n",
-        "    def __init__(self, model: ModelDF, n_agents: int):\n",
-        "        super().__init__(model)\n",
-        "        # Initialize agents\n",
-        "        self += pl.DataFrame({\n",
-        "            \"unique_id\": pl.arange(n_agents, eager=True),\n",
-        "            \"value\": pl.ones(n_agents, eager=True)\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",
-        "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(\n",
-        "            input1, input2, output\n",
-        "        ):\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"
-    }
+ "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."
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 4
+  {
+   "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
 }
\ No newline at end of file