103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
from shapely.geometry import Polygon,Point
|
|
from typing import Tuple, List, Dict, Set
|
|
from scipy.spatial import cKDTree
|
|
import sys
|
|
import numpy as np
|
|
from agent_setup import AgentSetup
|
|
|
|
|
|
class SpawnManager:
|
|
def __init__(self,spawn_area:Polygon,min_spacing:float=0.55):
|
|
"""
|
|
self.spawn_area: geometry where agents may spawn
|
|
self.min_spacing: minimum spacing for spawn_points
|
|
self.spawn_coords: all spawn points available
|
|
self.filled_coords: spawn points currently filled by agents
|
|
self.spawned_agents: all agents already spawned in sim
|
|
self.agent_pos: connects agent_id to spawn_index
|
|
self.rng: random number generator object
|
|
"""
|
|
|
|
self.spawn_area = spawn_area
|
|
self.min_spacing = min_spacing
|
|
self.spawn_coords: np.ndarray = np.array([])
|
|
self.filled_coords: np.ndarray = np.array([])
|
|
self.spawned_agents:Set[int] = set()
|
|
self.agent_pos:Dict[int,int] = {}
|
|
self.rng = np.random.default_rng()
|
|
|
|
def generate_coords(self,max_samples:int=1000)->None:
|
|
buffered_spawn = self.spawn_area.buffer(-self.min_spacing)
|
|
min_x,min_y,max_x,max_y = buffered_spawn.bounds
|
|
points = []
|
|
while len(points) ==0:
|
|
x = self.rng.uniform(min_x,max_x)
|
|
y = self.rng.uniform(min_y,max_y)
|
|
if buffered_spawn.contains(Point(x,y)):
|
|
points.append([x,y])
|
|
for _ in range(max_samples):
|
|
idx = self.rng.integers(0,len(points))
|
|
base = points[idx]
|
|
for _ in range(25):
|
|
angle = self.rng.uniform(0,2*np.pi)
|
|
radius = self.rng.uniform(self.min_spacing,2*self.min_spacing)
|
|
x = base[0]+radius*np.cos(angle)
|
|
y = base[1]+radius*np.sin(angle)
|
|
if not buffered_spawn.contains(Point(x,y)):
|
|
continue
|
|
if len(points)>0:
|
|
tree = cKDTree(points)
|
|
distance,_ = tree.query([[x,y]],k=1)
|
|
if distance[0]<self.min_spacing:
|
|
continue
|
|
points.append([x,y])
|
|
break
|
|
self.spawn_coords = np.array(points)
|
|
self.filled_coords = np.zeros(len(points),dtype=bool)
|
|
return self.spawn_coords
|
|
|
|
def get_coords(self)->Tuple[float,float]|None:
|
|
free_idx = np.where(~self.filled_coords)[0]
|
|
if len(free_idx) == 0:
|
|
return None
|
|
idx = self.rng.choice(free_idx)
|
|
return tuple(self.spawn_coords[idx])
|
|
|
|
|
|
def spawn_agent(self,all_agents:List[AgentSetup])->AgentSetup|None:
|
|
if len(self.spawned_agents) >= len(all_agents):
|
|
return None
|
|
spawn_point = self.get_coords()
|
|
if not spawn_point:
|
|
return None
|
|
|
|
free_agents = [agent for agent in all_agents \
|
|
if agent.id not in self.spawned_agents
|
|
]
|
|
if not free_agents:
|
|
return None
|
|
agent = self.rng.choice(free_agents)
|
|
self.spawned_agents.add(agent.id)
|
|
distances = np.linalg.norm(self.spawn_coords-spawn_point,axis=1)
|
|
spawn_idx = np.argmin(distances)
|
|
self.filled_coords[spawn_idx] = True
|
|
self.agent_pos[agent.id] = spawn_idx
|
|
return agent
|
|
|
|
def unfill_coords(self,agent_id:int)->None:
|
|
if agent_id in self.agent_pos:
|
|
spawn_idx = self.agent_pos[agent_id]
|
|
self.filled_coords[spawn_idx] = False
|
|
del self.agent_pos[agent_id]
|
|
self.spawned_agents.discard(agent_id)
|
|
|
|
def get_agent_pos(self,agent_id:int)->Tuple[float,float]:
|
|
if agent_id in self.agent_pos:
|
|
spawn_idx = self.agent_pos[agent_id]
|
|
return tuple(self.spawn_coords[spawn_idx])
|
|
return (0,0)
|
|
|
|
def check_spawn_complete(self,all_agents:List[AgentSetup])->bool:
|
|
return (len(self.spawned_agents)<len(all_agents) and\
|
|
np.sum(~self.filled_coords) > 0)
|