def in_cone2d()

in mae_envs/util/vision.py [0:0]

• 28 lines of code
• 3 McCabe index (conditional complexity)

def in_cone2d(origin_pts, origin_angles, cone_angle, target_pts):
    '''
        Computes whether 2D points target_pts are in the cones originating from
            origin_pts at angle origin_angles with cone spread angle cone_angle.
        Args:
            origin_pts (np.ndarray): array with shape (n_points, 2) of origin points
            origin_angles (np.ndarray): array with shape (n_points,) of origin angles
            cone_angle (float): cone angle width
            target_pts (np.ndarray): target points to check whether in cones
        Returns:
            np.ndarray of bools. Each row corresponds to origin cone, and columns to
                target points
    '''
    assert isinstance(origin_pts, np.ndarray)
    assert isinstance(origin_angles, np.ndarray)
    assert isinstance(cone_angle, float)
    assert isinstance(target_pts, np.ndarray)
    assert origin_pts.shape[0] == origin_angles.shape[0]
    assert len(origin_angles.shape) == 1, "Angles should only have 1 dimension"
    np.seterr(divide='ignore', invalid='ignore')
    cone_vec = np.array([np.cos(origin_angles), np.sin(origin_angles)]).T
    # Compute normed vectors between all pairs of agents
    pos_diffs = target_pts[None, ...] - origin_pts[:, None, :]
    norms = np.sqrt(np.sum(np.square(pos_diffs), -1, keepdims=True))
    unit_diffs = pos_diffs / norms
    # Dot product between unit vector in middle of cone and the vector
    dot_cone_diff = np.sum(unit_diffs * cone_vec[:, None, :], -1)
    angle_between = np.arccos(dot_cone_diff)
    # Right now the only thing that should be nan will be targets that are on the origin point
    # This can only happen for the origin looking at itself, so just make this always true
    angle_between[np.isnan(angle_between)] = 0.

    return np.abs(normalize_angles(angle_between)) <= cone_angle