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

Code before:

@make_xp_test_case(RigidTransform.__mul__)
def test_composition_validation(xp):
    tf2 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
    tf3 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))

    # Test incompatible shapes
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf3

    tf4 = RigidTransform.from_matrix(xp.tile(xp.eye(4), (1, 4, 1, 1)))
    # Test invalid broadcasting shape
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf4


@make_xp_test_case(RigidTransform.__mul__, RigidTransform.__rmul__)
def test_rotation_promotion(xp):
    """Test that Rotation is promoted to RigidTransform in composition."""
    atol = 1e-12
    dtype = xpx.default_dtype(xp)
    rng = np.random.default_rng(0)

    t = xp.asarray(rng.normal(size=(3,)), dtype=dtype)
    r1 = Rotation.random(rng=rng)
    r2 = Rotation.random(rng=rng)
    tf = RigidTransform.from_components(t, r2)

    # RigidTransform * Rotation
    result = tf * r1
    expected = tf * RigidTransform.from_rotation(r1)
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)

    # Rotation * RigidTransform
    result = r1 * tf
    expected = RigidTransform.from_rotation(r1) * tf
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)


@make_xp_test_case(RigidTransform.concatenate)
def test_concatenate_validation(xp):
    tf = RigidTransform.from_matrix(xp.eye(4))

    # Test invalid inputs
    with pytest.raises(TypeError,
                       match="input must contain RigidTransform objects"):
        RigidTransform.concatenate([tf, xp.eye(4)])

    # Test incompatible shapes

Review comment (suggestion):
```suggestion
    r1 = rotation_to_xp(Rotation.random(rng=rng), xp=xp)
    r2 = rotation_to_xp(Rotation.random(rng=rng), xp=xp)
```

Code after:

@make_xp_test_case(RigidTransform.__mul__)
def test_composition_validation(xp):
    tf2 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
    tf3 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))

    # Test incompatible shapes
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf3

    tf4 = RigidTransform.from_matrix(xp.tile(xp.eye(4), (1, 4, 1, 1)))
    # Test invalid broadcasting shape
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf4


@make_xp_test_case(RigidTransform.__mul__, RigidTransform.__rmul__)
def test_rotation_promotion(xp):
    """Test that Rotation is promoted to RigidTransform in composition."""
    atol = 1e-12
    dtype = xpx.default_dtype(xp)
    rng = np.random.default_rng(0)

    t = xp.asarray(rng.normal(size=(3,)), dtype=dtype)
    r1 = rotation_to_xp(Rotation.random(rng=rng), xp)
    r2 = rotation_to_xp(Rotation.random(rng=rng), xp)
    tf = RigidTransform.from_components(t, r2)

    # RigidTransform * Rotation
    result = tf * r1
    expected = tf * RigidTransform.from_rotation(r1)
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)

    # Rotation * RigidTransform
    result = r1 * tf
    expected = RigidTransform.from_rotation(r1) * tf
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)


@make_xp_test_case(RigidTransform.concatenate)
def test_concatenate_validation(xp):
    tf = RigidTransform.from_matrix(xp.eye(4))

    # Test invalid inputs
    with pytest.raises(TypeError,
                       match="input must contain RigidTransform objects"):
        RigidTransform.concatenate([tf, xp.eye(4)])

    # Test incompatible shapes