wrap.py

def _wrap_to_2pi(scalar_or_arr):
    """Offsets angles outside 0 <= x <= 2 * pi to lie within the interval"""
    if np.isscalar(scalar_or_arr):
        if scalar_or_arr < 0:
            scalar_or_arr += 2 * np.pi
        if scalar_or_arr > 2 * np.pi:
            scalar_or_arr -= 2 * np.pi
    else:
        j = np.where(scalar_or_arr < 0)
        if len(j[0]) > 0:
            scalar_or_arr[j] += 2 * np.pi

        j = np.where(scalar_or_arr > 2 * np.pi)
        if len(j[0]) > 0:
            scalar_or_arr[j] -= 2 * np.pi
    return scalar_or_arr

_wrap_to_2pi(3 * np.pi) == np.pi
_wrap_to_2pi(np.array([-1 * np.pi, 3 * np.pi])) == np.array([np.pi, np.pi])