def is_valid_matrix(cls, mat):
if mat.dim() < 3:
mat = mat.unsqueeze(dim=0)
# Check the shape
if mat.is_cuda:
shape_check = torch.cuda.ByteTensor(mat.shape[0]).fill_(False)
else:
shape_check = torch.ByteTensor(mat.shape[0]).fill_(False)
if mat.shape[1:3] != (cls.dim, cls.dim):
return shape_check
else:
shape_check.fill_(True)
# Determinants of each matrix in the batch should be 1
det_check = utils.isclose(
mat.__class__(np.linalg.det(mat.cpu().numpy())), 1.)
# The transpose of each matrix in the batch should be its inverse
inv_check = utils.isclose(mat.transpose(2, 1).bmm(mat),
torch.eye(cls.dim)).sum(dim=1).sum(dim=1) \
== cls.dim * cls.dim
return shape_check & det_check & inv_check
def to_rpy(self):
"""Convert a rotation matrix to RPY Euler angles."""
if self.mat.dim() < 3:
mat = self.mat.unsqueeze(dim=0)
else:
mat = self.mat
pitch = torch.atan2(-mat[:, 2, 0],
torch.sqrt(mat[:, 0, 0]**2 + mat[:, 1, 0]**2))
yaw = pitch.__class__(pitch.shape)
roll = pitch.__class__(pitch.shape)
near_pi_over_two_mask = utils.isclose(pitch, np.pi / 2.)
near_pi_over_two_inds = near_pi_over_two_mask.nonzero().squeeze_()
near_neg_pi_over_two_mask = utils.isclose(pitch, -np.pi / 2.)
near_neg_pi_over_two_inds = near_neg_pi_over_two_mask.nonzero(
).squeeze_()
remainder_inds = (1 -
(near_pi_over_two_mask
| near_neg_pi_over_two_mask)).nonzero().squeeze_()
if len(near_pi_over_two_inds) > 0:
yaw[near_pi_over_two_inds] = 0.
roll[near_pi_over_two_inds] = torch.atan2(
mat[near_pi_over_two_inds, 0, 1], mat[near_pi_over_two_inds, 1,
1])
if len(near_neg_pi_over_two_inds) > 0:
yaw[near_pi_over_two_inds] = 0.
roll[near_pi_over_two_inds] = -torch.atan2(
mat[near_pi_over_two_inds, 0, 1], mat[near_pi_over_two_inds, 1,
1])
if len(remainder_inds) > 0:
sec_pitch = 1. / pitch[remainder_inds].cos()
remainder_mats = mat[remainder_inds]
yaw = torch.atan2(remainder_mats[:, 1, 0] * sec_pitch,
remainder_mats[:, 0, 0] * sec_pitch)
roll = torch.atan2(remainder_mats[:, 2, 1] * sec_pitch,
remainder_mats[:, 2, 2] * sec_pitch)
return torch.cat([
roll.unsqueeze_(dim=1),
pitch.unsqueeze_(dim=1),
yaw.unsqueeze_(dim=1)
],
dim=1).squeeze_()
def left_jacobian(cls, phi):
"""(see Barfoot/Eade)."""
jac = phi.__class__(phi.shape[0], cls.dim, cls.dim)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(phi, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_(dim=1)
if len(small_angle_inds) > 0:
jac[small_angle_inds] = torch.eye(cls.dim).expand(
len(small_angle_inds), cls.dim, cls.dim) \
+ 0.5 * cls.wedge(phi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_(dim=1)
if len(large_angle_inds) > 0:
angle = phi[large_angle_inds]
s = angle.sin()
c = angle.cos()
A = (s / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
torch.eye(cls.dim).unsqueeze_(dim=0).expand_as(
jac[large_angle_inds])
B = ((1. - c) / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
cls.wedge(phi.__class__([1.]))
jac[large_angle_inds] = A + B
return jac.squeeze_()
def inv_left_jacobian(cls, phi):
"""(see Barfoot/Eade)."""
jac = phi.__class__(phi.shape[0], cls.dim, cls.dim)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(phi, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_(dim=1)
if len(small_angle_inds) > 0:
jac[small_angle_inds] = torch.eye(cls.dim).expand(
len(small_angle_inds), cls.dim, cls.dim) \
- 0.5 * cls.wedge(phi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_(dim=1)
if len(large_angle_inds) > 0:
angle = phi[large_angle_inds]
ha = 0.5 * angle # half angle
hacha = ha / ha.tan() # half angle * cot(half angle)
ha.unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
jac[large_angle_inds])
hacha.unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
jac[large_angle_inds])
A = hacha * \
torch.eye(cls.dim).unsqueeze_(
dim=0).expand_as(jac[large_angle_inds])
B = -ha * cls.wedge(phi.__class__([1.]))
jac[large_angle_inds] = A + B
return jac.squeeze_()
def log(self):
if self.mat.dim() < 3:
mat = self.mat.unsqueeze(dim=0)
else:
mat = self.mat
phi = mat.__class__(mat.shape[0], self.dof)
# The cosine of the rotation angle is related to the utils.trace of C
# Clamp to its proper domain to avoid NaNs from rounding errors
cos_angle = (0.5 * utils.trace(mat) - 0.5).clamp_(-1., 1.)
angle = cos_angle.acos()
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_()
if len(small_angle_inds) > 0:
phi[small_angle_inds, :] = \
self.vee(mat[small_angle_inds] -
torch.eye(self.dim).expand_as(mat[small_angle_inds]))
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_()
if len(large_angle_inds) > 0:
angle = angle[large_angle_inds]
sin_angle = angle.sin()
phi[large_angle_inds, :] = \
self.vee(
(0.5 * angle / sin_angle).unsqueeze_(dim=1).unsqueeze_(dim=1).expand_as(mat[large_angle_inds]) *
(mat[large_angle_inds] - mat[large_angle_inds].transpose(2, 1)))
return phi.squeeze_()
def inv_left_jacobian(cls, xi):
if xi.dim() < 2:
xi = xi.unsqueeze(dim=0)
if xi.shape[1] != cls.dof:
raise ValueError("xi must have shape ({},) or (N,{})".format(
cls.dof, cls.dof))
rho = xi[:, :3] # translation part
phi = xi[:, 3:] # rotation part
jac = phi.new_empty(phi.shape[0], cls.dof, cls.dof)
angle = phi.norm(p=2, dim=1)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_(dim=1)
if len(small_angle_inds) > 0:
#Create an identity matrix with a tensor type that matches the input
I = phi.new_empty(cls.dof, cls.dof)
torch.eye(cls.dof, out=I)
jac[small_angle_inds] = \
I.expand_as(jac[small_angle_inds]) - \
0.5 * cls.curlywedge(xi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_(dim=1)
if len(large_angle_inds) > 0:
so3_inv_jac = SO3.inv_left_jacobian(phi[large_angle_inds])
if so3_inv_jac.dim() < 3:
so3_inv_jac.unsqueeze_(dim=0)
Q_mat = cls.left_jacobian_Q_matrix(xi[large_angle_inds])
if Q_mat.dim() < 3:
Q_mat.unsqueeze_(dim=0)
zero_block = phi.new_empty(Q_mat.shape).zero_()
inv_jac_Q_inv_jac = so3_inv_jac.bmm(Q_mat).bmm(so3_inv_jac)
jac[large_angle_inds] = torch.cat([
torch.cat([so3_inv_jac, -inv_jac_Q_inv_jac], dim=2),
torch.cat([zero_block, so3_inv_jac], dim=2)
],
dim=1)
return jac.squeeze_()
def inv_left_jacobian(cls, phi):
if phi.dim() < 2:
phi = phi.unsqueeze(dim=0)
if phi.shape[1] != cls.dof:
raise ValueError("phi must have shape ({},) or (N,{})".format(
cls.dof, cls.dof))
jac = phi.__class__(phi.shape[0], cls.dof, cls.dof)
angle = phi.norm(p=2, dim=1)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_()
if len(small_angle_inds) > 0:
jac[small_angle_inds] = \
torch.eye(cls.dof).expand_as(jac[small_angle_inds]) - \
0.5 * cls.wedge(phi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_()
if len(large_angle_inds) > 0:
angle = angle[large_angle_inds]
axis = phi[large_angle_inds] / \
angle.unsqueeze(dim=1).expand(len(angle), cls.dof)
ha = 0.5 * angle # half angle
hacha = ha / ha.tan() # half angle * cot(half angle)
ha.unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
jac[large_angle_inds])
hacha.unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
jac[large_angle_inds])
A = hacha * \
torch.eye(cls.dof).unsqueeze_(
dim=0).expand_as(jac[large_angle_inds])
B = (1. - hacha) * utils.outer(axis, axis)
C = -ha * cls.wedge(axis)
jac[large_angle_inds] = A + B + C
return jac.squeeze_()
def left_jacobian(cls, xi):
if xi.dim() < 2:
xi = xi.unsqueeze(dim=0)
if xi.shape[1] != cls.dof:
raise ValueError("xi must have shape ({},) or (N,{})".format(
cls.dof, cls.dof))
rho = xi[:, :3] # translation part
phi = xi[:, 3:] # rotation part
jac = phi.__class__(phi.shape[0], cls.dof, cls.dof)
angle = phi.norm(p=2, dim=1)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_()
if len(small_angle_inds) > 0:
jac[small_angle_inds] = \
torch.eye(cls.dof).expand_as(jac[small_angle_inds]) + \
0.5 * cls.curlywedge(xi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_()
if len(large_angle_inds) > 0:
so3_jac = SO3.left_jacobian(phi[large_angle_inds])
if so3_jac.dim() < 3:
so3_jac.unsqueeze_(dim=0)
Q_mat = cls.left_jacobian_Q_matrix(xi[large_angle_inds])
if Q_mat.dim() < 3:
Q_mat.unsqueeze_(dim=0)
zero_block = phi.__class__(Q_mat.shape).zero_()
jac[large_angle_inds] = torch.cat([
torch.cat([so3_jac, Q_mat], dim=2),
torch.cat([zero_block, so3_jac], dim=2)
],
dim=1)
return jac.squeeze_()
def left_jacobian(cls, phi):
if phi.dim() < 2:
phi = phi.unsqueeze(dim=0)
if phi.shape[1] != cls.dof:
raise ValueError(
"phi must have shape ({},) or (N,{})".format(cls.dof, cls.dof))
jac = phi.new_empty(phi.shape[0], cls.dof, cls.dof)
angle = phi.norm(p=2, dim=1)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_(dim=1)
if len(small_angle_inds) > 0:
jac[small_angle_inds] = \
torch.eye(cls.dof, dtype=phi.dtype).expand_as(jac[small_angle_inds]) + \
0.5 * cls.wedge(phi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_(dim=1)
if len(large_angle_inds) > 0:
angle = angle[large_angle_inds]
axis = phi[large_angle_inds] / \
angle.unsqueeze(dim=1).expand(len(angle), cls.dof)
s = angle.sin()
c = angle.cos()
A = (s / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
torch.eye(cls.dof, dtype=phi.dtype).unsqueeze_(dim=0).expand_as(
jac[large_angle_inds])
B = (1. - s / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
utils.outer(axis, axis)
C = ((1. - c) / angle).unsqueeze_(dim=1).unsqueeze_(
dim=2).expand_as(jac[large_angle_inds]) * \
cls.wedge(axis.squeeze())
jac[large_angle_inds] = A + B + C
return jac.squeeze_()
def exp(cls, phi):
if phi.dim() < 2:
phi = phi.unsqueeze(dim=0)
if phi.shape[1] != cls.dof:
raise ValueError("phi must have shape ({},) or (N,{})".format(
cls.dof, cls.dof))
mat = phi.__class__(phi.shape[0], cls.dim, cls.dim)
angle = phi.norm(p=2, dim=1)
# Near phi==0, use first order Taylor expansion
small_angle_mask = utils.isclose(angle, 0.)
small_angle_inds = small_angle_mask.nonzero().squeeze_()
if len(small_angle_inds) > 0:
mat[small_angle_inds] = \
torch.eye(cls.dim).expand_as(mat[small_angle_inds]) + \
cls.wedge(phi[small_angle_inds])
# Otherwise...
large_angle_mask = 1 - small_angle_mask # element-wise not
large_angle_inds = large_angle_mask.nonzero().squeeze_()
if len(large_angle_inds) > 0:
angle = angle[large_angle_inds]
axis = phi[large_angle_inds] / \
angle.unsqueeze(dim=1).expand(len(angle), cls.dim)
s = angle.sin().unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
mat[large_angle_inds])
c = angle.cos().unsqueeze_(dim=1).unsqueeze_(dim=2).expand_as(
mat[large_angle_inds])
A = c * torch.eye(cls.dim).unsqueeze_(dim=0).expand_as(
mat[large_angle_inds])
B = (1. - c) * utils.outer(axis, axis)
C = s * cls.wedge(axis)
mat[large_angle_inds] = A + B + C
return cls(mat.squeeze_())
def test_isclose():
tol = 1e-6
mat = torch.Tensor([0, 1, tol, 10 * tol, 0.1 * tol])
ans = torch.ByteTensor([1, 0, 0, 0, 1])
assert (utils.isclose(mat, 0., tol=tol) == ans).all()
def to_quaternion(self, ordering='wxyz'):
"""Convert a rotation matrix to a unit length quaternion.
Valid orderings are 'xyzw' and 'wxyz'.
"""
if self.mat.dim() < 3:
R = self.mat.unsqueeze(dim=0)
else:
R = self.mat
qw = 0.5 * torch.sqrt(1. + R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2])
qx = qw.__class__(qw.shape)
qy = qw.__class__(qw.shape)
qz = qw.__class__(qw.shape)
near_zero_mask = utils.isclose(qw, 0.)
if sum(near_zero_mask) > 0:
cond1_mask = near_zero_mask & \
(R[:, 0, 0] > R[:, 1, 1]).squeeze_() & \
(R[:, 0, 0] > R[:, 2, 2]).squeeze_()
cond1_inds = cond1_mask.nonzero().squeeze_()
if len(cond1_inds) > 0:
R_cond1 = R[cond1_inds]
d = 2. * np.sqrt(1. + R_cond1[:, 0, 0] - R_cond1[:, 1, 1] -
R_cond1[:, 2, 2])
qw[cond1_inds] = (R_cond1[:, 2, 1] - R_cond1[:, 1, 2]) / d
qx[cond1_inds] = 0.25 * d
qy[cond1_inds] = (R_cond1[:, 1, 0] + R_cond1[:, 0, 1]) / d
qz[cond1_inds] = (R_cond1[:, 0, 2] + R_cond1[:, 2, 0]) / d
cond2_mask = near_zero_mask & (R[:, 1, 1] > R[:, 2, 2]).squeeze_()
cond2_inds = cond2_mask.nonzero().squeeze_()
if len(cond2_inds) > 0:
R_cond2 = R[cond2_inds]
d = 2. * np.sqrt(1. + R_cond2[:, 1, 1] - R_cond2[:, 0, 0] -
R_cond2[:, 2, 2])
qw[cond2_inds] = (R_cond2[:, 0, 2] - R_cond2[:, 2, 0]) / d
qx[cond2_inds] = (R_cond2[:, 1, 0] + R_cond2[:, 0, 1]) / d
qy[cond2_inds] = 0.25 * d
qz[cond2_inds] = (R_cond2[:, 2, 1] + R_cond2[:, 1, 2]) / d
cond3_mask = near_zero_mask & (1 - cond1_mask) & (1 - cond2_mask)
cond3_inds = cond3_mask.nonzero().squeeze_()
if len(cond3_inds) > 0:
R_cond3 = R[cond3_inds]
d = 2. * \
np.sqrt(1. + R_cond3[:, 2, 2] -
R_cond3[:, 0, 0] - R_cond3[:, 1, 1])
qw[cond3_inds] = (R_cond3[:, 1, 0] - R_cond3[:, 0, 1]) / d
qx[cond3_inds] = (R_cond3[:, 0, 2] + R_cond3[:, 2, 0]) / d
qy[cond3_inds] = (R_cond3[:, 2, 1] + R_cond3[:, 1, 2]) / d
qz[cond3_inds] = 0.25 * d
far_zero_mask = 1 - near_zero_mask
far_zero_inds = far_zero_mask.nonzero().squeeze_()
if len(far_zero_inds) > 0:
R_fz = R[far_zero_inds]
d = 4. * qw[far_zero_inds]
qx[far_zero_inds] = (R_fz[:, 2, 1] - R_fz[:, 1, 2]) / d
qy[far_zero_inds] = (R_fz[:, 0, 2] - R_fz[:, 2, 0]) / d
qz[far_zero_inds] = (R_fz[:, 2, 1] - R_fz[:, 1, 2]) / d
# Check ordering last
if ordering is 'xyzw':
quat = torch.cat([
qx.unsqueeze_(dim=1),
qy.unsqueeze_(dim=1),
qz.unsqueeze_(dim=1),
qw.unsqueeze_(dim=1)
],
dim=1).squeeze_()
elif ordering is 'wxyz':
quat = torch.cat([
qw.unsqueeze_(dim=1),
qx.unsqueeze_(dim=1),
qy.unsqueeze_(dim=1),
qz.unsqueeze_(dim=1)
],
dim=1).squeeze_()
else:
raise ValueError(
"Valid orderings are 'xyzw' and 'wxyz'. Got '{}'.".format(
ordering))
return quat
