diff --git a/mesa/space.py b/mesa/space.py index 18cf73145e2..5f95e733c6e 100644 --- a/mesa/space.py +++ b/mesa/space.py @@ -419,17 +419,59 @@ def place_agent(self, agent: Agent, pos: Coordinate) -> None: def remove_agent(self, agent: Agent) -> None: ... - def move_agent(self, agent: Agent, pos: Coordinate) -> None: - """Move an agent from its current position to a new position. + def move_agent( + self, + agent: Agent, + pos: Coordinate | list[Coordinate], + selection: str = "random", + ) -> None: + """ + Move an agent from its current position to a new position. Args: - agent: Agent object to move. Assumed to have its current location - stored in a 'pos' tuple. - pos: Tuple of new position to move the agent to. + agent: Agent object to move. Assumed to have its current location stored in a 'pos' tuple. + pos: A single position or a list of possible positions. + selection: String, either "random" or "closest". If "closest" is selected and multiple + cells are the same distance, one is chosen randomly. """ - pos = self.torus_adj(pos) + # Handle single position case quickly + if isinstance(pos, tuple): + pos = self.torus_adj(pos) + self.remove_agent(agent) + self.place_agent(agent, pos) + return + + # Handle list of positions + if selection == 'random': + chosen_pos = agent.random.choice(pos) + elif selection == 'closest': + current_pos = agent.pos + # Find the closest position without sorting all positions + closest_pos = None + min_distance = float('inf') + for p in pos: + distance = self.distance_squared(p, current_pos) + if distance < min_distance: + min_distance = distance + closest_pos = p + chosen_pos = closest_pos + else: + raise ValueError(f"Invalid selection method {selection}. Choose 'random' or 'closest'.") + + chosen_pos = self.torus_adj(chosen_pos) self.remove_agent(agent) - self.place_agent(agent, pos) + self.place_agent(agent, chosen_pos) + + def distance_squared(self, pos1: Coordinate, pos2: Coordinate) -> float: + """ + Calculate the squared Euclidean distance between two points for performance. + """ + # Use squared Euclidean distance to avoid sqrt operation + dx, dy = abs(pos1[0] - pos2[0]), abs(pos1[1] - pos2[1]) + if self.torus: + dx = min(dx, self.width - dx) + dy = min(dy, self.height - dy) + return dx ** 2 + dy ** 2 def swap_pos(self, agent_a: Agent, agent_b: Agent) -> None: """Swap agents positions""" diff --git a/tests/test_space.py b/tests/test_space.py index f8f2cc9440c..5d94b722802 100644 --- a/tests/test_space.py +++ b/tests/test_space.py @@ -327,6 +327,31 @@ def move_agent(self): assert self.space[initial_pos[0]][initial_pos[1]] is None assert self.space[final_pos[0]][final_pos[1]] == _agent + def test_move_agent_random_selection(self): + agent = self.agents[0] + possible_positions = [(10, 10), (20, 20), (30, 30)] + self.space.move_agent(agent, possible_positions, selection="random") + assert agent.pos in possible_positions + + def test_move_agent_closest_selection(self): + agent = self.agents[0] + agent.pos = (5, 5) + possible_positions = [(6, 6), (10, 10), (20, 20)] + self.space.move_agent(agent, possible_positions, selection="closest") + assert agent.pos == (6, 6) + + def test_move_agent_invalid_selection(self): + agent = self.agents[0] + possible_positions = [(10, 10), (20, 20), (30, 30)] + with self.assertRaises(ValueError): + self.space.move_agent(agent, possible_positions, selection="invalid_option") + + def test_distance_squared(self): + pos1 = (3, 4) + pos2 = (0, 0) + expected_distance_squared = 3**2 + 4**2 + assert self.space.distance_squared(pos1, pos2) == expected_distance_squared + def test_iter_cell_list_contents(self): """ Test neighborhood retrieval @@ -350,6 +375,42 @@ def test_iter_cell_list_contents(self): assert len(cell_list_4) == 1 +class TestSingleGridTorus(unittest.TestCase): + def setUp(self): + self.space = SingleGrid(50, 50, True) # Torus is True here + self.agents = [] + for i, pos in enumerate(TEST_AGENTS_GRID): + a = MockAgent(i, None) + self.agents.append(a) + self.space.place_agent(a, pos) + + def test_move_agent_random_selection(self): + agent = self.agents[0] + possible_positions = [(49, 49), (1, 1), (25, 25)] + self.space.move_agent(agent, possible_positions, selection="random") + assert agent.pos in possible_positions + + def test_move_agent_closest_selection(self): + agent = self.agents[0] + agent.pos = (0, 0) + possible_positions = [(3, 3), (49, 49), (25, 25)] + self.space.move_agent(agent, possible_positions, selection="closest") + # Expecting (49, 49) to be the closest in a torus grid + assert agent.pos == (49, 49) + + def test_move_agent_invalid_selection(self): + agent = self.agents[0] + possible_positions = [(10, 10), (20, 20), (30, 30)] + with self.assertRaises(ValueError): + self.space.move_agent(agent, possible_positions, selection="invalid_option") + + def test_distance_squared_torus(self): + pos1 = (0, 0) + pos2 = (49, 49) + expected_distance_squared = 1**2 + 1**2 # In torus, these points are close + assert self.space.distance_squared(pos1, pos2) == expected_distance_squared + + class TestSingleNetworkGrid(unittest.TestCase): GRAPH_SIZE = 10