Ehancement: Add names to AgentSetDF and AgentSetsAccessor for name/type/index lookup and container-aware keying

AGENTS.md

@@ -13,7 +13,7 @@
 
 - Install (dev stack): `uv sync` (always use uv)
 - Lint & format: `uv run ruff check . --fix && uv run ruff format .`
-- Tests (quiet + coverage): `export MESA_FRAMES_RUNTIME_TYPECHECKING = 1 && uv run pytest -q --cov=mesa_frames --cov-report=term-missing`
+- Tests (quiet + coverage): `export MESA_FRAMES_RUNTIME_TYPECHECKING=1 && uv run pytest -q --cov=mesa_frames --cov-report=term-missing`
 - Pre-commit (all files): `uv run pre-commit run -a`
 - Docs preview: `uv run mkdocs serve`
 
@@ -36,7 +36,7 @@ Always run tools via uv: `uv run <command>`.
 ## Commit & Pull Request Guidelines
 
 - Commits: Imperative mood, concise subject, meaningful body when needed.
-  Example: `Fix AgentsDF.sets copy binding and tests`.
+  Example: `Fix AgentSetRegistry.sets copy binding and tests`.
 - PRs: Link issues, summarize changes, note API impacts, add/adjust tests and docs.
 - CI hygiene: Run `ruff`, `pytest`, and `pre-commit` locally before pushing.
 

README.md

@@ -88,13 +88,13 @@ pip install -e .
 
 ### Creation of an Agent
 
-The agent implementation differs from base mesa. Agents are only defined at the AgentSet level. You can import `AgentSetPolars`. As in mesa, you subclass and make sure to call `super().__init__(model)`. You can use the `add` method or the `+=` operator to add agents to the AgentSet. Most methods mirror the functionality of `mesa.AgentSet`. Additionally, `mesa-frames.AgentSet` implements many dunder methods such as `AgentSet[mask, attr]` to get and set items intuitively. All operations are by default inplace, but if you'd like to use functional programming, mesa-frames implements a fast copy method which aims to reduce memory usage, relying on reference-only and native copy methods.
+The agent implementation differs from base mesa. Agents are only defined at the AgentSet level. You can import `AgentSet`. As in mesa, you subclass and make sure to call `super().__init__(model)`. You can use the `add` method or the `+=` operator to add agents to the AgentSet. Most methods mirror the functionality of `mesa.AgentSet`. Additionally, `mesa-frames.AgentSet` implements many dunder methods such as `AgentSet[mask, attr]` to get and set items intuitively. All operations are by default inplace, but if you'd like to use functional programming, mesa-frames implements a fast copy method which aims to reduce memory usage, relying on reference-only and native copy methods.
 
 ```python
-from mesa-frames import AgentSetPolars
+from mesa-frames import AgentSet
 
-class MoneyAgentPolars(AgentSetPolars):
-    def __init__(self, n: int, model: ModelDF):
+class MoneyAgentDF(AgentSet):
+    def __init__(self, n: int, model: Model):
         super().__init__(model)
         # Adding the agents to the agent set
         self += pl.DataFrame(
@@ -126,20 +126,20 @@ class MoneyAgentPolars(AgentSetPolars):
 
 ### Creation of the Model
 
-Creation of the model is fairly similar to the process in mesa. You subclass `ModelDF` and call `super().__init__()`. The `model.agents` attribute has the same interface as `mesa-frames.AgentSet`. You can use `+=` or `self.agents.add` with a `mesa-frames.AgentSet` (or a list of `AgentSet`) to add agents to the model.
+Creation of the model is fairly similar to the process in mesa. You subclass `Model` and call `super().__init__()`. The `model.sets` attribute has the same interface as `mesa-frames.AgentSet`. You can use `+=` or `self.sets.add` with a `mesa-frames.AgentSet` (or a list of `AgentSet`) to add agents to the model.
 
 ```python
-from mesa-frames import ModelDF
+from mesa-frames import Model
 
-class MoneyModelDF(ModelDF):
+class MoneyModelDF(Model):
     def __init__(self, N: int, agents_cls):
         super().__init__()
         self.n_agents = N
-        self.agents += MoneyAgentPolars(N, self)
+        self.sets += MoneyAgentDF(N, self)
 
     def step(self):
-        # Executes the step method for every agentset in self.agents
-        self.agents.do("step")
+        # Executes the step method for every agentset in self.sets
+        self.sets.do("step")
 
     def run_model(self, n):
         for _ in range(n):

ROADMAP.md

@@ -49,7 +49,7 @@ The Sugarscape example demonstrates the need for this abstraction, as multiple a
 
 #### Progress and Next Steps
 
-- Create utility functions in `DiscreteSpaceDF` and `AgentContainer` to move agents optimally based on specified attributes
+- Create utility functions in `AbstractDiscreteSpace` and `AbstractAgentSetRegistry` to move agents optimally based on specified attributes
 - Provide built-in resolution strategies for common concurrency scenarios
 - Ensure the implementation works efficiently with the vectorized approach of mesa-frames
 

docs/api/conf.py

@@ -4,14 +4,15 @@
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here.
 import sys
+from datetime import datetime
 from pathlib import Path
 
 sys.path.insert(0, str(Path("..").resolve()))
 
 # -- Project information -----------------------------------------------------
 project = "mesa-frames"
 author = "Project Mesa, Adam Amer"
-copyright = f"2023, {author}"
+copyright = f"{datetime.now().year}, {author}"
 
 # -- General configuration ---------------------------------------------------
 extensions = [

docs/api/reference/agents/index.rst

@@ -4,14 +4,14 @@ Agents
 .. currentmodule:: mesa_frames
 
 
-.. autoclass:: AgentSetPolars
+.. autoclass:: AgentSet
     :members:
     :inherited-members:
     :autosummary:
     :autosummary-nosignatures:
 
-.. autoclass:: AgentsDF
+.. autoclass:: AgentSetRegistry
     :members:
     :inherited-members:
     :autosummary:
-    :autosummary-nosignatures:
+    :autosummary-nosignatures:

docs/api/reference/model.rst

@@ -3,7 +3,7 @@ Model
 
 .. currentmodule:: mesa_frames
 
-.. autoclass:: ModelDF
+.. autoclass:: Model
     :members:
     :inherited-members:
     :autosummary:

docs/api/reference/space/index.rst

@@ -4,7 +4,7 @@ This page provides a high-level overview of possible space objects for mesa-fram
 
 .. currentmodule:: mesa_frames
 
-.. autoclass:: GridPolars
+.. autoclass:: Grid
     :members:
     :inherited-members:
     :autosummary:

docs/general/index.md

@@ -41,11 +41,11 @@ pip install -e .
 Here's a quick example of how to create a model using mesa-frames:
 
 ```python
-from mesa_frames import AgentSetPolars, ModelDF
+from mesa_frames import AgentSet, Model
 import polars as pl
 
-class MoneyAgentPolars(AgentSetPolars):
-    def __init__(self, n: int, model: ModelDF):
+class MoneyAgents(AgentSet):
+    def __init__(self, n: int, model: Model):
         super().__init__(model)
         self += pl.DataFrame(
             {"wealth": pl.ones(n, eager=True)}
@@ -57,13 +57,13 @@ class MoneyAgentPolars(AgentSetPolars):
     def give_money(self):
         # ... (implementation details)
 
-class MoneyModelDF(ModelDF):
+class MoneyModel(Model):
     def __init__(self, N: int):
         super().__init__()
-        self.agents += MoneyAgentPolars(N, self)
+        self.sets += MoneyAgentDF(N, self)
 
     def step(self):
-        self.agents.do("step")
+        self.sets.do("step")
 
     def run_model(self, n):
         for _ in range(n):

docs/general/user-guide/0_getting-started.md

@@ -35,14 +35,14 @@ Here's a comparison between mesa-frames and mesa:
 === "mesa-frames"
 
     ```python
-    class MoneyAgentPolarsConcise(AgentSetPolars):
+    class MoneyAgents(AgentSet):
         # initialization...
         def give_money(self):
             # Active agents are changed to wealthy agents
             self.select(self.wealth > 0)
 
             # Receiving agents are sampled (only native expressions currently supported)
-            other_agents = self.agents.sample(
+            other_agents = self.sets.sample(
                 n=len(self.active_agents), with_replacement=True
             )
 
@@ -64,7 +64,7 @@ Here's a comparison between mesa-frames and mesa:
         def give_money(self):
             # Verify agent has some wealth
             if self.wealth > 0:
-                other_agent = self.random.choice(self.model.agents)
+                other_agent = self.random.choice(self.model.sets)
                 if other_agent is not None:
                     other_agent.wealth += 1
                     self.wealth -= 1
@@ -84,15 +84,15 @@ If you're familiar with mesa, this guide will help you understand the key differ
 === "mesa-frames"
 
     ```python
-    class MoneyAgentSet(AgentSetPolars):
+    class MoneyAgents(AgentSet):
         def __init__(self, n, model):
             super().__init__(model)
             self += pl.DataFrame({
                 "wealth": pl.ones(n)
                 })
         def step(self):
             givers = self.wealth > 0
-            receivers = self.agents.sample(n=len(self.active_agents))
+            receivers = self.sets.sample(n=len(self.active_agents))
             self[givers, "wealth"] -= 1
             new_wealth = receivers.groupby("unique_id").count()
             self[new_wealth["unique_id"], "wealth"] += new_wealth["count"]
@@ -121,13 +121,13 @@ If you're familiar with mesa, this guide will help you understand the key differ
 === "mesa-frames"
 
     ```python
-    class MoneyModel(ModelDF):
+    class MoneyModel(Model):
         def __init__(self, N):
             super().__init__()
-            self.agents += MoneyAgentSet(N, self)
+            self.sets += MoneyAgents(N, self)
 
         def step(self):
-            self.agents.do("step")
+            self.sets.do("step")
 
     ```
 

docs/general/user-guide/1_classes.md

@@ -1,18 +1,18 @@
 # Classes 📚
 
-## AgentSetDF 👥
+## AgentSet 👥
 
-To create your own AgentSetDF class, you need to subclass the AgentSetPolars class and make sure to call `super().__init__(model)`.
+To create your own AgentSet class, you need to subclass the AgentSet class and make sure to call `super().__init__(model)`.
 
-Typically, the next step would be to populate the class with your agents. To do that, you need to add a DataFrame to the AgentSetDF. You can do `self += agents` or `self.add(agents)`, where `agents` is a DataFrame or something that could be passed to a DataFrame constructor, like a dictionary or lists of lists. You need to make sure your DataFrame doesn't have a 'unique_id' column because IDs are generated automatically, otherwise you will get an error raised. In the DataFrame, you should also put any attribute of the agent you are using.
+Typically, the next step would be to populate the class with your agents. To do that, you need to add a DataFrame to the AgentSet. You can do `self += agents` or `self.add(agents)`, where `agents` is a DataFrame or something that could be passed to a DataFrame constructor, like a dictionary or lists of lists. You need to make sure your DataFrame doesn't have a 'unique_id' column because IDs are generated automatically, otherwise you will get an error raised. In the DataFrame, you should also put any attribute of the agent you are using.
 
 How can you choose which agents should be in the same AgentSet? The idea is that you should minimize the missing values in the DataFrame (so they should have similar/same attributes) and mostly everybody should do the same actions.
 
 Example:
 
 ```python
-class MoneyAgent(AgentSetPolars):
-    def __init__(self, n: int, model: ModelDF):
+class MoneyAgents(AgentSet):
+    def __init__(self, n: int, model: Model):
         super().__init__(model)
         self.initial_wealth = pl.ones(n)
         self += pl.DataFrame({
@@ -25,28 +25,28 @@ class MoneyAgent(AgentSetPolars):
 
 You can access the underlying DataFrame where agents are stored with `self.df`. This allows you to use DataFrame methods like `self.df.sample` or `self.df.group_by("wealth")` and more.
 
-## ModelDF 🏗️
+## Model 🏗️
 
-To add your AgentSetDF to your ModelDF, you should also add it to the agents with `+=` or `add`.
+To add your AgentSet to your Model, you should also add it to the sets with `+=` or `add`.
 
-NOTE: ModelDF.agents are stored in a class which is entirely similar to AgentSetDF called AgentsDF. The API of the two are the same. If you try accessing AgentsDF.df, you will get a dictionary of `[AgentSetDF, DataFrame]`.
+NOTE: Model.sets are stored in a class which is entirely similar to AgentSet called AgentSetRegistry. The API of the two are the same. If you try accessing AgentSetRegistry.df, you will get a dictionary of `[AgentSet, DataFrame]`.
 
 Example:
 
 ```python
-class EcosystemModel(ModelDF):
+class EcosystemModel(Model):
     def __init__(self, n_prey, n_predators):
         super().__init__()
-        self.agents += Preys(n_prey, self)
-        self.agents += Predators(n_predators, self)
+        self.sets += Preys(n_prey, self)
+        self.sets += Predators(n_predators, self)
 
     def step(self):
-        self.agents.do("move")
-        self.agents.do("hunt")
+        self.sets.do("move")
+        self.sets.do("hunt")
         self.prey.do("reproduce")
 ```
 
-## Space: GridDF 🌐
+## Space: Grid 🌐
 
 mesa-frames provides efficient implementations of spatial environments:
 
@@ -55,12 +55,12 @@ mesa-frames provides efficient implementations of spatial environments:
 Example:
 
 ```python
-class GridWorld(ModelDF):
+class GridWorld(Model):
     def __init__(self, width, height):
         super().__init__()
-        self.space = GridPolars(self, (width, height))
-        self.agents += AgentSet(100, self)
-        self.space.place_to_empty(self.agents)
+        self.space = Grid(self, (width, height))
+        self.sets += AgentSet(100, self)
+        self.space.place_to_empty(self.sets)
 ```
 
 A continuous GeoSpace, NetworkSpace, and a collection to have multiple spaces in the models are in the works! 🚧
@@ -73,10 +73,10 @@ You configure what to collect, how to store it, and when to trigger collection.
 Example:
 
 ```python
-class ExampleModel(ModelDF):
+class ExampleModel(Model):
     def __init__(self):
         super().__init__()
-        self.agents = MoneyAgent(self)
+        self.sets = MoneyAgent(self)
         self.datacollector = DataCollector(
             model=self,
             model_reporters={"total_wealth": lambda m: m.agents["wealth"].sum()},
@@ -87,7 +87,7 @@ class ExampleModel(ModelDF):
         )
 
     def step(self):
-        self.agents.step()
+        self.sets.step()
         self.datacollector.conditional_collect()
         self.datacollector.flush()
 ```