liealg.py

def hat(w):
     return torch.tensor([[0, -w[2], w[1]],
                    [w[2], 0, -w[0]],
                    [-w[1], w[0], 0]])


def hat_inv(R):
     return torch.tensor([R[2][1], R[0][2], R[1][0]])


def exp_so3(w):
     theta = torch.norm(w)
     t1 = torch.eye(3)
     t2 = (torch.sin(theta) / theta) * hat(w)
     t3 = ((1 - torch.cos(theta)) / (theta * theta)) * torch.mm(hat(w), hat(w))
     return t1 + t2 + t3


def log_so3(R):
     theta = torch.acos((torch.trace(R) - 1) / 2)
     skew_sym = (theta / (2 * torch.sin(theta))) * (R - torch.transpose(R, 0, 1))
     return hat_inv(skew_sym)


def exp_se3(v):
     t = torch.tensor([v[0], v[1], v[2]], dtype = torch.float)
     w = torch.tensor([v[3], v[4], v[5]], dtype = torch.float)

     theta = torch.norm(w)
     t1 = torch.eye(3)
     t2 = ((1 - torch.cos(theta)) / (theta * theta)) * hat(w)
     t3 = ((theta - torch.sin(theta)) / (theta * theta * theta)) * torch.mm(hat(w), hat(w))
     V = t1 + t2 + t3
     Vt = torch.unsqueeze(torch.mv(V, t), 1)

     top = torch.cat((exp_so3(w), Vt), 1)
     bottom = torch.tensor([[0, 0, 0, 1]], dtype = torch.float)
     return torch.cat((top, bottom), 0)


def log_se3(A):
     R = torch.narrow(A, 0, 0, 3)
     t = torch.squeeze(torch.narrow(R, 1, 3, 1))
     R = torch.narrow(R, 1, 0, 3)

     w = log_so3(R)
     theta = torch.norm(w)
     
     t1 = torch.eye(3)
     t2 = -0.5 * hat(w)
     t3 = ((1 - ((theta * torch.cos(theta / 2)) / (2 * torch.sin(theta / 2)))) / (theta * theta)) * torch.mm(hat(w), hat(w))
     V_ = t1 + t2 + t3

     t_ = torch.mv(V_, t)

     return torch.cat((t_, w), 0)

Add epslions for stability