def test_integ():
g = DiffeoGroup((16, 16))
dim = 2
vel = torch.randn(dim, dtype=torch.float64)
velocity = vel.reshape(-1, 1, 1) * torch.ones(
(dim, 16, 16), dtype=torch.float64)
g.exponential(velocity)
def test_identity():
small_shape = [16] * 2
group = DiffeoGroup(small_shape)
vol_16 = get_density_action(small_shape)
idall_ = group.element()
x = torch.from_numpy(np.random.randn(*small_shape))
res = vol_16(x, idall_)
assert pytest.approx(torch.max((x - res).abs()).numpy()) == 0
def test_function():
small_shape = [16] * 2
group = DiffeoGroup(small_shape)
comp_16 = get_composition_action(small_shape, compute_id=True)
idall = group.get_raw_identity(requires_grad=True)
I16 = torch.randn(*small_shape, dtype=torch.float64)
torch.autograd.gradcheck(comp_16, inputs=(I16, idall))
def test_grad_loss():
small_shape = [16] * 2
group = DiffeoGroup(small_shape)
comp_16 = get_composition_action(small_shape)
idall = group.get_raw_identity(requires_grad=True)
I16 = torch.randn(*small_shape, dtype=torch.float64)
loss = mse(comp_16(I16, idall),
torch.zeros(small_shape, dtype=torch.float64))
torch.autograd.backward(loss)
def test_density():
small_shape = [16] * 2
group = DiffeoGroup(small_shape)
vol_16 = get_density_action(small_shape, compute_id=True)
idall = group.get_raw_identity(requires_grad=True)
x = torch.randn(*small_shape, dtype=torch.float64)
torch.autograd.gradcheck(
vol_16,
inputs=(x, idall),
atol=1e-10,
)
def test_grad_img():
small_shape = [16] * 2
comp_16 = get_composition_action(small_shape)
group = DiffeoGroup(small_shape)
v = group.zero()
v[0] = 1.5
v[1] = .5
elt = group.exponential(v)
I16 = torch.ones(*small_shape, dtype=torch.float64)
I16.requires_grad = True
torch.autograd.gradcheck(comp_16, inputs=(I16, elt))
def test_one_jacobian():
"""
Jacobian of translation is one.
"""
shape = [16] * 2
group = DiffeoGroup(shape)
act = get_density_action(shape)
idall = group.element()
vel = group.zero()
vel[0] += 3
trans = group.exponential(vel)
x = torch.ones(shape, dtype=torch.float64)
y = act(x, trans)
assert pytest.approx(x.numpy()) == y.numpy()
def test_deep_optimisation():
shape = (16, 16)
I0, I1 = [1 + torch.rand(*shape, dtype=torch.float64) for i in range(2)]
group = DiffeoGroup(I0.shape)
cometric = get_laplace_cometric(group, s=2)
seq = Sequential(*[FunctionRepresentation(group) for i in range(3)])
do = VelocityOptimizer(seq.parameters(),
lr=.1,
cometric=cometric,
weight_decay=1.)
for i in range(2):
do.zero_grad()
I_ = seq(I0)
loss = mse(I_, I1)
loss.backward()
do.step()
def test_orbit_optimisation():
shape = (16, 16)
I0, I1 = [1 + torch.rand(*shape, dtype=torch.float64) for i in range(2)]
group = DiffeoGroup(I0.shape)
cometric = get_laplace_cometric(group, s=2)
sum_rep = get_sum_representation(FunctionRepresentation(group),
DensityRepresentation(group))
oo = GroupOptimizer(sum_rep.parameters(), lr=.1, cometric=cometric)
vol = normalize(torch.ones_like(I1))
vol__ = vol + 1e-2 * torch.randn_like(vol)
for i in range(2):
oo.zero_grad()
I_, vol_ = sum_rep(I0, vol)
loss = mse(I_, I0) + information_distance(vol_, vol__)
loss.backward()
oo.step()
def test_element():
g = DiffeoGroup((16, 16))
g.element()
def test_inverse():
group = DiffeoGroup((16, 16))
defm = get_random_diffeo(group, scale=.1)
aid = defm.compose(defm.inverse())
rid = group.get_raw_identity()
assert np.allclose(rid.numpy(), aid.forward.numpy(), rtol=1e-9, atol=1e-1)
def test_exponential():
g = DiffeoGroup((16, 16))
z = torch.zeros((2, 16, 16), dtype=torch.float64)
g.exponential(z)
def test_compose():
g = DiffeoGroup((16, 16))
id1 = g.element()
id2 = g.element()
id3 = id1.compose(id2)
assert torch.allclose(id1.forward, id3.forward)