|
import copy |
|
|
|
import torch |
|
import numpy as np |
|
|
|
from liegroups.torch import SO3, utils |
|
|
|
|
|
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_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_identity(): |
|
C = SO3.identity() |
|
assert isinstance(C, SO3) \ |
|
and C.mat.dim() == 2 \ |
|
and C.mat.shape == (3, 3) |
|
|
|
|
|
def test_identity_batch(): |
|
C = SO3.identity(5) |
|
assert isinstance(C, SO3) \ |
|
and C.mat.dim() == 3 \ |
|
and C.mat.shape == (5, 3, 3) |
|
|
|
C_copy = SO3.identity(5, copy=True) |
|
assert isinstance(C_copy, SO3) \ |
|
and C_copy.mat.dim() == 3 \ |
|
and C_copy.mat.shape == (5, 3, 3) |
|
|
|
|
|
def test_dot(): |
|
C = SO3(torch.Tensor([[0, -1, 0], |
|
[1, 0, 0], |
|
[0, 0, 1]])) |
|
pt = torch.Tensor([1, 2, 3]) |
|
|
|
CC = C.mat.mm(C.mat) |
|
assert utils.allclose(C.dot(C).mat, CC) |
|
|
|
Cpt = C.mat.matmul(pt) |
|
assert utils.allclose(C.dot(pt), Cpt) |
|
|
|
|
|
def test_dot_batch(): |
|
C1 = SO3(torch.Tensor([[0, -1, 0], |
|
[1, 0, 0], |
|
[0, 0, 1]]).expand(5, 3, 3)) |
|
C3 = SO3(torch.Tensor([[0, -1, 0], |
|
[1, 0, 0], |
|
[0, 0, 1]])) |
|
pt1 = torch.Tensor([1, 2, 3]) |
|
pt3 = torch.Tensor([4, 5, 6]) |
|
pt3 = torch.Tensor([7, 8, 9]) |
|
pts = torch.cat([pt1.unsqueeze(dim=0), |
|
pt3.unsqueeze(dim=0), |
|
pt3.unsqueeze(dim=0)], dim=0) |
|
ptsbatch = pts.unsqueeze(dim=0).expand(5, 3, 3) |
|
|
|
C1C1 = torch.bmm(C1.mat, C1.mat) |
|
C1C1_SO3 = C1.dot(C1).mat |
|
assert C1C1_SO3.shape == C1.mat.shape and utils.allclose(C1C1_SO3, C1C1) |
|
|
|
C1C3 = torch.matmul(C1.mat, C3.mat) |
|
C1C3_SO3 = C1.dot(C3).mat |
|
assert C1C3_SO3.shape == C1.mat.shape and utils.allclose(C1C3_SO3, C1C3) |
|
|
|
C1pt1 = torch.matmul(C1.mat, pt1) |
|
C1pt1_SO3 = C1.dot(pt1) |
|
assert C1pt1_SO3.shape == (C1.mat.shape[0], pt1.shape[0]) \ |
|
and utils.allclose(C1pt1_SO3, C1pt1) |
|
|
|
C1pt3 = torch.matmul(C1.mat, pt3) |
|
C1pt3_SO3 = C1.dot(pt3) |
|
assert C1pt3_SO3.shape == (C1.mat.shape[0], pt3.shape[0]) \ |
|
and utils.allclose(C1pt3_SO3, C1pt3) |
|
|
|
C1pts = torch.matmul(C1.mat, pts.transpose(1, 0)).transpose(2, 1) |
|
C1pts_SO3 = C1.dot(pts) |
|
assert C1pts_SO3.shape == (C1.mat.shape[0], pts.shape[0], pts.shape[1]) \ |
|
and utils.allclose(C1pts_SO3, C1pts) \ |
|
and utils.allclose(C1pt1, C1pts[:, 0, :]) \ |
|
and utils.allclose(C1pt3, C1pts[:, 1, :]) |
|
|
|
C1ptsbatch = torch.bmm(C1.mat, ptsbatch.transpose(2, 1)).transpose(2, 1) |
|
C1ptsbatch_SO3 = C1.dot(ptsbatch) |
|
assert C1ptsbatch_SO3.shape == ptsbatch.shape \ |
|
and utils.allclose(C1ptsbatch_SO3, C1ptsbatch) \ |
|
and utils.allclose(C1pt1, C1ptsbatch[:, 0, :]) \ |
|
and utils.allclose(C1pt3, C1ptsbatch[:, 1, :]) |
|
|
|
C3ptsbatch = torch.matmul(C3.mat, ptsbatch.transpose(2, 1)).transpose(2, 1) |
|
C3ptsbatch_SO3 = C3.dot(ptsbatch) |
|
assert C3ptsbatch_SO3.shape == ptsbatch.shape \ |
|
and utils.allclose(C3ptsbatch_SO3, C3ptsbatch) \ |
|
and utils.allclose(C3.dot(pt1), C3ptsbatch[:, 0, :]) \ |
|
and utils.allclose(C3.dot(pt3), C3ptsbatch[:, 1, :]) |
|
|
|
|
|
def test_wedge(): |
|
phi = torch.Tensor([1, 2, 3]) |
|
Phi = SO3.wedge(phi) |
|
assert (Phi == -Phi.t()).all() |
|
|
|
|
|
def test_wedge_batch(): |
|
phis = torch.Tensor([[1, 2, 3], |
|
[4, 5, 6]]) |
|
Phis = SO3.wedge(phis) |
|
assert (Phis[0, :, :] == SO3.wedge(phis[0])).all() |
|
assert (Phis[1, :, :] == SO3.wedge(phis[1])).all() |
|
|
|
|
|
def test_wedge_vee(): |
|
phi = torch.Tensor([1, 2, 3]) |
|
Phi = SO3.wedge(phi) |
|
assert (phi == SO3.vee(Phi)).all() |
|
|
|
|
|
def test_wedge_vee_batch(): |
|
phis = torch.Tensor([[1, 2, 3], |
|
[4, 5, 6]]) |
|
Phis = SO3.wedge(phis) |
|
assert (phis == SO3.vee(Phis)).all() |
|
|
|
|
|
def test_left_jacobians(): |
|
phi_small = torch.Tensor([0., 0., 0.]) |
|
phi_big = torch.Tensor([np.pi / 2, np.pi / 3, np.pi / 4]) |
|
|
|
left_jacobian_small = SO3.left_jacobian(phi_small) |
|
inv_left_jacobian_small = SO3.inv_left_jacobian(phi_small) |
|
assert utils.allclose( |
|
torch.mm(left_jacobian_small, inv_left_jacobian_small), |
|
torch.eye(3)) |
|
|
|
left_jacobian_big = SO3.left_jacobian(phi_big) |
|
inv_left_jacobian_big = SO3.inv_left_jacobian(phi_big) |
|
assert utils.allclose( |
|
torch.mm(left_jacobian_big, inv_left_jacobian_big), |
|
torch.eye(3)) |
|
|
|
|
|
def test_left_jacobians_batch(): |
|
phis = torch.Tensor([[0., 0., 0.], |
|
[np.pi / 2, np.pi / 3, np.pi / 4]]) |
|
|
|
left_jacobian = SO3.left_jacobian(phis) |
|
inv_left_jacobian = SO3.inv_left_jacobian(phis) |
|
assert utils.allclose(torch.bmm(left_jacobian, inv_left_jacobian), |
|
torch.eye(3).unsqueeze_(dim=0).expand(2, 3, 3)) |
|
|
|
|
|
def test_exp_log(): |
|
C_big = SO3.exp(0.25 * np.pi * torch.ones(3)) |
|
assert utils.allclose(SO3.exp(SO3.log(C_big)).mat, C_big.mat) |
|
|
|
C_small = SO3.exp(torch.zeros(3)) |
|
assert utils.allclose(SO3.exp(SO3.log(C_small)).mat, C_small.mat) |
|
|
|
|
|
def test_exp_log_batch(): |
|
C = SO3.exp(torch.Tensor([[1, 2, 3], |
|
[0, 0, 0]])) |
|
assert utils.allclose(SO3.exp(SO3.log(C)).mat, C.mat) |
|
|
|
|
|
def test_perturb(): |
|
C = SO3.exp(0.25 * np.pi * torch.ones(3)) |
|
C_copy = copy.deepcopy(C) |
|
phi = torch.Tensor([0.1, 0.2, 0.3]) |
|
C.perturb(phi) |
|
assert utils.allclose( |
|
C.as_matrix(), (SO3.exp(phi).dot(C_copy)).as_matrix()) |
|
|
|
|
|
def test_perturb_batch(): |
|
C = SO3.exp(torch.Tensor([[1, 2, 3], |
|
[4, 5, 6]])) |
|
C_copy1 = copy.deepcopy(C) |
|
C_copy2 = copy.deepcopy(C) |
|
|
|
phi = torch.Tensor([0.1, 0.2, 0.3]) |
|
C_copy1.perturb(phi) |
|
assert utils.allclose(C_copy1.as_matrix(), |
|
(SO3.exp(phi).dot(C)).as_matrix()) |
|
|
|
phis = torch.Tensor([[0.1, 0.2, 0.3], |
|
[0.4, 0.5, 0.6]]) |
|
C_copy2.perturb(phis) |
|
assert utils.allclose(C_copy2.as_matrix(), |
|
(SO3.exp(phis).dot(C)).as_matrix()) |
|
|
|
|
|
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() |
|
|
|
|
|
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_inv(): |
|
C = SO3.exp(0.25 * np.pi * torch.ones(3)) |
|
assert utils.allclose(C.dot(C.inv()).mat, SO3.identity().mat) |
|
|
|
|
|
def test_inv_batch(): |
|
C = SO3.exp(torch.Tensor([[1, 2, 3], |
|
[4, 5, 6]])) |
|
assert utils.allclose(C.dot(C.inv()).mat, SO3.identity(C.mat.shape[0]).mat) |
|
|
|
|
|
def test_adjoint(): |
|
C = SO3.exp(0.25 * np.pi * torch.ones(3)) |
|
assert (C.adjoint() == C.mat).all() |
|
|
|
|
|
def test_adjoint_batch(): |
|
C = SO3.exp(torch.Tensor([[1, 2, 3], |
|
[4, 5, 6]])) |
|
assert (C.adjoint() == C.mat).all() |
|
|
|
|
|
def test_rotx(): |
|
C_got = SO3.rotx(torch.Tensor([np.pi / 2])) |
|
C_expected = torch.Tensor([[1, 0, 0], |
|
[0, 0, -1], |
|
[0, 1, 0]]) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_rotx_batch(): |
|
C_got = SO3.rotx(torch.Tensor([np.pi / 2, np.pi])) |
|
C_expected = torch.cat([torch.Tensor([[1, 0, 0], |
|
[0, 0, -1], |
|
[0, 1, 0]]).unsqueeze_(dim=0), |
|
torch.Tensor([[1, 0, 0], |
|
[0, -1, 0], |
|
[0, 0, -1]]).unsqueeze_(dim=0)], dim=0) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_roty(): |
|
C_got = SO3.roty(torch.Tensor([np.pi / 2])) |
|
C_expected = torch.Tensor([[0, 0, 1], |
|
[0, 1, 0], |
|
[-1, 0, 0]]) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_roty_batch(): |
|
C_got = SO3.roty(torch.Tensor([np.pi / 2, np.pi])) |
|
C_expected = torch.cat([torch.Tensor([[0, 0, 1], |
|
[0, 1, 0], |
|
[-1, 0, 0]]).unsqueeze_(dim=0), |
|
torch.Tensor([[-1, 0, 0], |
|
[0, 1, 0], |
|
[0, 0, -1]]).unsqueeze_(dim=0)], dim=0) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_rotz(): |
|
C_got = SO3.rotz(torch.Tensor([np.pi / 2])) |
|
C_expected = torch.Tensor([[0, -1, 0], |
|
[1, 0, 0], |
|
[0, 0, 1]]) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_rotz_batch(): |
|
C_got = SO3.rotz(torch.Tensor([np.pi / 2, np.pi])) |
|
C_expected = torch.cat([torch.Tensor([[0, -1, 0], |
|
[1, 0, 0], |
|
[0, 0, 1]]).unsqueeze_(dim=0), |
|
torch.Tensor([[-1, 0, 0], |
|
[0, -1, 0], |
|
[0, 0, 1]]).unsqueeze_(dim=0)], dim=0) |
|
assert utils.allclose(C_got.mat, C_expected) |
|
|
|
|
|
def test_rpy(): |
|
rpy = torch.Tensor([np.pi / 12, np.pi / 6, np.pi / 3]) |
|
C_got = SO3.from_rpy(rpy) |
|
C_expected = SO3.rotz(torch.Tensor([rpy[2]])).dot( |
|
SO3.roty(torch.Tensor([rpy[1]])).dot( |
|
SO3.rotx(torch.Tensor([rpy[0]])) |
|
) |
|
) |
|
assert utils.allclose(C_got.mat, C_expected.mat) |
|
|
|
|
|
def test_rpy_batch(): |
|
rpy = torch.Tensor([[np.pi / 12, np.pi / 6, np.pi / 3], |
|
[0, 0, 0]]) |
|
C_got = SO3.from_rpy(rpy) |
|
C_expected = SO3.rotz(rpy[:, 2]).dot( |
|
SO3.roty(rpy[:, 1]).dot( |
|
SO3.rotx(rpy[:, 0]) |
|
) |
|
) |
|
assert utils.allclose(C_got.mat, C_expected.mat) |
|
|
|
|
|
def test_quaternion(): |
|
q1 = torch.Tensor([1, 0, 0, 0]) |
|
q2 = torch.Tensor([0, 1, 0, 0]) |
|
q3 = torch.Tensor([0, 0, 1, 0]) |
|
q4 = torch.Tensor([0, 0, 0, 1]) |
|
q5 = 0.5 * torch.ones(4) |
|
q6 = -q5 |
|
|
|
assert utils.allclose(SO3.from_quaternion(q1).to_quaternion(), q1) |
|
assert utils.allclose(SO3.from_quaternion(q2).to_quaternion(), q2) |
|
assert utils.allclose(SO3.from_quaternion(q3).to_quaternion(), q3) |
|
assert utils.allclose(SO3.from_quaternion(q4).to_quaternion(), q4) |
|
assert utils.allclose(SO3.from_quaternion(q5).to_quaternion(), q5) |
|
assert utils.allclose(SO3.from_quaternion(q5).mat, |
|
SO3.from_quaternion(q6).mat) |
|
|
|
|
|
def test_quaternion_batch(): |
|
quats = torch.Tensor([[1, 0, 0, 0], |
|
[0, 1, 0, 0], |
|
[0, 0, 1, 0], |
|
[0, 0, 0, 1], |
|
[0.5, 0.5, 0.5, 0.5]]) |
|
|
|
assert utils.allclose(SO3.from_quaternion(quats).to_quaternion(), quats) |
|
|
