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

Code before:
    # Test left handed rotation matrix
    matrix = xp.eye(4)
    matrix = xpx.at(matrix)[0, 0].set(-1)
    if is_lazy_array(matrix):
        matrix = RigidTransform.from_matrix(matrix).as_matrix()
        assert xp.all(xp.isnan(matrix[..., :3, :3]))
    else:
        with pytest.raises(ValueError, match="Non-positive determinant"):
            RigidTransform(matrix, normalize=True)

    # Test non-Rotation input
    with pytest.raises(TypeError,
                       match="Expected `rotation` to be a `Rotation` instance"):
        RigidTransform.from_rotation(xp.eye(3))


@make_xp_test_case(RigidTransform.mean)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_mean(xp, ndim: int):
    atol = 1e-12
    rng = np.random.default_rng(123)

    dtype = xpx.default_dtype(xp)
    t = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (4,)), dtype=dtype)
    r = rotation_to_xp(Rotation.from_quat(q), xp=xp)
    tf = RigidTransform.from_components(t, r)

    # Unweighted mean
    axis = tuple(range(t.ndim - 1))
    t_mean = xp.mean(t, axis=axis)
    r_mean = r.mean()
    tf_mean = tf.mean()
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),
                    RigidTransform.from_components(t_mean, r_mean).as_matrix(),
                    atol=atol)

    # Weighted mean
    if ndim == 1:
        weights = None
        t_mean = t
    else:
        weights = xp.asarray(rng.random(size=(ndim,) * (ndim - 1)), dtype=dtype)
        norm = xp.sum(weights[..., None], axis=axis)
        wsum = xp.sum(t * weights[..., None], axis=axis)
        t_mean = wsum/norm
    r_mean = r.mean(weights=weights)
    tf_mean = tf.mean(weights=weights)
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),

Review comment (suggestion):
Should not be necessary since Rotation was already created from an xp array of correct dtype.

Code after:
    # Test left handed rotation matrix
    matrix = xp.eye(4)
    matrix = xpx.at(matrix)[0, 0].set(-1)
    if is_lazy_array(matrix):
        matrix = RigidTransform.from_matrix(matrix).as_matrix()
        assert xp.all(xp.isnan(matrix[..., :3, :3]))
    else:
        with pytest.raises(ValueError, match="Non-positive determinant"):
            RigidTransform(matrix, normalize=True)

    # Test non-Rotation input
    with pytest.raises(TypeError,
                       match="Expected `rotation` to be a `Rotation` instance"):
        RigidTransform.from_rotation(xp.eye(3))


@make_xp_test_case(RigidTransform.mean)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_mean(xp, ndim: int):
    atol = 1e-12
    rng = np.random.default_rng(123)

    dtype = xpx.default_dtype(xp)
    t = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (4,)), dtype=dtype)
    r = Rotation.from_quat(q)
    tf = RigidTransform.from_components(t, r)

    # Unweighted mean
    axis = tuple(range(t.ndim - 1))
    t_mean = xp.mean(t, axis=axis)
    r_mean = r.mean()
    tf_mean = tf.mean()
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),
                    RigidTransform.from_components(t_mean, r_mean).as_matrix(),
                    atol=atol)

    # Weighted mean
    if ndim == 1:
        weights = None
        t_mean = t
    else:
        weights = xp.asarray(rng.random(size=(ndim,) * (ndim - 1)), dtype=dtype)
        norm = xp.sum(weights[..., None], axis=axis)
        wsum = xp.sum(t * weights[..., None], axis=axis)
        t_mean = wsum/norm
    r_mean = r.mean(weights=weights)
    tf_mean = tf.mean(weights=weights)
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),