fix ss_polars

Author: adamamer20

examples/sugarscape_ig/ss_polars/agents.py

@@ -35,12 +35,13 @@ 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