Skip to content

Commit 73179f3

Browse files
committed
space: Let move_agent choose from multiple positions
- Add functionality to move_agent to handle a list of positions. - Implement selection criteria: 'random' for random selection and 'closest' for selecting the nearest position. - Include error handling for invalid selection methods. - Optimize distance calculations using squared Euclidean distance, considering toroidal grid adjustments. - Update tests in TestSingleGrid to cover new move_agent functionalities, including tests for random and closest selection and handling of invalid selection methods.
1 parent 9495a5a commit 73179f3

File tree

2 files changed

+110
-7
lines changed

2 files changed

+110
-7
lines changed

mesa/space.py

Lines changed: 49 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -419,17 +419,59 @@ def place_agent(self, agent: Agent, pos: Coordinate) -> None:
419419
def remove_agent(self, agent: Agent) -> None:
420420
...
421421

422-
def move_agent(self, agent: Agent, pos: Coordinate) -> None:
423-
"""Move an agent from its current position to a new position.
422+
def move_agent(
423+
self,
424+
agent: Agent,
425+
pos: Coordinate | list[Coordinate],
426+
selection: str = "random",
427+
) -> None:
428+
"""
429+
Move an agent from its current position to a new position.
424430
425431
Args:
426-
agent: Agent object to move. Assumed to have its current location
427-
stored in a 'pos' tuple.
428-
pos: Tuple of new position to move the agent to.
432+
agent: Agent object to move. Assumed to have its current location stored in a 'pos' tuple.
433+
pos: A single position or a list of possible positions.
434+
selection: String, either "random" or "closest". If "closest" is selected and multiple
435+
cells are the same distance, one is chosen randomly.
429436
"""
430-
pos = self.torus_adj(pos)
437+
# Handle single position case quickly
438+
if isinstance(pos, tuple):
439+
pos = self.torus_adj(pos)
440+
self.remove_agent(agent)
441+
self.place_agent(agent, pos)
442+
return
443+
444+
# Handle list of positions
445+
if selection == 'random':
446+
chosen_pos = agent.random.choice(pos)
447+
elif selection == 'closest':
448+
current_pos = agent.pos
449+
# Find the closest position without sorting all positions
450+
closest_pos = None
451+
min_distance = float('inf')
452+
for p in pos:
453+
distance = self.distance_squared(p, current_pos)
454+
if distance < min_distance:
455+
min_distance = distance
456+
closest_pos = p
457+
chosen_pos = closest_pos
458+
else:
459+
raise ValueError(f"Invalid selection method {selection}. Choose 'random' or 'closest'.")
460+
461+
chosen_pos = self.torus_adj(chosen_pos)
431462
self.remove_agent(agent)
432-
self.place_agent(agent, pos)
463+
self.place_agent(agent, chosen_pos)
464+
465+
def distance_squared(self, pos1: Coordinate, pos2: Coordinate) -> float:
466+
"""
467+
Calculate the squared Euclidean distance between two points for performance.
468+
"""
469+
# Use squared Euclidean distance to avoid sqrt operation
470+
dx, dy = abs(pos1[0] - pos2[0]), abs(pos1[1] - pos2[1])
471+
if self.torus:
472+
dx = min(dx, self.width - dx)
473+
dy = min(dy, self.height - dy)
474+
return dx ** 2 + dy ** 2
433475

434476
def swap_pos(self, agent_a: Agent, agent_b: Agent) -> None:
435477
"""Swap agents positions"""

tests/test_space.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,31 @@ def move_agent(self):
327327
assert self.space[initial_pos[0]][initial_pos[1]] is None
328328
assert self.space[final_pos[0]][final_pos[1]] == _agent
329329

330+
def test_move_agent_random_selection(self):
331+
agent = self.agents[0]
332+
possible_positions = [(10, 10), (20, 20), (30, 30)]
333+
self.space.move_agent(agent, possible_positions, selection="random")
334+
assert agent.pos in possible_positions
335+
336+
def test_move_agent_closest_selection(self):
337+
agent = self.agents[0]
338+
agent.pos = (5, 5)
339+
possible_positions = [(6, 6), (10, 10), (20, 20)]
340+
self.space.move_agent(agent, possible_positions, selection="closest")
341+
assert agent.pos == (6, 6)
342+
343+
def test_move_agent_invalid_selection(self):
344+
agent = self.agents[0]
345+
possible_positions = [(10, 10), (20, 20), (30, 30)]
346+
with self.assertRaises(ValueError):
347+
self.space.move_agent(agent, possible_positions, selection="invalid_option")
348+
349+
def test_distance_squared(self):
350+
pos1 = (3, 4)
351+
pos2 = (0, 0)
352+
expected_distance_squared = 3**2 + 4**2
353+
assert self.space.distance_squared(pos1, pos2) == expected_distance_squared
354+
330355
def test_iter_cell_list_contents(self):
331356
"""
332357
Test neighborhood retrieval
@@ -350,6 +375,42 @@ def test_iter_cell_list_contents(self):
350375
assert len(cell_list_4) == 1
351376

352377

378+
class TestSingleGridTorus(unittest.TestCase):
379+
def setUp(self):
380+
self.space = SingleGrid(50, 50, True) # Torus is True here
381+
self.agents = []
382+
for i, pos in enumerate(TEST_AGENTS_GRID):
383+
a = MockAgent(i, None)
384+
self.agents.append(a)
385+
self.space.place_agent(a, pos)
386+
387+
def test_move_agent_random_selection(self):
388+
agent = self.agents[0]
389+
possible_positions = [(49, 49), (1, 1), (25, 25)]
390+
self.space.move_agent(agent, possible_positions, selection="random")
391+
assert agent.pos in possible_positions
392+
393+
def test_move_agent_closest_selection(self):
394+
agent = self.agents[0]
395+
agent.pos = (0, 0)
396+
possible_positions = [(3, 3), (49, 49), (25, 25)]
397+
self.space.move_agent(agent, possible_positions, selection="closest")
398+
# Expecting (49, 49) to be the closest in a torus grid
399+
assert agent.pos == (49, 49)
400+
401+
def test_move_agent_invalid_selection(self):
402+
agent = self.agents[0]
403+
possible_positions = [(10, 10), (20, 20), (30, 30)]
404+
with self.assertRaises(ValueError):
405+
self.space.move_agent(agent, possible_positions, selection="invalid_option")
406+
407+
def test_distance_squared_torus(self):
408+
pos1 = (0, 0)
409+
pos2 = (49, 49)
410+
expected_distance_squared = 1**2 + 1**2 # In torus, these points are close
411+
assert self.space.distance_squared(pos1, pos2) == expected_distance_squared
412+
413+
353414
class TestSingleNetworkGrid(unittest.TestCase):
354415
GRAPH_SIZE = 10
355416

0 commit comments

Comments
 (0)