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

Code before:
    mat[0] = np.array([
        [0, -1, 0],
        [1, 0, 0],
        [0, 0, 1]
    ])
    mat[1] = np.array([
        [1, 0, 0],
        [0, 0, -1],
        [0, 1, 0]
    ])
    mat = xp.asarray(mat)
    r = Rotation.from_matrix(mat)

    assert_equal(len(r), 2)
    assert_equal(len(r[:-1]), 1)


def test_random_rotation():
    # No xp testing since random rotations are always using NumPy
    rng = np.random.default_rng(0)
    assert_equal(Rotation.random(rng=rng).as_quat().shape, (4,))
    assert_equal(Rotation.random(None, rng=rng).as_quat().shape, (4,))
    assert_equal(Rotation.random(1, rng=rng).as_quat().shape, (1, 4))
    assert_equal(Rotation.random(5, rng=rng).as_quat().shape, (5, 4))
    # Shape argument
    assert_equal(Rotation.random(rng=rng,shape=()).as_quat().shape, (4,))
    assert_equal(Rotation.random(rng=rng, shape=(3,)).as_quat().shape, (3, 4))
    assert_equal(Rotation.random(rng=rng, shape=(2, 3)).as_quat().shape, (2, 3, 4))
    # Values should be the same for num=prod(shape)
    rng1, rng2 = np.random.default_rng(42), np.random.default_rng(42)
    r_num = Rotation.random(6, rng=rng1)
    r_shape = Rotation.random(rng=rng2, shape=(2, 3))
    xp_assert_equal(r_num.as_quat(), r_shape.as_quat().reshape(6, 4))
    # Errors
    with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"):
        Rotation.random(num=3,rng=rng, shape=(2, 2))
    with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"):
        Rotation.random(rng=rng, shape=2.5)
    with pytest.raises(TypeError, match="takes from 0 to 2 positional arguments"):
        Rotation.random(1, rng, None)  # Shape should be kwarg only


@make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix)
def test_align_vectors_no_rotation(xp):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-5
    x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype)
    y = xp.asarray(x, copy=True)

    r, rssd = Rotation.align_vectors(x, y)
    xp_assert_close(r.as_matrix(), xp.eye(3), atol=atol)

Review comment (suggestion):
```suggestion
    assert_equal(Rotation.random(rng=rng, shape=()).as_quat().shape, (4,))
```

Code after:
    mat[0] = np.array([
        [0, -1, 0],
        [1, 0, 0],
        [0, 0, 1]
    ])
    mat[1] = np.array([
        [1, 0, 0],
        [0, 0, -1],
        [0, 1, 0]
    ])
    mat = xp.asarray(mat)
    r = Rotation.from_matrix(mat)

    assert_equal(len(r), 2)
    assert_equal(len(r[:-1]), 1)


def test_random_rotation():
    # No xp testing since random rotations are always using NumPy
    rng = np.random.default_rng(0)
    assert_equal(Rotation.random(rng=rng).as_quat().shape, (4,))
    assert_equal(Rotation.random(None, rng=rng).as_quat().shape, (4,))
    assert_equal(Rotation.random(1, rng=rng).as_quat().shape, (1, 4))
    assert_equal(Rotation.random(5, rng=rng).as_quat().shape, (5, 4))
    # Shape argument
    assert_equal(Rotation.random(rng=rng, shape=()).as_quat().shape, (4,))
    assert_equal(Rotation.random(rng=rng, shape=(3,)).as_quat().shape, (3, 4))
    assert_equal(Rotation.random(rng=rng, shape=(2, 3)).as_quat().shape, (2, 3, 4))
    # Values should be the same for num=prod(shape)
    rng1, rng2 = np.random.default_rng(42), np.random.default_rng(42)
    r_num = Rotation.random(6, rng=rng1)
    r_shape = Rotation.random(rng=rng2, shape=(2, 3))
    xp_assert_equal(r_num.as_quat(), r_shape.as_quat().reshape(6, 4))
    # Errors
    with pytest.raises(ValueError, match="Only one of `num` or `shape` can be"):
        Rotation.random(num=3,rng=rng, shape=(2, 2))
    with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"):
        Rotation.random(rng=rng, shape=2.5)
    with pytest.raises(TypeError, match="takes from 0 to 2 positional arguments"):
        Rotation.random(1, rng, None)  # Shape should be kwarg only


@make_xp_test_case(Rotation.align_vectors, Rotation.as_matrix)
def test_align_vectors_no_rotation(xp):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-5
    x = xp.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype)
    y = xp.asarray(x, copy=True)

    r, rssd = Rotation.align_vectors(x, y)
    xp_assert_close(r.as_matrix(), xp.eye(3), atol=atol)