def inner(self,
z: torch.Tensor,
u: torch.Tensor,
v=None,
*,
keepdim=False) -> torch.Tensor:
"""
Inner product for tangent vectors at point :math:`z`.
For the upper half space model, the inner product at point z = x + iy of the vectors u, v
it is (z, u, v are complex symmetric matrices):
g_{z}(u, v) = tr[ y^-1 u y^-1 \ov{v} ]
:param z: torch.Tensor point on the manifold: b x 2 x n x n
:param u: torch.Tensor tangent vector at point :math:`z`: b x 2 x n x n
:param v: Optional[torch.Tensor] tangent vector at point :math:`z`: b x 2 x n x n
:param keepdim: bool keep the last dim?
:return: torch.Tensor inner product (broadcastable): b x 2 x 1 x 1
"""
if v is None:
v = u
inv_imag_z = torch.inverse(sm.imag(z))
inv_imag_z = sm.stick(inv_imag_z, torch.zeros_like(inv_imag_z))
res = sm.bmm3(inv_imag_z, u, inv_imag_z)
res = sm.bmm(res, sm.conjugate(v))
real_part = sm.trace(sm.real(res), keepdim=True) # b x 1
real_part = torch.unsqueeze(real_part, -1) # b x 1 x 1
return sm.stick(real_part, real_part) # b x 2 x 1 x 1
def test_takagi_factorization_real_neg_imag_neg(self):
a = get_random_symmetric_matrices(3, 3)
a = sm.stick(sm.real(a) * -1, sm.imag(a) * -1)
eigenvalues, s = TakagiFactorization(3).factorize(a)
diagonal = torch.diag_embed(eigenvalues)
diagonal = sm.stick(diagonal, torch.zeros_like(diagonal))
self.assertAllClose(
a, sm.bmm3(sm.conjugate(s), diagonal, sm.conj_trans(s)))
def test_distance_is_symmetric_only_imaginary_matrices(self):
x = self.manifold.random(10)
y = self.manifold.random(10)
zeros = torch.zeros_like(sm.real(x))
x = sm.stick(zeros, sm.imag(x))
y = sm.stick(zeros, sm.imag(y))
dist_xy = self.manifold.dist(x, y)
dist_yx = self.manifold.dist(y, x)
self.assertAllClose(dist_xy, dist_yx)
def test_pow_square(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[-80, -112], [-171, -176]]])
expected_imag = torch.Tensor([[[18, 66], [-140, -210]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.pow(x, 2)
self.assertAllClose(expected, result)
def test_pow_cube(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[-242, 1062], [-2815, 4382]]])
expected_imag = torch.Tensor([[[-702, 1034], [1694, -1170]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.pow(x, 3)
self.assertAllClose(expected, result)
def test_pow_square_root(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[2.24225167, 2.04960413],
[3.15167167, 2.18551428]]])
expected_imag = torch.Tensor([[[2.00691120, -2.68344501],
[-2.22104353, 3.43168656]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.pow(x, 0.5)
self.assertAllClose(expected, result)
def test_pow_inverse(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[0.01219512, -0.02307692],
[0.02262443, -0.02554744]]])
expected_imag = torch.Tensor([[[-0.10975609, 0.08461538],
[0.06334841, -0.05474452]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.pow(x, -1)
self.assertAllClose(expected, result)
def test_hermitian_matrix_sqrt(self):
# expected result from https://www.wolframalpha.com/
x_real = torch.Tensor([[[0.9408, 0.1332], [0.1332, 0.5936]]])
x_imag = torch.Tensor([[[0.0000, 0.5677], [-0.5677, 0.0000]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[0.896153, 0.0847679],
[0.0847679, 0.675196]]])
expected_imag = torch.Tensor([[[0, 0.361282], [-0.361282, 0]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.hermitian_matrix_sqrt(x)
self.assertAllClose(expected, result, rtol=1e-05, atol=1e-06)
def test_inverse_symmetric_2d(self):
# expected result from https://adrianstoll.com/linear-algebra/matrix-inversion.html
x_real = torch.Tensor([[[1, -3], [-3, 7]]])
x_imag = torch.Tensor([[[-9, 11], [11, 15]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor([[[256 / 8105, 141 / 16210],
[141 / 16210, 23 / 16210]]])
expected_imag = torch.Tensor([[[921 / 16210, -356 / 8105],
[-356 / 8105, -288 / 8105]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.inverse(x)
self.assertAllClose(expected, result)
def projx(self, z: torch.Tensor) -> torch.Tensor:
"""
Project point :math:`z` on the manifold.
In this space, we need to ensure that Y = Im(X) is positive definite.
Since the matrix Y is symmetric, it is possible to diagonalize it.
For a diagonal matrix the condition is just that all diagonal entries are positive, so we clamp the values
that are <=0 in the diagonal to an epsilon, and then restore the matrix back into non-diagonal form using
the base change matrix that was obtained from the diagonalization.
Steps to project: Y = Im(z)
1) Y = SDS^-1
2) D_tilde = clamp(D, min=epsilon)
3) Y_tilde = SD_tildeS^-1
:param z: points to be projected: (b, 2, n, n)
"""
z = super().projx(z)
y = sm.imag(z)
y_tilde, batchwise_mask = sm.positive_conjugate_projection(y)
self.projected_points += len(z) - sum(batchwise_mask).item()
return sm.stick(sm.real(z), y_tilde)
def test_bmm(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
y_real = torch.Tensor([[[9, -11], [-14, 15]]])
y_imag = torch.Tensor([[[1, -3], [5, -7]]])
y = sm.stick(y_real, y_imag)
expected_real = torch.Tensor([[[97, -106], [82, -97]]])
expected_imag = torch.Tensor([[[221, -246], [-366, 413]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.bmm(x, y)
self.assertAllClose(expected, result)
def symmetric_svd_with_eigenvectors(self, z: torch.Tensor):
"""
:param z: complex symmetric matrix
:return:
"""
compound_z = sm.to_compound_symmetric(
z) # b x 2 x 2n x 2n. Z = A + iB, then [(A, B),(B, -A)]
evalues, q = self.symeig(
compound_z,
eigenvectors=True) # b x n in ascending order, b x 2n x 2n
# I can think of Q as 4 n x n matrices.
# Q = [(X, Re(U)),
# (Y, -Im(U))] where X, Y are irrelevant and I need to build U
real_u_on_top_of_minus_imag_u = torch.chunk(q, 2, dim=-1)[-1]
real_u, minus_imag_u = torch.chunk(real_u_on_top_of_minus_imag_u,
2,
dim=-2)
u = sm.stick(real_u, -minus_imag_u) # b x 2 x n x n
sing_values = evalues[:, z.shape[-1]:]
# sing_values_matrix = sm.bmm3(sm.transpose(u), z, u) # b x 2 x n x n. imag part should be zero
# sing_values = torch.diagonal(sm.real(sing_values_matrix), offset=0, dim1=-2, dim2=-1)
return sing_values, u
def test_stick(self):
x_real = torch.rand(10, 2, 4, 4)
x_imag = torch.rand(10, 2, 4, 4)
x = sm.stick(x_real, x_imag)
self.assertAllEqual(x_real, sm.real(x))
self.assertAllEqual(x_imag, sm.imag(x))
def test_distance_is_symmetric_real_neg_imag_pos(self):
x = self.manifold.random(10)
x = sm.stick(sm.real(x) * -1, sm.imag(x))
y = self.manifold.random(10)
dist_xy = self.manifold.dist(x, y)
dist_yx = self.manifold.dist(y, x)
self.assertAllClose(dist_xy, dist_yx)
def test_takagi_factorization_very_large_values(self):
a = get_random_symmetric_matrices(3, 3) * 1000
eigenvalues, s = TakagiFactorization(3).factorize(a)
diagonal = torch.diag_embed(eigenvalues)
diagonal = sm.stick(diagonal, torch.zeros_like(diagonal))
self.assertAllClose(
a, sm.bmm3(sm.conjugate(s), diagonal, sm.conj_trans(s)))
def test_inverse_symmetric_3d(self):
# expected result from https://adrianstoll.com/linear-algebra/matrix-inversion.html
x_real = torch.Tensor([[[-1, -3, 9], [-3, 5, 7], [9, 7, 11]]])
x_imag = torch.Tensor([[[9, 4, -6], [4, 7, 9], [-6, 9, -3]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor(
[[[-36251 / 845665, -27631 / 845665, 188 / 9949],
[-27631 / 845665, 251611 / 3382660, -689 / 39796],
[188 / 9949, -689 / 39796, 1299 / 39796]]])
expected_imag = torch.Tensor(
[[[-18757 / 845665, -35642 / 845665, 532 / 9949],
[-35642 / 845665, -112703 / 3382660, -1103 / 39796],
[532 / 9949, -1103 / 39796, 289 / 39796]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.inverse(x)
self.assertAllClose(expected, result)
def test_inverse_nonsymmetric_3d(self):
# expected result from https://adrianstoll.com/linear-algebra/matrix-inversion.html
x_real = torch.Tensor([[[-1, -3, 9], [3, 5, 7], [2, 9, 11]]])
x_imag = torch.Tensor([[[9, 4, -6], [-4, 7, 9], [-2, 7, -3]]])
x = sm.stick(x_real, x_imag)
expected_real = torch.Tensor(
[[[951 / 16589, 3223 / 16589, -4496 / 16589],
[5029 / 66356, 2486 / 16589, 1387 / 33178],
[2137 / 66356, 1030 / 16589, -1828 / 16589]]])
expected_imag = torch.Tensor(
[[[-3143 / 16589, -3622 / 16589, 1532 / 16589],
[5533 / 66356, 6925 / 33178, -17977 / 66356],
[-4167 / 66356, -5779 / 33178, 6565 / 66356]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.inverse(x)
self.assertAllClose(expected, result)
def test_to_hermitian_from_compound_real_symmetric_2d(self):
x = torch.Tensor([[[1, -3, 0, -5], [-3, 7, 5, 0], [0, 5, 1, -3],
[-5, 0, -3, 7]]])
expected_real = torch.Tensor([[[1, -3], [-3, 7]]])
expected_imag = torch.Tensor([[[0, 5], [-5, 0]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.to_hermitian_from_compound_real_symmetric(x)
self.assertAllEqual(expected, result)
def test_bmm3(self):
x_real = torch.Tensor([[[1, -3], [5, -7]]])
x_imag = torch.Tensor([[[9, -11], [-14, 15]]])
x = sm.stick(x_real, x_imag)
y_real = torch.Tensor([[[9, -11], [-14, 15]]])
y_imag = torch.Tensor([[[1, -3], [5, -7]]])
y = sm.stick(y_real, y_imag)
z_real = torch.Tensor([[[-3, -1], [-2, 5]]])
z_imag = torch.Tensor([[[-1, 3], [0, -2]]])
z = sm.stick(z_real, z_imag)
expected_real = torch.Tensor([[[142, -1782], [-418, 1357]]])
expected_imag = torch.Tensor([[[-268, -948], [190, 2871]]])
expected = sm.stick(expected_real, expected_imag)
result = sm.bmm3(x, y, z)
self.assertAllClose(expected, result)
def test_hermitian_eig(self):
# expected result from https://www.arndt-bruenner.de/mathe/scripts/engl_eigenwert2.htm
x_real = torch.Tensor([[[1, -3], [-3, 7]]])
x_imag = torch.Tensor([[[0, 5], [-5, 0]]])
x = sm.stick(x_real, x_imag)
expected = torch.Tensor([[[-2.55743852, 10.55743852]]])
_, _, result = sm.hermitian_eig(x)
self.assertAllClose(expected, result)
def test_takagi_factorization_real_identity(self):
a = sm.identity_like(get_random_symmetric_matrices(3, 3))
eigenvalues, s = TakagiFactorization(3).factorize(a)
diagonal = torch.diag_embed(eigenvalues)
diagonal = sm.stick(diagonal, torch.zeros_like(diagonal))
self.assertAllClose(
a, sm.bmm3(sm.conjugate(s), diagonal, sm.conj_trans(s)))
self.assertAllClose(a, s)
self.assertAllClose(torch.ones_like(eigenvalues), eigenvalues)
def test_proj_x_real_neg_imag_neg(self):
x = get_random_symmetric_matrices(10, self.dims)
x = sm.stick(sm.real(x) * -1, sm.imag(x) * -1)
proj_x = self.manifold.projx(x)
# assert symmetry
self.assertAllClose(proj_x, sm.transpose(proj_x))
# assert all points belong to the manifold
for point in proj_x:
self.assertTrue(self.manifold.check_point_on_manifold(point))
def test_distance_is_symmetric_with_diagonal_matrices(self):
x = self.manifold.random(10)
y = self.manifold.random(10)
diagonal_mask = torch.eye(self.dims).unsqueeze(0).repeat(10, 1,
1).bool()
diagonal_mask = sm.stick(diagonal_mask, diagonal_mask)
x = torch.where(diagonal_mask, x, torch.zeros_like(x))
y = torch.where(diagonal_mask, y, torch.zeros_like(y))
dist_xy = self.manifold.dist(x, y)
dist_yx = self.manifold.dist(y, x)
self.assertAllClose(dist_xy, dist_yx)
def test_conj_transpose(self):
x = get_random_symmetric_matrices(10, 4)
x_real = sm.real(x)
x_imag = sm.imag(x)
x_real_transpose = x_real.transpose(-1, -2)
x_imag_transpose = x_imag.transpose(-1, -2)
x_expected_conj_transpose = sm.stick(x_real_transpose,
x_imag_transpose * -1)
x_result_conj_transpose = sm.conj_trans(x)
self.assertAllEqual(x_expected_conj_transpose, x_result_conj_transpose)
def test_transpose(self):
x = get_random_symmetric_matrices(10, 4)
x_real = sm.real(x)
x_imag = sm.imag(x)
x_real_transpose = x_real.transpose(-1, -2)
x_imag_transpose = x_imag.transpose(-1, -2)
x_expected_transpose = sm.stick(x_real_transpose, x_imag_transpose)
x_result_transpose = sm.transpose(x)
self.assertAllEqual(x_expected_transpose, x_result_transpose)
self.assertAllEqual(
x, x_result_transpose) # because they are symmetric matrices
def test_cayley_transform_real_neg_imag_pos(self):
x = self.upper_half_manifold.random(10)
x = sm.stick(sm.real(x) * -1, sm.imag(x))
tran_x = cayley_transform(x)
result = inverse_cayley_transform(tran_x)
self.assertAllClose(x, result)
# the intermediate result belongs to the Bounded domain manifold
for point in tran_x:
self.assertTrue(self.bounded_manifold.check_point_on_manifold(point))
# the final result belongs to the Upper Half Space manifold
for point in result:
self.assertTrue(self.upper_half_manifold.check_point_on_manifold(point))
def cayley_transform(z: torch.Tensor) -> torch.Tensor:
"""
T(Z): Upper Half Space model -> Bounded Domain Model
T(Z) = (Z - i Id)(Z + i Id)^-1
:param z: b x 2 x n x n: PRE: z \in Upper Half Space Manifold
:return: y \in Bounded Domain Manifold
"""
identity = sm.identity_like(z)
i_identity = sm.stick(sm.imag(identity), sm.real(identity))
z_minus_id = sm.subtract(z, i_identity)
inv_z_plus_id = sm.inverse(sm.add(z, i_identity))
return sm.bmm(z_minus_id, inv_z_plus_id)
def test_takagi_factorization_real_diagonal(self):
a = get_random_symmetric_matrices(3, 3) * 10
a = torch.where(sm.identity_like(a) == 1, a, torch.zeros_like(a))
eigenvalues, s = TakagiFactorization(3).factorize(a)
diagonal = torch.diag_embed(eigenvalues)
diagonal = sm.stick(diagonal, torch.zeros_like(diagonal))
self.assertAllClose(
a, sm.bmm3(sm.conjugate(s), diagonal, sm.conj_trans(s)))
# real part of eigenvectors is made of vectors with one 1 and all zeros
real_part = torch.sum(torch.abs(sm.real(s)), dim=-1)
self.assertAllClose(torch.ones_like(real_part), real_part)
# imaginary part of eigenvectors is all zeros
self.assertAllClose(torch.zeros(1), torch.sum(sm.imag(s)))
def egrad2rgrad(self, z: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
"""
Transform gradient computed using autodiff to the correct Riemannian gradient for the point :math:`x`.
If you have a function f(z) on Hn, then the gradient is the y * grad_eucl(f(z)) * y,
where y is the imaginary part of z, and multiplication is just matrix multiplication.
:param z: point on the manifold. Shape: (b, 2, n, n)
:param u: gradient to be projected: Shape: same than z
:return grad vector in the Riemannian manifold. Shape: same than z
"""
real_grad, imag_grad = sm.real(u), sm.imag(u)
y = sm.imag(z)
real_grad = y.bmm(real_grad).bmm(y)
imag_grad = y.bmm(imag_grad).bmm(y)
return sm.stick(real_grad, imag_grad)
def random(self, *size, dtype=None, device=None, **kwargs) -> torch.Tensor:
"""
Random sampling on the manifold.
The exact implementation depends on manifold and usually does not follow all
assumptions about uniform measure, etc.
"""
from_ = kwargs.get("from_", -0.001)
to = kwargs.get("to", 0.001)
dims = self.dims
perturbation = sm.squared_to_symmetric(
torch.Tensor(size[0], dims, dims).uniform_(from_, to))
identity = torch.eye(dims).unsqueeze(0).repeat(size[0], 1, 1)
imag_part = identity + perturbation
real_part = sm.squared_to_symmetric(
torch.Tensor(size[0], dims, dims).uniform_(from_, to))
return sm.stick(real_part, imag_part).to(device=device, dtype=dtype)
