Testing various methods of wrapping angles to [-pi,pi].

test_constrain_angle.py

import math

import pytest
import wpimath
from hypothesis import given
from hypothesis.strategies import floats


def constrain_angle_atan(angle: float) -> float:
    """Wrap an angle to the interval [-pi,pi]."""
    return math.atan2(math.sin(angle), math.cos(angle))


def constrain_angle_while(angle: float) -> float:
    while angle < -math.pi:
        angle += math.tau
    while angle > math.pi:
        angle -= math.tau
    return angle


def constrain_angle_mod(angle: float) -> float:
    angle %= math.tau
    if angle > math.pi:
        angle -= math.tau
    return angle


@pytest.fixture(
    scope="module",
    params=[
        constrain_angle_atan,
        constrain_angle_mod,
        constrain_angle_while,
        wpimath.angleModulus,
    ],
    ids=["atan", "mod", "while", "wpimath"],
)
def constrain_angle(request):
    return request.param


@given(angle=floats(-math.pi, math.pi))
def test_happy(angle: float, constrain_angle):
    """Test the happy path: the angle is in [-pi,pi]."""
    result = constrain_angle(angle)
    assert -math.pi <= result <= math.pi
    assert result == pytest.approx(angle)


def test_zero(benchmark, constrain_angle):
    assert benchmark(constrain_angle, 0) == 0


def test_edge_pos(benchmark, constrain_angle):
    result = benchmark(constrain_angle, math.pi)
    assert -math.pi <= result <= math.pi
    assert result == pytest.approx(math.pi)


def test_edge_neg(benchmark, constrain_angle):
    result = benchmark(constrain_angle, -math.pi)
    assert -math.pi <= result <= math.pi
    assert result == pytest.approx(-math.pi)


def test_revolution_pos(benchmark, constrain_angle):
    assert benchmark(constrain_angle, math.tau) == pytest.approx(0)


def test_revolution_neg(benchmark, constrain_angle):
    assert benchmark(constrain_angle, -math.tau) == pytest.approx(0)


@given(angle=floats(-math.tau, -math.pi, exclude_max=True))
def test_one_wrap_positive_half(angle: float, constrain_angle):
    result = constrain_angle(angle)
    assert -math.pi <= result <= math.pi
    assert result == pytest.approx(angle + math.tau)


@given(angle=floats(math.pi, math.tau, exclude_min=True))
def test_one_wrap_negative_half(angle: float, constrain_angle):
    result = constrain_angle(angle)
    assert -math.pi <= result <= math.pi
    assert result == pytest.approx(angle - math.tau)


@given(floats(-2 * math.tau, 2 * math.tau))
def test_agreement(angle):
    assert math.isclose(constrain_angle_atan(angle), constrain_angle_while(angle))