First

_source_/spawning.py

@@ -0,0 +1,102 @@
+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)