109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy
|
|
from scipy.spatial import cKDTree
|
|
|
|
if TYPE_CHECKING:
|
|
from shapely.geometry import Polygon
|
|
|
|
|
|
_COST_CACHE_SIZE = 100000
|
|
|
|
|
|
class DangerMap:
|
|
"""
|
|
A proximity cost evaluator using a KD-Tree of obstacle boundary points.
|
|
Scales with obstacle perimeter rather than design area.
|
|
"""
|
|
__slots__ = ('minx', 'miny', 'maxx', 'maxy', 'resolution', 'safety_threshold', 'k', 'tree', '_cost_cache')
|
|
|
|
def __init__(
|
|
self,
|
|
bounds: tuple[float, float, float, float],
|
|
resolution: float = 5.0,
|
|
safety_threshold: float = 10.0,
|
|
k: float = 1.0,
|
|
) -> None:
|
|
"""
|
|
Initialize the Danger Map.
|
|
|
|
Args:
|
|
bounds: (minx, miny, maxx, maxy) in um.
|
|
resolution: Sampling resolution for obstacle boundaries (um).
|
|
safety_threshold: Proximity limit (um).
|
|
k: Penalty multiplier.
|
|
"""
|
|
self.minx, self.miny, self.maxx, self.maxy = bounds
|
|
self.resolution = resolution
|
|
self.safety_threshold = safety_threshold
|
|
self.k = k
|
|
self.tree: cKDTree | None = None
|
|
self._cost_cache: OrderedDict[tuple[int, int], float] = OrderedDict()
|
|
|
|
def precompute(self, obstacles: list[Polygon]) -> None:
|
|
"""
|
|
Pre-compute the proximity tree by sampling obstacle boundaries.
|
|
"""
|
|
all_points = []
|
|
for poly in obstacles:
|
|
# Sample exterior
|
|
exterior = poly.exterior
|
|
dist = 0
|
|
while dist < exterior.length:
|
|
pt = exterior.interpolate(dist)
|
|
all_points.append((pt.x, pt.y))
|
|
dist += self.resolution
|
|
# Sample interiors (holes)
|
|
for interior in poly.interiors:
|
|
dist = 0
|
|
while dist < interior.length:
|
|
pt = interior.interpolate(dist)
|
|
all_points.append((pt.x, pt.y))
|
|
dist += self.resolution
|
|
|
|
if all_points:
|
|
self.tree = cKDTree(numpy.array(all_points))
|
|
else:
|
|
self.tree = None
|
|
|
|
self._cost_cache.clear()
|
|
|
|
def is_within_bounds(self, x: float, y: float) -> bool:
|
|
"""
|
|
Check if a coordinate is within the design bounds.
|
|
"""
|
|
return self.minx <= x <= self.maxx and self.miny <= y <= self.maxy
|
|
|
|
def get_cost(self, x: float, y: float) -> float:
|
|
"""
|
|
Get the proximity cost at a specific coordinate using the KD-Tree.
|
|
Coordinates are quantized to 1nm to improve cache performance.
|
|
"""
|
|
qx_milli = int(round(x * 1000))
|
|
qy_milli = int(round(y * 1000))
|
|
key = (qx_milli, qy_milli)
|
|
if key in self._cost_cache:
|
|
self._cost_cache.move_to_end(key)
|
|
return self._cost_cache[key]
|
|
|
|
cost = self._compute_cost_quantized(qx_milli, qy_milli)
|
|
self._cost_cache[key] = cost
|
|
if len(self._cost_cache) > _COST_CACHE_SIZE:
|
|
self._cost_cache.popitem(last=False)
|
|
return cost
|
|
|
|
def _compute_cost_quantized(self, qx_milli: int, qy_milli: int) -> float:
|
|
qx = qx_milli / 1000.0
|
|
qy = qy_milli / 1000.0
|
|
if not self.is_within_bounds(qx, qy):
|
|
return 1e15
|
|
if self.tree is None:
|
|
return 0.0
|
|
dist, _ = self.tree.query([qx, qy], distance_upper_bound=self.safety_threshold)
|
|
if dist >= self.safety_threshold:
|
|
return 0.0
|
|
safe_dist = max(dist, 0.1)
|
|
return float(self.k / (safe_dist ** 2))
|