def exp_map()

in theseus/geometry/se2.py [0:0]

• 29 lines of code
• 4 McCabe index (conditional complexity)

    def exp_map(tangent_vector: torch.Tensor) -> LieGroup:
        u = tangent_vector[:, :2]
        theta = tangent_vector[:, 2]
        rotation = SO2(theta=theta)

        cosine, sine = rotation.to_cos_sin()

        sin_theta_by_theta = torch.zeros_like(sine)
        one_minus_cos_theta_by_theta = torch.zeros_like(sine)

        # Compute above quantities when theta is not near zero
        idx_regular_thetas = theta.abs() > theseus.constants.EPS
        if idx_regular_thetas.any():
            sin_theta_by_theta[idx_regular_thetas] = (
                sine[idx_regular_thetas] / theta[idx_regular_thetas]
            )
            one_minus_cos_theta_by_theta[idx_regular_thetas] = (
                -cosine[idx_regular_thetas] + 1
            ) / theta[idx_regular_thetas]

        # Same as above three lines but for small angles
        idx_small_thetas = theta.abs() < theseus.constants.EPS
        if idx_small_thetas.any():
            small_theta = theta[idx_small_thetas]
            small_theta_sq = small_theta ** 2
            sin_theta_by_theta[idx_small_thetas] = -small_theta_sq / 6 + 1
            one_minus_cos_theta_by_theta[idx_small_thetas] = (
                0.5 * small_theta - small_theta / 24 * small_theta_sq
            )

        new_x = sin_theta_by_theta * u[:, 0] - one_minus_cos_theta_by_theta * u[:, 1]
        new_y = one_minus_cos_theta_by_theta * u[:, 0] + sin_theta_by_theta * u[:, 1]
        translation = Point2(data=torch.stack([new_x, new_y], dim=1))

        se2 = SE2(dtype=tangent_vector.dtype)
        se2.update_from_rot_and_trans(rotation, translation)
        return se2