Ehancement: Named AgentSet + flexible lookup & stricter separation of concerns

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

@@ -27,9 +27,9 @@ You can access the underlying DataFrame where agents are stored with `self.df`.
 
 ## Model 🏗️
 
-To add your AgentSet to your Model, you should also add it to the sets with `+=` or `add`.
+To add your AgentSet to your Model, use the registry `self.sets` with `+=` or `add`.
 
-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]`.
+Note: All agent sets live inside `AgentSetRegistry` (available as `model.sets`). Access sets through the registry, and access DataFrames from the set itself. For example: `self.sets["Preys"].df`.
 
 Example:
 
@@ -43,7 +43,8 @@ class EcosystemModel(Model):
     def step(self):
         self.sets.do("move")
         self.sets.do("hunt")
-        self.prey.do("reproduce")
+        # Access specific sets via the registry
+        self.sets["Preys"].do("reproduce")
 ```
 
 ## Space: Grid 🌐
@@ -76,18 +77,23 @@ Example:
 class ExampleModel(Model):
     def __init__(self):
         super().__init__()
-        self.sets = MoneyAgent(self)
+        # Add the set to the registry
+        self.sets.add(MoneyAgents(100, self))
+        # Configure reporters: use the registry to locate sets; get df from the set
         self.datacollector = DataCollector(
             model=self,
-            model_reporters={"total_wealth": lambda m: lambda m: list(m.sets.df.values())[0]["wealth"].sum()},
+            model_reporters={
+                "total_wealth": lambda m: m.sets["MoneyAgents"].df["wealth"].sum(),
+            },
             agent_reporters={"wealth": "wealth"},
             storage="csv",
             storage_uri="./data",
-            trigger=lambda m: m.schedule.steps % 2 == 0
+            trigger=lambda m: m.steps % 2 == 0,
         )
 
     def step(self):
-        self.sets.step()
+        # Step all sets via the registry
+        self.sets.do("step")
         self.datacollector.conditional_collect()
         self.datacollector.flush()
 ```

docs/general/user-guide/2_introductory-tutorial.ipynb

@@ -74,7 +74,9 @@
     "        self.sets += agents_cls(N, self)\n",
     "        self.datacollector = DataCollector(\n",
     "            model=self,\n",
-    "            model_reporters={\"total_wealth\": lambda m: m.agents[\"wealth\"].sum()},\n",
+    "            model_reporters={\n",
+    "                \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum()\n",
+    "            },\n",
     "            agent_reporters={\"wealth\": \"wealth\"},\n",
     "            storage=\"csv\",\n",
     "            storage_uri=\"./data\",\n",

docs/general/user-guide/4_datacollector.ipynb

@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 6,
    "id": "9a63283cbaf04dbcab1f6479b197f3a8",
    "metadata": {
     "editable": true
@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "72eea5119410473aa328ad9291626812",
    "metadata": {
     "editable": true
@@ -63,33 +63,33 @@
        " │ ---  ┆ ---                             ┆ ---   ┆ ---          ┆ ---      │\n",
        " │ i64  ┆ str                             ┆ i64   ┆ f64          ┆ i64      │\n",
        " ╞══════╪═════════════════════════════════╪═══════╪══════════════╪══════════╡\n",
-       " │ 2    ┆ 162681765859364298619846106603… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
-       " │ 4    ┆ 162681765859364298619846106603… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
-       " │ 6    ┆ 162681765859364298619846106603… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
-       " │ 8    ┆ 162681765859364298619846106603… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
-       " │ 10   ┆ 162681765859364298619846106603… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
+       " │ 2    ┆ 332212815818606584686857770936… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
+       " │ 4    ┆ 332212815818606584686857770936… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
+       " │ 6    ┆ 332212815818606584686857770936… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
+       " │ 8    ┆ 332212815818606584686857770936… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
+       " │ 10   ┆ 332212815818606584686857770936… ┆ 0     ┆ 1000.0       ┆ 1000     │\n",
        " └──────┴─────────────────────────────────┴───────┴──────────────┴──────────┘,\n",
        " 'agent': shape: (5_000, 4)\n",
        " ┌────────────────────┬──────┬─────────────────────────────────┬───────┐\n",
        " │ wealth_MoneyAgents ┆ step ┆ seed                            ┆ batch │\n",
        " │ ---                ┆ ---  ┆ ---                             ┆ ---   │\n",
        " │ f64                ┆ i32  ┆ str                             ┆ i32   │\n",
        " ╞════════════════════╪══════╪═════════════════════════════════╪═══════╡\n",
-       " │ 0.0                ┆ 2    ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 3.0                ┆ 2    ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 1.0                ┆ 2    ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 3.0                ┆ 2    ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 6.0                ┆ 2    ┆ 162681765859364298619846106603… ┆ 0     │\n",
+       " │ 3.0                ┆ 2    ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 2    ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 2.0                ┆ 2    ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 1.0                ┆ 2    ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 2    ┆ 332212815818606584686857770936… ┆ 0     │\n",
        " │ …                  ┆ …    ┆ …                               ┆ …     │\n",
-       " │ 4.0                ┆ 10   ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 1.0                ┆ 10   ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 0.0                ┆ 10   ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 0.0                ┆ 10   ┆ 162681765859364298619846106603… ┆ 0     │\n",
-       " │ 0.0                ┆ 10   ┆ 162681765859364298619846106603… ┆ 0     │\n",
+       " │ 0.0                ┆ 10   ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 10   ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 10   ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 10   ┆ 332212815818606584686857770936… ┆ 0     │\n",
+       " │ 0.0                ┆ 10   ┆ 332212815818606584686857770936… ┆ 0     │\n",
        " └────────────────────┴──────┴─────────────────────────────────┴───────┘}"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -120,8 +120,8 @@
     "        self.dc = DataCollector(\n",
     "            model=self,\n",
     "            model_reporters={\n",
-    "                \"total_wealth\": lambda m: list(m.sets.df.values())[0][\"wealth\"].sum(),\n",
-    "                \"n_agents\": lambda m: len(list(m.sets.df.values())[0]),\n",
+    "                \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n",
+    "                \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n",
     "            },\n",
     "            agent_reporters={\n",
     "                \"wealth\": \"wealth\",  # pull existing column\n",
@@ -175,7 +175,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": null,
    "id": "5f14f38c",
    "metadata": {},
    "outputs": [
@@ -185,7 +185,7 @@
        "[]"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -198,8 +198,8 @@
     "model_csv.dc = DataCollector(\n",
     "    model=model_csv,\n",
     "    model_reporters={\n",
-    "        \"total_wealth\": lambda m: list(m.sets.df.values())[0][\"wealth\"].sum(),\n",
-    "        \"n_agents\": lambda m: len(list(m.sets.df.values())[0]),\n",
+    "        \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n",
+    "        \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n",
     "    },\n",
     "    agent_reporters={\n",
     "        \"wealth\": \"wealth\",\n",
@@ -226,7 +226,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "id": "8763a12b2bbd4a93a75aff182afb95dc",
    "metadata": {
     "editable": true
@@ -238,7 +238,7 @@
        "[]"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -249,8 +249,8 @@
     "model_parq.dc = DataCollector(\n",
     "    model=model_parq,\n",
     "    model_reporters={\n",
-    "        \"total_wealth\": lambda m: list(m.sets.df.values())[0][\"wealth\"].sum(),\n",
-    "        \"n_agents\": lambda m: len(list(m.sets.df.values())[0]),\n",
+    "        \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n",
+    "        \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n",
     "    },\n",
     "    agent_reporters={\n",
     "        \"wealth\": \"wealth\",\n",
@@ -279,7 +279,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "id": "7cdc8c89c7104fffa095e18ddfef8986",
    "metadata": {
     "editable": true
@@ -290,8 +290,8 @@
     "model_s3.dc = DataCollector(\n",
     "    model=model_s3,\n",
     "    model_reporters={\n",
-    "        \"total_wealth\": lambda m: list(m.sets.df.values())[0][\"wealth\"].sum(),\n",
-    "        \"n_agents\": lambda m: len(list(m.sets.df.values())[0]),\n",
+    "        \"total_wealth\": lambda m: m.sets[\"MoneyAgents\"].df[\"wealth\"].sum(),\n",
+    "        \"n_agents\": lambda m: len(m.sets[\"MoneyAgents\"]),\n",
     "    },\n",
     "    agent_reporters={\n",
     "        \"wealth\": \"wealth\",\n",
@@ -319,7 +319,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 11,
    "id": "938c804e27f84196a10c8828c723f798",
    "metadata": {
     "editable": true
@@ -381,7 +381,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 12,
    "id": "59bbdb311c014d738909a11f9e486628",
    "metadata": {
     "editable": true
@@ -410,7 +410,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 13,
    "id": "8a65eabff63a45729fe45fb5ade58bdc",
    "metadata": {
     "editable": true
@@ -426,7 +426,7 @@
        "  white-space: pre-wrap;\n",
        "}\n",
        "</style>\n",
-       "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>step</th><th>seed</th><th>batch</th><th>total_wealth</th><th>n_agents</th></tr><tr><td>i64</td><td>str</td><td>i64</td><td>f64</td><td>i64</td></tr></thead><tbody><tr><td>2</td><td>&quot;732054881101029867447298951813…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>4</td><td>&quot;732054881101029867447298951813…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>6</td><td>&quot;732054881101029867447298951813…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>8</td><td>&quot;732054881101029867447298951813…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>10</td><td>&quot;732054881101029867447298951813…</td><td>0</td><td>100.0</td><td>100</td></tr></tbody></table></div>"
+       "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>step</th><th>seed</th><th>batch</th><th>total_wealth</th><th>n_agents</th></tr><tr><td>i64</td><td>str</td><td>i64</td><td>f64</td><td>i64</td></tr></thead><tbody><tr><td>2</td><td>&quot;540832786058427425452319829502…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>4</td><td>&quot;540832786058427425452319829502…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>6</td><td>&quot;540832786058427425452319829502…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>8</td><td>&quot;540832786058427425452319829502…</td><td>0</td><td>100.0</td><td>100</td></tr><tr><td>10</td><td>&quot;540832786058427425452319829502…</td><td>0</td><td>100.0</td><td>100</td></tr></tbody></table></div>"
       ],
       "text/plain": [
        "shape: (5, 5)\n",
@@ -435,15 +435,15 @@
        "│ ---  ┆ ---                             ┆ ---   ┆ ---          ┆ ---      │\n",
        "│ i64  ┆ str                             ┆ i64   ┆ f64          ┆ i64      │\n",
        "╞══════╪═════════════════════════════════╪═══════╪══════════════╪══════════╡\n",
-       "│ 2    ┆ 732054881101029867447298951813… ┆ 0     ┆ 100.0        ┆ 100      │\n",
-       "│ 4    ┆ 732054881101029867447298951813… ┆ 0     ┆ 100.0        ┆ 100      │\n",
-       "│ 6    ┆ 732054881101029867447298951813… ┆ 0     ┆ 100.0        ┆ 100      │\n",
-       "│ 8    ┆ 732054881101029867447298951813… ┆ 0     ┆ 100.0        ┆ 100      │\n",
-       "│ 10   ┆ 732054881101029867447298951813… ┆ 0     ┆ 100.0        ┆ 100      │\n",
+       "│ 2    ┆ 540832786058427425452319829502… ┆ 0     ┆ 100.0        ┆ 100      │\n",
+       "│ 4    ┆ 540832786058427425452319829502… ┆ 0     ┆ 100.0        ┆ 100      │\n",
+       "│ 6    ┆ 540832786058427425452319829502… ┆ 0     ┆ 100.0        ┆ 100      │\n",
+       "│ 8    ┆ 540832786058427425452319829502… ┆ 0     ┆ 100.0        ┆ 100      │\n",
+       "│ 10   ┆ 540832786058427425452319829502… ┆ 0     ┆ 100.0        ┆ 100      │\n",
        "└──────┴─────────────────────────────────┴───────┴──────────────┴──────────┘"
       ]
      },
-     "execution_count": 25,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }

examples/sugarscape_ig/ss_polars/agents.py

@@ -35,12 +35,15 @@ def __init__(
         self.add(agents)
 
     def eat(self):
+        # Only consider cells currently occupied by agents of this set
         cells = self.space.cells.filter(pl.col("agent_id").is_not_null())
-        self[cells["agent_id"], "sugar"] = (
-            self[cells["agent_id"], "sugar"]
-            + cells["sugar"]
-            - self[cells["agent_id"], "metabolism"]
-        )
+        mask_in_set = cells["agent_id"].is_in(self.index)
+        if mask_in_set.any():
+            cells = cells.filter(mask_in_set)
+            ids = cells["agent_id"]
+            self[ids, "sugar"] = (
+                self[ids, "sugar"] + cells["sugar"] - self[ids, "metabolism"]
+            )
 
     def step(self):
         self.shuffle().do("move").do("eat")

examples/sugarscape_ig/ss_polars/model.py

@@ -33,15 +33,19 @@ def __init__(
             sugar=sugar_grid.flatten(), max_sugar=sugar_grid.flatten()
         )
         self.space.set_cells(sugar_grid)
-        self.sets += agent_type(self, n_agents, initial_sugar, metabolism, vision)
+        # Create and register the main agent set; keep its name for later lookups
+        main_set = agent_type(self, n_agents, initial_sugar, metabolism, vision)
+        self.sets += main_set
+        self._main_set_name = main_set.name
         if initial_positions is not None:
             self.space.place_agents(self.sets, initial_positions)
         else:
             self.space.place_to_empty(self.sets)
 
     def run_model(self, steps: int) -> list[int]:
         for _ in range(steps):
-            if len(list(self.sets.df.values())[0]) == 0:
+            # Stop if the main agent set is empty
+            if len(self.sets[self._main_set_name]) == 0:  # type: ignore[index]
                 return
             empty_cells = self.space.empty_cells
             full_cells = self.space.full_cells