Adds move_to_optimal function in DiscreteSpaceDF class

examples/sugarscape_ig/ss_polars/agents.py

@@ -272,6 +272,7 @@ def get_best_moves(self, neighborhood: pl.DataFrame):
         )
         return best_moves
 
+    # Resolved method with proper docstring
     def _prepare_cells(
         self, neighborhood: pl.DataFrame
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -286,7 +287,14 @@ def _prepare_cells(
         Returns
         -------
         tuple[np.ndarray, np.ndarray, np.ndarray]
+
+            A tuple containing:
+            - occupied_cells: Array of currently occupied cell positions
+            - free_cells: Boolean array indicating which cells are free
+            - target_cells: Array of target cell positions for each agent
+
             occupied_cells, free_cells, target_cells
+
         """
         occupied_cells = (
             neighborhood[["agent_id_center", "agent_order"]]

mesa_frames/abstract/agents.py

@@ -616,6 +616,19 @@
         """
         ...
 
+    @abstractmethod
+    def move_to_optimal(
+        self,
+        attr_names: str | list[str],
+        rank_order: str | list[str] = "max",
+        radius: int | Series | None = None,
+        include_center: bool = True,
+        shuffle: bool = True,
+        inplace: bool = True,
+    ) -> Self:
+        """Move agents to optimal cells based on neighborhood ranking."""
+        ...
+
     @property
     def model(self) -> ModelDF:
         """The model that the AgentContainer belongs to.
@@ -1039,6 +1052,31 @@
     def __reversed__(self) -> Iterator:
         return reversed(self._agents)
 
+    def move_to_optimal(
+        self,
+        attr_names: str | list[str],
+        rank_order: str | list[str] = "max",
+        radius: int | Series | None = None,
+        include_center: bool = True,
+        shuffle: bool = True,
+        inplace: bool = True,
+    ) -> Self:
+        """Move all agent sets to optimal cells based on neighborhood ranking."""
+        obj = self._get_obj(inplace)
+
+        # Apply move_to_optimal to each agent set in the container
+        for agent_set in obj.agent_sets.values():
+            agent_set.move_to_optimal(
+                attr_names=attr_names,
+                rank_order=rank_order,
+                radius=radius,
+                include_center=include_center,
+                shuffle=shuffle,
+                inplace=True,
+            )
+
+        return obj
+
     @property
     def agents(self) -> DataFrame:
         return self._agents