added docstrings

Author: suryanshgargbpgc

mesa_frames/abstract/space.py

@@ -59,7 +59,6 @@ def __init__(self, model, dimensions, torus, capacity, neighborhood_type):
 from typing_extensions import Any, Self
 
 
-from mesa_frames.concrete.polars.agentset import AgentSetPolars
 from mesa_frames.concrete.agents import AgentsDF
 from mesa_frames.abstract.agents import AgentContainer, AgentSetDF
 from mesa_frames.abstract.mixin import CopyMixin, DataFrameMixin
@@ -77,12 +76,13 @@ def __init__(self, model, dimensions, torus, capacity, neighborhood_type):
     Series,
     SpaceCoordinate,
     SpaceCoordinates,
+    AgentLike
 )
 
 ESPG = int
 
 
-AgentLike = Union[AgentSetPolars, pl.DataFrame]
+
 
 if TYPE_CHECKING:
     from mesa_frames.concrete.model import ModelDF
@@ -1050,36 +1050,98 @@ def move_to(
         include_center: bool = True,
         shuffle: bool = True
     ) -> None:
+        """
+        Move agents to new positions based on neighborhood ranking.
+
+        This method determines each agent's potential moves by computing their
+        local neighborhood (with optional inclusion of the center cell). For each
+        agent, the method ranks all possible moves according to specified attribute(s)
+        and rank order(s). If multiple agents contend for the same cell, the method
+        applies a tie-breaking approach. Agents can optionally be processed in a
+        randomized order to break ties. The final position of each agent is then
+        updated in-place.
+
+        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
+            - 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.
+            If multiple attributes are provided, each should have a corresponding
+            entry in ``rank_order``.
+        rank_order : str or list of str, optional
+            The ranking order for each attribute. Accepts:
+            - ``"max"`` (default) for descending order
+            - ``"min"`` for ascending order
+
+            If a single string is provided, it is applied to all attributes in
+            ``attr_names``.
+        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
+            ``agents``. If ``vision`` is not found, a ``ValueError`` is raised.
+        include_center : bool, optional
+            If ``True`` (default), the agent's current position is included in its
+            neighborhood.
+        shuffle : bool, optional
+            If ``True`` (default), the order of agents is randomized to break ties.
+            If ``False``, agents are processed in the order they appear in the data.
+
+        Returns
+        -------
+        None
+            This method updates agent positions in-place based on the computed best moves.
+        """
+        # Ensure attr_names and rank_order are lists of the same length
         if isinstance(attr_names, str):
             attr_names = [attr_names]
         if isinstance(rank_order, str):
             rank_order = [rank_order] * len(attr_names)
         if len(attr_names) != len(rank_order):
             raise ValueError("attr_names and rank_order must have the same length")
+
+        # Handle the neighborhood radius
         if radius is None:
             if "vision" in agents.columns:
                 radius = agents["vision"]
             else:
-                raise ValueError("radius must be specified if agents do not have a 'vision' attribute")
+                raise ValueError(
+                    "radius must be specified if agents do not have a 'vision' attribute"
+                )
+
+        # Get neighborhood and join with cell information
         neighborhood = self.get_neighborhood(
-            radius=radius, 
-            agents=agents, 
+            radius=radius,
+            agents=agents,
             include_center=include_center
         )
         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"]
+        )
+        
         neighborhood = neighborhood.with_columns(
-            agent_id_center=neighborhood.join(
-                agents.pos,
-                left_on=["dim_0_center", "dim_1_center"],
-                right_on=["dim_0", "dim_1"],
-            )["unique_id"]
+            agent_id_center=join_result[agent_id_col]
         )
+
+        # Determine the processing order of agents
         if shuffle:
             agent_order = (
                 neighborhood
                 .unique(subset=["agent_id_center"], keep="first")
                 .select("agent_id_center")
-                .sample(fraction=1.0, seed=self.model.random.integers(0, 2**31-1))
+                .sample(fraction=1.0, seed=self.model.random.integers(0, 2**31 - 1))
                 .with_row_index("agent_order")
             )
         else:
@@ -1089,16 +1151,24 @@ def move_to(
                 .with_row_index("agent_order")
                 .select(["agent_id_center", "agent_order"])
             )
+
+        # Join the processing order with the neighborhood
         neighborhood = neighborhood.join(agent_order, on="agent_id_center")
+
+        # Prepare sorting columns and order
         sort_cols = []
         sort_desc = []
         for attr, order in zip(attr_names, rank_order):
             sort_cols.append(attr)
             sort_desc.append(order.lower() == "max")
+
+        # Sort the neighborhood cells by specified attributes and then by location
         neighborhood = neighborhood.sort(
             sort_cols + ["radius", "dim_0", "dim_1"],
             descending=sort_desc + [False, False, False]
         )
+
+        # Join to track if another agent has blocked a cell
         neighborhood = neighborhood.join(
             agent_order.select(
                 pl.col("agent_id_center").alias("agent_id"),
@@ -1107,39 +1177,71 @@ def move_to(
             on="agent_id",
             how="left",
         ).rename({"agent_id": "blocking_agent_id"})
+
+        # Iteratively select the best moves
         best_moves = pl.DataFrame()
-        while len(best_moves) < len(agents):
+        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
+            
+            # 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")
+            
+            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"])
             )
             neighborhood = neighborhood.join(
                 best_moves.select(["dim_0", "dim_1"]), on=["dim_0", "dim_1"], how="anti"
             )
+
+        # Move agents to their new positions
         if len(best_moves) > 0:
-            self.move_agents(
-                best_moves.sort("agent_order")["agent_id_center"],
-                best_moves.sort("agent_order").select(["dim_0", "dim_1"])
-            )
+            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}")
+
+
 
 
     @property

mesa_frames/types_.py

@@ -1,14 +1,15 @@
 """Type aliases for the mesa_frames package."""
 
 from collections.abc import Collection, Sequence
-from typing import Literal
+from typing import Literal, Union
 
 # import geopandas as gpd
 # import geopolars as gpl
 import pandas as pd
 import polars as pl
 from numpy import ndarray
 from typing_extensions import Any
+from mesa_frames.concrete.polars.agentset import AgentSetPolars
 
 ####----- Agnostic Types -----####
 AgnosticMask = (
@@ -76,3 +77,4 @@
 NetworkCapacity = DataFrame
 
 DiscreteSpaceCapacity = GridCapacity | NetworkCapacity
+AgentLike = Union[AgentSetPolars, pl.DataFrame]