robosim.physics.collision

 1"""collision.py - Collision submodule between basic shapes."""
 2import numpy as np
 3from microecs import World, Entity
 4from robosim.utils import Point3D, logger
 5from robosim.components import HasModel, HasCollision, ColliderKinds
 6
 7Key3D = tuple[int, int, int]
 8
 9def make_grid_3d(entities: list[Entity], cell_size: Point3D) -> dict[Key3D, list[Entity]]:
10    """
11    Creates a 3d grid given all the objects based on their bounding boxes. Cell size is provided from outside.
12    Output dict is provided (not returned), so we can create it outside (e.g. prefill fixed ones)!
13    """
14    res: dict[Key3D, list[Entity]] = {}
15    assert (cell_size > 0).all(), cell_size
16    for entity in entities:
17        top_left, bottom_right = entity.collider_bbox[0:2]
18        first_cell = (top_left / cell_size).astype(int) # Note: possible bug with negative positions??
19        last_cell = (bottom_right / cell_size).astype(int)
20        nx, ny, nz = last_cell - first_cell + 1
21
22        for i in range(nx):
23            for j in range(ny):
24                for k in range(nz):
25                    key = (first_cell[0].item() + i, first_cell[1].item() + j, first_cell[2].item() + k)
26                    res.setdefault(key, []).append(entity)
27    return res
28
29def make_collision_cell_size(world: World) -> Point3D:
30    """the collision cell size is twice them edian of all bboxes of all collidables with a model"""
31    world.update() # update so we know for sure we use all entities even those recently added via world.add_entity
32    qr = world.query(HasModel, HasCollision)
33    bboxes = qr.model_bbox * qr.scale[..., None]
34    diffs = (bboxes[:, 1] - bboxes[:, 0]).numpy()
35    median = np.median(diffs, axis=0) # median across all 3 dimensions
36    res = median * 2 # heuristic: twice the median
37    logger.info(f"Median bbox size: {median}. Cell size: {res}")
38    return res
39
40def _get_closest_point_on_box(top_left: Point3D, bottom_right: Point3D, other_point: Point3D) -> Point3D:
41    _f = lambda l, p, r: l if p < l else min(p, r) # returns one of 3 cases: l, p or r on 1D line: [pl, [l, p, r], pr]
42    px = _f(top_left[0], other_point[0], bottom_right[0])
43    py = _f(top_left[1], other_point[1], bottom_right[1])
44    pz = _f(top_left[2], other_point[2], bottom_right[2])
45    return np.float32([px, py, pz])
46
47def sphere_sphere_collision(pos1: Point3D, radius1: float, pos2: Point3D, radius2: float, eps: float=1e-5) -> bool:
48    """checks if two spheres are colliding"""
49    max_dist = radius1 + radius2
50    diff_vec = pos1 - pos2
51    dist: float = ((diff_vec**2).sum() - max_dist**2).item() # instead of np.linalg.norm(a.pos-b.pos) < max_dist
52    if check := dist < -eps:
53        logger.log_every_s(f"Collision. {pos1=} {radius1=} {pos2=} {radius2=}", "DEBUG", True)
54    return check
55
56def sphere_aabb_collision(center: Point3D, radius: float, top_left: Point3D,
57                          bottom_right: Point3D, eps: float=1e-5) -> bool:
58    """checks if a sphere is collidng with an aabb"""
59    xr, yr, zr = _get_closest_point_on_box(top_left, bottom_right, center)
60    dist_sq = (center[0] - xr)**2 + (center[1] - yr)**2 + (center[2] - zr)**2
61    dist: float = (dist_sq - radius ** 2).item()
62    if check := dist < -eps:
63        logger.log_every_s(f"Collision. {center=} {radius=} {top_left=} {bottom_right=}", "DEBUG", True)
64    return check
65
66# pylint: disable=unused-argument
67def aabb_aabb_collision(top_left1: Point3D, bottom_right1: Point3D,
68                        top_left2: Point3D, bottom_right2: Point3D, eps: float=1e-5) -> bool:
69    """checks if two aabbs are colliding"""
70    return False
71
72def check_collision(e1: Entity | HasCollision, e2: Entity | HasCollision) -> bool:
73    """checks the collision between two entities with the HasCollision component"""
74    if e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.SPHERE:
75        return sphere_sphere_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
76                                       e2.candidate_pose[0:3, 3], e2.collider_radii.item())
77    elif e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.AABB:
78        return sphere_aabb_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
79                                     e2.collider_bbox[0], e2.collider_bbox[1])
80    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.SPHERE:
81        return sphere_aabb_collision(e2.candidate_pose[0:3, 3], e2.collider_radii.item(),
82                                     e1.collider_bbox[0], e1.collider_bbox[1])
83    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.AABB:
84        return aabb_aabb_collision(e1.collider_bbox[0], e1.collider_bbox[1],
85                                   e2.collider_bbox[0], e2.collider_bbox[1])
86    else:
87        raise NotImplementedError((e1.collider_kind, e2.collider_kind))