Update space.py

Author: suryanshgargbpgc

mesa_frames/abstract/space.py

@@ -1061,13 +1061,21 @@ def move_to(
         randomized order to break ties. The final position of each agent is then
         updated in-place.
 
+        In practice, this method leverages Polars for efficient DataFrame operations.
+        For very large models or specialized data structures (like ``AgentContainer``),
+        consider referencing the minimal UDP approach in
+        ``space/utils/spatial_utils/grid_pairs.py``. You can adapt similar logic
+        (e.g., using lazy frames or partial computations) to maintain performance
+        at scale. If you observe performance degradation for large neighborhoods,
+        you may need to refactor or optimize further.
+
         Parameters
         ----------
         agents : AgentLike
             A DataFrame-like structure containing agent information. Must include
             at least the following columns:
-            - ``agent_id``: a unique identifier for each agent
-            - ``dim_0``, ``dim_1``: the current positions of agents
+            - ``unique_id``: a unique identifier for each agent
+            - ``dim_0``, ``dim_1``: the current positions of agents (in ``agents.pos``)
             - Optionally ``vision`` if ``radius`` is not provided
         attr_names : str or list of str
             The name(s) of the attribute(s) used for ranking the neighborhood cells.
@@ -1079,7 +1087,8 @@ def move_to(
             - ``"min"`` for ascending order
 
             If a single string is provided, it is applied to all attributes in
-            ``attr_names``.
+            ``attr_names``. **Note**: this method strongly assumes that the length
+            of ``attr_names`` matches the length of ``rank_order``.
         radius : int or pl.Series, optional
             The radius (or per-agent radii) defining the neighborhood around agents.
             If not provided, this method attempts to use the ``vision`` column from
@@ -1095,7 +1104,52 @@ def move_to(
         -------
         None
             This method updates agent positions in-place based on the computed best moves.
+
+        Raises
+        ------
+        ValueError
+            If the lengths of ``attr_names`` and ``rank_order`` do not match, or if
+            ``radius`` is not provided and the ``agents`` DataFrame does not contain
+            a ``vision`` attribute.
+
+        Notes
+        -----
+        - Type definitions (e.g., ``AgentLike``) are typically maintained in
+        ``mesa_frames.types``.
+        - For advanced usage or adapting this method to a custom data structure
+        (e.g., an ``AgentContainer``), you may reference the logic in
+        ``space/utils/spatial_utils/grid_pairs.py`` for a minimal UDP approach.
+        Using lazy frames or partial computations might help if performance
+        scales poorly with large data sets.
+
+        This method performs the following steps:
+        1. Compute the neighborhood for each agent using ``self.get_neighborhood``.
+        2. Join the neighborhood data with cell information (``self.cells``) and agent
+        positions (from ``agents.pos``).
+        3. Optionally shuffle the agent order to break ties.
+        4. Sort the neighborhood cells by the specified attribute(s) and order(s).
+        5. Iteratively select the best moves, ensuring no cell is claimed by multiple agents.
+        6. Call ``self.move_agents`` to finalize the updated positions of each agent.
+
+        Examples
+        --------
+        >>> # Assume we have a DataFrame 'agents' with columns:
+        >>> # ['unique_id', 'dim_0', 'dim_1', 'vision', 'food_availability', 'safety_score']
+        >>> # and a space object 'space' in a mesa-frames model.
+        >>>
+        >>> # We want to move each agent to the best available cell within its vision
+        >>> # radius, prioritizing cells with higher 'food_availability' and 'safety_score'.
+        >>> space.move_to(
+        ...     agents=agents,
+        ...     attr_names=["food_availability", "safety_score"],
+        ...     rank_order=["max", "max"],  # rank both attributes in descending order
+        ...     radius=None,               # use each agent's 'vision' column
+        ...     include_center=False,      # do not include the agent's current cell
+        ...     shuffle=True               # randomize the order in which agents move
+        ... )
+        >>> # After this call, each agent's position in 'agents' will be updated in-place.
         """
+
         # Ensure attr_names and rank_order are lists of the same length
         if isinstance(attr_names, str):
             attr_names = [attr_names]
@@ -1121,18 +1175,13 @@ def move_to(
         )
         neighborhood = neighborhood.join(self.cells, on=["dim_0", "dim_1"])
 
-        # Determine the agent identifier column
-        agent_id_col = "agent_id" if "agent_id" in agents.columns else "unique_id"
-        
-        # Add a column to identify the center agent
-        join_result = neighborhood.join(
-            agents.select(["dim_0", "dim_1", agent_id_col]),
-            left_on=["dim_0_center", "dim_1_center"],
-            right_on=["dim_0", "dim_1"]
-        )
-        
+        # Add a column to identify the center agent (the one evaluating moves)
         neighborhood = neighborhood.with_columns(
-            agent_id_center=join_result[agent_id_col]
+            agent_id_center=neighborhood.join(
+                agents.pos,
+                left_on=["dim_0_center", "dim_1_center"],
+                right_on=["dim_0", "dim_1"],
+            )["unique_id"]
         )
 
         # Determine the processing order of agents
@@ -1180,43 +1229,30 @@ def move_to(
 
         # Iteratively select the best moves
         best_moves = pl.DataFrame()
-        max_iterations = min(len(agents) * 2, 1000)  # Safeguard against infinite loops
-        iteration_count = 0
-        
-        while len(best_moves) < len(agents) and iteration_count < max_iterations:
-            iteration_count += 1
-            
+        while len(best_moves) < len(agents):
             # Count how many times each (dim_0, dim_1) is being claimed
             neighborhood = neighborhood.with_columns(
                 priority=pl.col("agent_order").cum_count().over(["dim_0", "dim_1"])
             )
-            
             new_best_moves = (
                 neighborhood.group_by("agent_id_center", maintain_order=True)
                 .first()
                 .unique(subset=["dim_0", "dim_1"], keep="first", maintain_order=True)
             )
-            
             condition = (
                 pl.col("blocking_agent_id").is_null()
                 | (pl.col("blocking_agent_id") == pl.col("agent_id_center"))
             )
-            
             if len(best_moves) > 0:
                 condition = condition | pl.col("blocking_agent_id").is_in(
                     best_moves["agent_id_center"]
                 )
-                
             condition = condition & (pl.col("priority") == 1)
             new_best_moves = new_best_moves.filter(condition)
-            
             if len(new_best_moves) == 0:
                 break
-                
+
             best_moves = pl.concat([best_moves, new_best_moves])
-            
-            # Update neighborhood to exclude agents that already have a move
-            # and cells that are already claimed
             neighborhood = neighborhood.filter(
                 ~pl.col("agent_id_center").is_in(best_moves["agent_id_center"])
             )
@@ -1226,20 +1262,11 @@ def move_to(
 
         # Move agents to their new positions
         if len(best_moves) > 0:
-            try:
-                self.move_agents(
-                    best_moves.sort("agent_order")["agent_id_center"],
-                    best_moves.sort("agent_order").select(["dim_0", "dim_1"])
-                )
-            except Exception as e:
-                # Check if the agent exists in the model
-                available_agents = set(self.model.agents[agent_id_col].to_list()) if hasattr(self.model, 'agents') else set()
-                missing_agents = [a for a in best_moves["agent_id_center"].to_list() if a not in available_agents]
-                
-                if missing_agents and available_agents:
-                    raise ValueError(f"Some agents are not present in the model: {missing_agents}")
-                else:
-                    raise ValueError(f"Error moving agents: {e}")
+            self.move_agents(
+                best_moves.sort("agent_order")["agent_id_center"],
+                best_moves.sort("agent_order").select(["dim_0", "dim_1"])
+            )
+