def test_from_matrix():
C_good = SO3.from_matrix(torch.eye(3))
assert isinstance(C_good, SO3) \
and C_good.mat.dim() == 2 \
and C_good.mat.shape == (3, 3) \
and SO3.is_valid_matrix(C_good.mat).all()
C_bad = SO3.from_matrix(torch.eye(3).add_(1e-3), normalize=True)
assert isinstance(C_bad, SO3) \
and C_bad.mat.dim() == 2 \
and C_bad.mat.shape == (3, 3) \
and SO3.is_valid_matrix(C_bad.mat).all()
def test_normalize_batch():
C = SO3.exp(torch.Tensor([[1, 2, 3], [4, 5, 6], [0, 0, 0]]))
assert (SO3.is_valid_matrix(C.mat) == torch.ByteTensor([1, 1, 1])).all()
C.mat.add_(0.1)
assert (SO3.is_valid_matrix(C.mat) == torch.ByteTensor([0, 0, 0])).all()
C.normalize(inds=[0, 2])
assert (SO3.is_valid_matrix(C.mat) == torch.ByteTensor([1, 0, 1])).all()
C.normalize()
assert SO3.is_valid_matrix(C.mat).all()
def test_from_matrix_batch():
C_good = SO3.from_matrix(torch.eye(3).repeat(5, 1, 1))
assert isinstance(C_good, SO3) \
and C_good.mat.dim() == 3 \
and C_good.mat.shape == (5, 3, 3) \
and SO3.is_valid_matrix(C_good.mat).all()
C_bad = copy.deepcopy(C_good.mat)
C_bad[3].add_(0.1)
C_bad = SO3.from_matrix(C_bad, normalize=True)
assert isinstance(C_bad, SO3) \
and C_bad.mat.dim() == 3 \
and C_bad.mat.shape == (5, 3, 3) \
and SO3.is_valid_matrix(C_bad.mat).all()
def test_normalize():
C = SO3.exp(0.25 * np.pi * torch.ones(3))
C.mat.add_(0.1)
C.normalize()
assert SO3.is_valid_matrix(C.mat).all()