Repository: scipy/scipy
File: scipy/spatial/transform/tests/test_rotation.py (Python)

Code before:
        r = Rotation.from_rotvec(t * axes)
        xp_assert_close(r.mean().magnitude(), desired, atol=atol)


@make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.magnitude)
@pytest.mark.parametrize("ndim", range(1, 5))
def test_mean_axis(xp, ndim: int):
    axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1))
    theta = xp.pi / 4
    r = Rotation.from_rotvec(theta * axes)
    desired = xp.full(axes.shape[:-2], 0.0)
    if ndim == 1:
        desired = desired[()]
    atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10
    xp_assert_close(r.mean(axis=-1).magnitude(), desired, atol=atol)
    # Test tuple axes
    desired = xp.full(axes.shape[1:-2], 0.0)
    if ndim < 3:
        desired = desired[()]
    xp_assert_close(r.mean(axis=(0, -1)).magnitude(), desired, atol=atol)
    # Empty axis tuple should return Rotation unchanged
    r_mean = r.mean(axis=())
    xp_assert_close(r_mean.as_quat(canonical=True), r.as_quat(canonical=True),
                    atol=atol)


@make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.inv,
                   Rotation.magnitude)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_weighted_mean(xp, ndim: int):
    # test that doubling a weight is equivalent to including a rotation twice.
    thetas = xp.linspace(0, xp.pi / 2, 100)

    # Create batched copies of the same setup
    batch_shape = (ndim,) * (ndim - 1)
    axes = xp.asarray([[0.0, 0, 0], [1, 0, 0], [1, 0, 0]])
    weights = xp.asarray([1, 2])
    axes = xp.tile(axes, batch_shape + (1, 1))
    weights = xp.tile(weights, batch_shape + (1,))

    expected = xp.asarray(0.0)[()]
    for t in thetas:
        rw = Rotation.from_rotvec(t * axes[..., :2, :])
        mw = rw.mean(weights=weights)

        r = Rotation.from_rotvec(t * axes)
        m = r.mean()
        xp_assert_close((m * mw.inv()).magnitude(), expected, atol=1e-6)


@make_xp_test_case(Rotation.mean)

Review comment (suggestion):
The expected mean being zero in all these cases is a useful check, but doesn't fully exercise things IMO. Could you add a test with random rotations, that compares a reduction along several axes to a list of axis=None mean of rotations formed from those slices? That will prove out the equivalence to the existing algorithm.

Code after:
    q = Rotation.from_quat(xp.ones((3, 2, 4)))
    with pytest.raises(ValueError, match="broadcastable shapes"):
        p.approx_equal(q)

    p = Rotation.from_quat(xp.ones((2, 4)))
    q = Rotation.from_quat(xp.ones((3, 4)))
    with pytest.raises(ValueError, match="broadcastable shapes"):
        p.approx_equal(q)


@make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.magnitude)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_mean(xp, ndim: int):
    axes = xp.concat((-xp.eye(3), xp.eye(3)))
    axes = xp.reshape(axes, (1,) * (ndim - 1) + (6, 3))
    thetas = xp.linspace(0, xp.pi / 2, 100)
    desired = xp.asarray(0.0)[()]
    atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10
    for t in thetas:
        r = Rotation.from_rotvec(t * axes)
        xp_assert_close(r.mean().magnitude(), desired, atol=atol)


@make_xp_test_case(Rotation.from_rotvec, Rotation.mean, Rotation.magnitude)
@pytest.mark.parametrize("ndim", range(1, 5))
def test_mean_axis(xp, ndim: int):
    axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1))
    theta = xp.pi / 4
    r = Rotation.from_rotvec(theta * axes)
    desired = xp.full(axes.shape[:-2], 0.0)
    if ndim == 1:
        desired = desired[()]
    atol = 1e-6 if xp_default_dtype(xp) is xp.float32 else 1e-10
    xp_assert_close(r.mean(axis=-1).magnitude(), desired, atol=atol)
    # Test tuple axes
    desired = xp.full(axes.shape[1:-2], 0.0)
    if ndim < 3:
        desired = desired[()]
    xp_assert_close(r.mean(axis=(0, -1)).magnitude(), desired, atol=atol)
    # Empty axis tuple should return Rotation unchanged
    r_mean = r.mean(axis=())
    xp_assert_close(r_mean.as_quat(canonical=True), r.as_quat(canonical=True),
                    atol=atol)


@make_xp_test_case(Rotation.mean, Rotation.magnitude)
def test_mean_compare_axis(xp):
    # Create a random set of rotations and compare the mean over an axis with the
    # mean without axis of the sliced quaternion
    atol = 1e-10 if xpx.default_dtype(xp) == xp.float64 else 1e-6
    rng = np.random.default_rng(0)