robosim.physics

14def integrate_velocity_into_pose(pose: Arr, velocity: Arr, dt: float) -> np.ndarray:
15    """integrates pose[t+1] = pose[t] + v[t] * dt with v=(v, w) twist. Uses rodrigues formula (batched)."""
16    pose = pose.numpy() if isinstance(pose, Field) else pose                 # (N, 4, 4)
17    velocity = velocity.numpy() if isinstance(velocity, Field) else velocity # (N, 6) [m/s]
18    res = np.float32(pose) @ btrexp(bskewa(velocity * dt)).astype("float32") # (N, 4, 4) [m]
19    if random.random() >= 0.99: # renormalize every now and then to ensure rotation matrix is still reliable
20        for i in range(len(pose)):
21            res[i, 0:3, 0:3] = renormalize_rotation_matrix(res[i, 0:3, 0:3])
22    return res

24def motion_level1(motion_input: Arr, max_velocities: Arr) -> Arr:
25    """level 1 physics: velocity is the motion input. Returns it after clamping"""
26    velocity = np.clip(motion_input, -1, 1) * max_velocities              # [m/s]
27    return velocity

29def motion_level2(motion_input: Arr, dt: float, current_velocity: Arr,
30                  max_accelerations: Arr, drag_coefficient: Arr) -> tuple[Arr, Arr]:
31    """
32    level 2 physics: integrates acceleration (motion input) + drag (ct) into the current velocity (A = (a[t] - k * v[t])
33    - v[t+1] = v[t] + dt * (a[t] - k * v[t])
34    """
35    clamped_acceleration = np.clip(motion_input, -1, 1) * max_accelerations
36    drag = current_velocity * drag_coefficient                            # [m/s**2]; (N, 6) * (N, )
37    acceleration = clamped_acceleration - drag                            # [m/s**2]; (N, 6) - (N, 6)
38    velocity = current_velocity + dt * acceleration                       # [m/s]
39    # clamp off super small velocities so they don't go "forever" in the tests
40    velocity = np.where(np.abs(velocity) < EPS, 0, velocity)              # [m/s]
41    return velocity, acceleration

10def make_grid_3d(entities: list[Entity], cell_size: Point3D) -> dict[Key3D, list[Entity]]:
11    """
12    Creates a 3d grid given all the objects based on their bounding boxes. Cell size is provided from outside.
13    Output dict is provided (not returned), so we can create it outside (e.g. prefill fixed ones)!
14    """
15    res: dict[Key3D, list[Entity]] = {}
16    assert (cell_size > 0).all(), cell_size
17    for entity in entities:
18        top_left, bottom_right = entity.collider_bbox[0:2]
19        first_cell = (top_left / cell_size).astype(int) # Note: possible bug with negative positions??
20        last_cell = (bottom_right / cell_size).astype(int)
21        nx, ny, nz = last_cell - first_cell + 1
22
23        for i in range(nx):
24            for j in range(ny):
25                for k in range(nz):
26                    key = (first_cell[0].item() + i, first_cell[1].item() + j, first_cell[2].item() + k)
27                    res.setdefault(key, []).append(entity)
28    return res

30def make_collision_cell_size(world: World) -> Point3D:
31    """the collision cell size is twice them edian of all bboxes of all collidables with a model"""
32    world.update() # update so we know for sure we use all entities even those recently added via world.add_entity
33    qr = world.query(HasModel, HasCollision)
34    bboxes = qr.model_bbox * qr.scale[..., None]
35    diffs = (bboxes[:, 1] - bboxes[:, 0]).numpy()
36    median = np.median(diffs, axis=0) # median across all 3 dimensions
37    res = median * 2 # heuristic: twice the median
38    logger.info(f"Median bbox size: {median}. Cell size: {res}")
39    return res

73def check_collision(e1: Entity | HasCollision, e2: Entity | HasCollision) -> bool:
74    """checks the collision between two entities with the HasCollision component"""
75    if e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.SPHERE:
76        return sphere_sphere_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
77                                       e2.candidate_pose[0:3, 3], e2.collider_radii.item())
78    elif e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.AABB:
79        return sphere_aabb_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
80                                     e2.collider_bbox[0], e2.collider_bbox[1])
81    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.SPHERE:
82        return sphere_aabb_collision(e2.candidate_pose[0:3, 3], e2.collider_radii.item(),
83                                     e1.collider_bbox[0], e1.collider_bbox[1])
84    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.AABB:
85        return aabb_aabb_collision(e1.collider_bbox[0], e1.collider_bbox[1],
86                                   e2.collider_bbox[0], e2.collider_bbox[1])
87    else:
88        raise NotImplementedError((e1.collider_kind, e2.collider_kind))

68def aabb_aabb_collision(top_left1: Point3D, bottom_right1: Point3D,
69                        top_left2: Point3D, bottom_right2: Point3D, eps: float=1e-5) -> bool:
70    """checks if two aabbs are colliding"""
71    return False

57def sphere_aabb_collision(center: Point3D, radius: float, top_left: Point3D,
58                          bottom_right: Point3D, eps: float=1e-5) -> bool:
59    """checks if a sphere is collidng with an aabb"""
60    xr, yr, zr = _get_closest_point_on_box(top_left, bottom_right, center)
61    dist_sq = (center[0] - xr)**2 + (center[1] - yr)**2 + (center[2] - zr)**2
62    dist: float = (dist_sq - radius ** 2).item()
63    if check := dist < -eps:
64        logger.log_every_s(f"Collision. {center=} {radius=} {top_left=} {bottom_right=}", "DEBUG", True)
65    return check

48def sphere_sphere_collision(pos1: Point3D, radius1: float, pos2: Point3D, radius2: float, eps: float=1e-5) -> bool:
49    """checks if two spheres are colliding"""
50    max_dist = radius1 + radius2
51    diff_vec = pos1 - pos2
52    dist: float = ((diff_vec**2).sum() - max_dist**2).item() # instead of np.linalg.norm(a.pos-b.pos) < max_dist
53    if check := dist < -eps:
54        logger.log_every_s(f"Collision. {pos1=} {radius1=} {pos2=} {radius2=}", "DEBUG", True)
55    return check