def forward(self):
if self.lie_multiply:
alg_distr = Normal(torch.zeros(3).double(), self.scale)
loc = so3_exp(self.loc)
transforms = [self.transform, SO3MultiplyTransform(loc)]
else:
alg_distr = Normal(self.loc * 1, self.scale)
transforms = [self.transform]
return LDTD(alg_distr, transforms)
def se3_exp(v):
"""
Exponential map of SE(3) with Rordigues formula.
:param v: algebra vector of shape (..., 6)
:return: group element of shape (..., 4, 4)
"""
v_so3, v_r = v.split([3, 3], dim=-1)
so3 = so3_exp(v_so3)
theta = v_so3.norm(p=2, dim=-1, keepdim=True)
k = so3_hat(v_so3)
eye = torch.eye(3, device=v.device, dtype=v.dtype)
V = (eye + ((1 - torch.cos(theta)) / theta**2)[..., None] * k +
((theta - torch.sin(theta)) / theta**3)[..., None] * (k @ k))
r3 = V @ v_r.unsqueeze(-1)
return se3_fill(so3, r3.squeeze(-1), "group")
# Distribution to create noisy observations: p(G'|G)=n @ G, n \sim exp^*(N(0, .1))
noise_alg_distr = Normal(
torch.zeros(3).double().to(device),
torch.full((3, ), 0.1).double().to(device))
noise_group_distr = LDTD(noise_alg_distr, SO3ExpTransform())
# Sample true and noisy group actions
num_samples = 100_000
noise_samples = noise_group_distr.sample((num_samples, ))
# Make symmetrical
symmetry_group_size = 3
symmetry_group = so3_exp(
torch.tensor(
[[2 * np.pi / symmetry_group_size * i, 0, 0]
for i in range(symmetry_group_size)],
device=device,
).double())
group_data = symmetry_group.repeat(num_samples // symmetry_group_size, 1, 1)
group_data_noised = noise_samples[:len(group_data)] @ group_data
dataset = TensorDataset(group_data_noised)
loader = TensorLoader(dataset, 640, True)
loader_iter = cycle(loader)
class Flow(nn.Module):
def __init__(self, d, n_layers):
super().__init__()
self.d = d
from torch.distributions import Normal
from torch.utils.data import TensorDataset
from relie import LocalDiffeoTransformedDistribution, SO3ExpTransform
from relie.utils.data import TensorLoader, cycle
from relie.utils.so3_tools import so3_exp
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device = torch.device("cpu")
target_alg_distr = Normal(
torch.tensor([0.0, 0.5, -0.5]).double(),
torch.tensor([1.5, 0.1, 1.0]).double())
target_group_distr = LocalDiffeoTransformedDistribution(
target_alg_distr, SO3ExpTransform(k_max=10))
data = so3_exp(target_alg_distr.sample((10000, )))
target_entropy = -target_group_distr.log_prob(data).mean()
dataset = TensorDataset(data)
loader = TensorLoader(dataset, 64, True)
loader_iter = cycle(loader)
loc = nn.Parameter(torch.zeros(3).double())
scale = nn.Parameter(torch.ones(3).double())
optimizer = torch.optim.Adam([loc, scale])
transform = SO3ExpTransform(k_max=10)
for it in range(100_000):
batch, = next(loader_iter)
alg_distr = Normal(loc * 1, scale * 1)
def _call(self, x):
return so3_exp(x)
def gen_data(symmetry_group_size=3, noise=0.1, num_samples=100_000):
# Prior distribution p(G)
generation_group_distr = SO3Prior(dtype=torch.double, device=device)
# Distribution to create noisy observations: p(G'|G)=n @ G, n \sim exp^*(N(0, .1))
noise_alg_distr = Normal(
torch.zeros(3).double().to(device), torch.full((3,), noise).double().to(device)
)
noise_group_distr = LDTD(noise_alg_distr, SO3ExpTransform())
# Sample true and noisy group actions
noise_samples = noise_group_distr.sample((num_samples,))
group_data = generation_group_distr.sample((num_samples,))
group_data_noised = noise_samples @ group_data
# Create original data
max_rep_degree = 3
rep_copies = 1
x_dims = (max_rep_degree + 1) ** 2 * rep_copies
# x_zero = torch.randn((max_rep_degree + 1)**2, rep_copies, device=device)
x_zero = [
0.974_945_008_754_730_2,
1.320_214_033_126_831,
2.740_009_784_698_486_3,
0.629_877_269_268_035_9,
0.859_345_138_072_967_5,
0.679_915_964_603_424_1,
-0.487_856_239_080_429_1,
1.209_429_860_115_051_3,
0.009_278_437_122_702_599,
-1.521_781_563_758_85,
1.634_827_971_458_435,
-1.268_676_042_556_762_7,
1.858_604_192_733_764_6,
1.052_274_703_979_492_2,
-0.713_051_140_308_380_1,
0.341_978_967_189_788_8,
]
x_zero = torch.tensor(x_zero, device=device)[:, None]
# Make symmetrical
symmetry_group = so3_exp(
torch.tensor(
[
[2 * np.pi / symmetry_group_size * i, 0, 0]
for i in range(symmetry_group_size)
],
device=device,
).double()
)
x_data = block_wigner_matrix_multiply(
so3_matrix_to_eazyz(symmetry_group).float(),
x_zero.expand(symmetry_group_size, -1, -1),
max_rep_degree,
)
x_zero = x_data.mean(0)
# Act with group
angles = so3_matrix_to_eazyz(group_data_noised)
x_data = block_wigner_matrix_multiply(
angles.float(), x_zero.expand(num_samples, -1, -1), max_rep_degree
)
dataset = TensorDataset(x_data, group_data, group_data)
loader = TensorLoader(dataset, 640, True)
loader_iter = cycle(loader)
return Bunch(
loader_iter=loader_iter,
x_dims=x_dims,
generation_group_distr=generation_group_distr,
symmetry_group=symmetry_group,
x_zero=x_zero,
max_rep_degree=max_rep_degree,
)