Skip to content

Commit

Permalink
space: Let move_agent choose from multiple positions
Browse files Browse the repository at this point in the history
- 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.
  • Loading branch information
EwoutH committed Dec 21, 2023
1 parent 9495a5a commit 73179f3
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 7 deletions.
56 changes: 49 additions & 7 deletions mesa/space.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down
61 changes: 61 additions & 0 deletions tests/test_space.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down

0 comments on commit 73179f3

Please sign in to comment.