def retraction(self, u, vec):
"""Transports a set of points from the complex Stiefel
manifold via a retraction map.
Args:
u: complex valued tensor of shape (..., n, p), a set
of points to be transported.
vec: complex valued tensor of shape (..., n, p),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., n, p),
a set of transported points."""
if self._retraction == 'svd':
new_u = u + vec
_, v, w = tf.linalg.svd(new_u)
return v @ adj(w)
elif self._retraction == 'cayley':
W = vec @ adj(u) - 0.5 * u @ (adj(u) @ vec @ adj(u))
W = W - adj(W)
Id = tf.eye(W.shape[-1], dtype=W.dtype)
return tf.linalg.inv(Id - W / 2) @ (Id + W / 2) @ u
elif self._retraction == 'qr':
new_u = u + vec
q, r = tf.linalg.qr(new_u)
diag = tf.linalg.diag_part(r)
sign = tf.math.sign(diag)[..., tf.newaxis, :]
return q * sign
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., m, n, n),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., m, n, n),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., m, n, n),
the set of Reimannian gradients."""
n = u.shape[-1]
m = u.shape[-3]
shape = u.shape[:-3]
size = len(shape)
idx = tuple(range(size))
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.transpose(egrad, idx + (size + 2, size, size + 1))
vec_mod = tf.reshape(vec_mod, shape + (n * m, n))
u_mod = tf.transpose(u, idx + (size + 2, size, size + 1))
u_mod = tf.reshape(u_mod, shape + (n * m, n))
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.reshape(vec_mod, shape + (n, m, n))
vec_mod = tf.transpose(vec_mod, idx + (size + 1, size + 2, size))
return vec_mod
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n ** 2, k),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., n ** 2, k),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n ** 2, k),
the set of Reimannian gradients."""
k = u.shape[-1]
n = int(math.sqrt(u.shape[-2]))
shape = u.shape[:-2]
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.reshape(egrad, shape + (n, k * n))
vec_mod = tf.linalg.matrix_transpose(vec_mod)
u_mod = tf.reshape(u, shape + (n, k * n))
u_mod = tf.linalg.matrix_transpose(u_mod)
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.linalg.matrix_transpose(vec_mod)
vec_mod = tf.reshape(vec_mod, shape + (n ** 2, k))
return vec_mod
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n ** 2, k),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., n ** 2, k),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n ** 2, k),
the set of Reimannian gradients.
Note:
The complexity O(kn^3)"""
k = tf.shape(u)[-1]
n = tf.cast(tf.math.sqrt(tf.cast(tf.shape(u)[-2], dtype=tf.float32)),
dtype=tf.int32)
shape = tf.shape(u)[:-2]
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.reshape(
egrad, shape_conc(shape, n[tf.newaxis], (k * n)[tf.newaxis]))
vec_mod = tf.linalg.matrix_transpose(vec_mod)
u_mod = tf.reshape(
u, shape_conc(shape, n[tf.newaxis], (k * n)[tf.newaxis]))
u_mod = tf.linalg.matrix_transpose(u_mod)
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.linalg.matrix_transpose(vec_mod)
vec_mod = tf.reshape(
vec_mod, shape_conc(shape, (n**2)[tf.newaxis], k[tf.newaxis]))
return vec_mod
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., n, n),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n, n),
the set of Reimannian gradients."""
if self.metric == 'log_euclidean':
lmbd, U = tf.linalg.eigh(u)
f = _f_matrix(lmbd)
# Riemannian gradient
E = adj(U) @ ((egrad + adj(egrad)) / 2) @ U
R = U @ (E * f * f) @ adj(U)
return R
elif self.metric == 'log_cholesky':
n = u.shape[-1]
dtype = u.dtype
L = tf.linalg.cholesky(u)
mask = tf.ones((n, n), dtype=dtype)
mask = _lower(mask)
G_inv = mask + tf.linalg.diag(tf.linalg.diag_part(L) ** 2)
R = G_inv * ((egrad + adj(egrad)) @ L)
R_diag = tf.linalg.diag(tf.linalg.diag_part(R))
R = R - 1j * tf.cast(tf.math.imag(R_diag), dtype=u.dtype)
return _push_forward_chol(R, L)
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., n ** 2, k),
a set of points from the manifold.
vec: complex valued tensor of shape (..., n ** 2, k),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n ** 2, k),
a set of projected vectors."""
k = u.shape[-1]
n = int(math.sqrt(u.shape[-2]))
shape = u.shape[:-2]
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.reshape(vec, shape + (n, k * n))
vec_mod = tf.linalg.matrix_transpose(vec_mod)
u_mod = tf.reshape(u, shape + (n, k * n))
u_mod = tf.linalg.matrix_transpose(u_mod)
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.linalg.matrix_transpose(vec_mod)
vec_mod = tf.reshape(vec_mod, shape + (n ** 2, k))
# projection onto the horizontal space
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_mod - adj(vec_mod) @ u)
return vec_mod - u @ Omega
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., n, n),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n, n),
the set of Reimannian gradients.
Note:
The complexity O(n^3)."""
if self._metric == 'log_euclidean':
lmbd, U = tf.linalg.eigh(u)
f = _f_matrix(lmbd)
# Riemannian gradient
E = adj(U) @ ((egrad + adj(egrad)) / 2) @ U
R = U @ (E * f * f) @ adj(U)
return R
elif self._metric == 'log_cholesky':
L = tf.linalg.cholesky(u)
sym_fl = (egrad + adj(egrad)) @ L
L_sq_diag = tf.linalg.diag_part(L)**2
term2 = 0.5 * L_sq_diag * tf.linalg.diag_part(sym_fl + adj(sym_fl))
term2 = tf.linalg.diag(term2)
term1 = _lower(sym_fl)
R = term1 + term2
return _push_forward_chol(R, L)
def inner(self, u, vec1, vec2):
"""Returns manifold wise inner product of vectors from
a tangent space.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
vec1: complex valued tensor of shape (..., n, p),
a set of tangent vectors from the complex
Stiefel manifold.
vec2: complex valued tensor of shape (..., n, p),
a set of tangent vectors from the complex
Stiefel manifold.
Returns:
complex valued tensor of shape (..., 1, 1),
manifold wise inner product"""
if self._metric == 'euclidean':
s_sq = tf.linalg.trace(adj(vec1) @ vec2)[..., tf.newaxis,
tf.newaxis]
elif self._metric == 'canonical':
G = tf.eye(u.shape[-2], dtype=u.dtype) - u @ adj(u) / 2
s_sq = tf.linalg.trace(adj(vec1) @ G @ vec2)[..., tf.newaxis,
tf.newaxis]
return tf.cast(tf.math.real(s_sq), dtype=u.dtype)
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., n, r),
a set of points from the manifold.
vec: complex valued tensor of shape (..., n, r),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n, r),
a set of projected vectors.
Note:
The complexity is O(nr^2)."""
# projection onto the tangent space of ||u||_F = 1
vec_proj = vec - u * tf.reduce_sum(
tf.math.conj(u) * vec, axis=(-2, -1))[..., tf.newaxis, tf.newaxis]
# projection onto the horizontal space
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_proj - adj(vec_proj) @ u)
return vec_proj - u @ Omega
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., m, n, n),
a set of points from the manifold.
vec: complex valued tensor of shape (..., m, n, n),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., m, n, n),
a set of projected vectors.
Note:
The complexity is O(mn^3)."""
n = tf.shape(u)[-1]
m = tf.shape(u)[-3]
shape = tf.shape(u)[:-3]
size = tf.size(shape)
idx = tf.range(size)
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.transpose(
vec,
shape_conc(idx, (size + 2)[tf.newaxis], size[tf.newaxis],
(size + 1)[tf.newaxis]))
vec_mod = tf.reshape(
vec_mod, shape_conc(shape, (n * m)[tf.newaxis], n[tf.newaxis]))
u_mod = tf.transpose(
u,
shape_conc(idx, (size + 2)[tf.newaxis], size[tf.newaxis],
(size + 1)[tf.newaxis]))
u_mod = tf.reshape(
u_mod, shape_conc(shape, (n * m)[tf.newaxis], n[tf.newaxis]))
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.reshape(
vec_mod,
shape_conc(shape, n[tf.newaxis], m[tf.newaxis], n[tf.newaxis]))
vec_mod = tf.transpose(
vec_mod,
shape_conc(idx, (size + 1)[tf.newaxis], (size + 2)[tf.newaxis],
size[tf.newaxis]))
# projection onto the horizontal space (POVM element-wise)
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_mod - adj(vec_mod) @ u)
return vec_mod - u @ Omega
def vector_transport(self, u, vec1, vec2):
"""Returns a vector tranported along an another vector
via vector transport.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold, starting points.
vec1: complex valued tensor of shape (..., n, n),
a set of vectors to be transported.
vec2: complex valued tensor of shape (..., n, n),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., n, n),
a set of transported vectors.
Note:
The complexity O(n^3)."""
if self._metric == 'log_euclidean':
lmbd, U = tf.linalg.eigh(u)
# geoidesic in S
Su = U @ tf.linalg.diag(tf.math.log(lmbd)) @ adj(U)
Svec2 = _pull_back_log(vec2, U, lmbd)
Sresult = Su + Svec2
# eig decomposition of a new point from S
log_new_lmbd, new_U = tf.linalg.eigh(Sresult)
# new lmbd
new_lmbd = tf.exp(log_new_lmbd)
# transported vector
new_vec1 = _push_forward_log(_pull_back_log(vec1, U, lmbd), new_U,
new_lmbd)
return new_vec1
elif self._metric == 'log_cholesky':
v = self.retraction(u, vec2)
L = tf.linalg.cholesky(u)
inv_L = tf.linalg.inv(L)
inv_diag_L = tf.linalg.diag(1 / tf.linalg.diag_part(L))
X = _pull_back_chol(vec1, L, inv_L)
K = tf.linalg.cholesky(v)
L_transport = _lower(X) + tf.linalg.band_part(K, 0, 0) *\
inv_diag_L * tf.linalg.band_part(X, 0, 0)
return K @ adj(L_transport) + L_transport @ adj(K)
def retraction_transport(self, u, vec1, vec2):
"""Performs a retraction and a vector transport simultaneously.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold, starting points.
vec1: complex valued tensor of shape (..., n, n),
a set of vectors to be transported.
vec2: complex valued tensor of shape (..., n, n),
a set of direction vectors.
Returns:
two complex valued tensors of shape (..., n, n),
a set of transported points and vectors."""
if self.metric == 'log_euclidean':
lmbd, U = tf.linalg.eigh(u)
# geoidesic in S
Su = U @ tf.linalg.diag(tf.math.log(lmbd)) @ adj(U)
Svec2 = _pull_back_log(vec2, U, lmbd)
Sresult = Su + Svec2
# eig decomposition of new point from S
log_new_lmbd, new_U = tf.linalg.eigh(Sresult)
# new point from S++
new_point = new_U @ tf.linalg.diag(tf.exp(log_new_lmbd)) @\
adj(new_U)
# new lmbd
new_lmbd = tf.exp(log_new_lmbd)
# transported vector
new_vec1 = _push_forward_log(_pull_back_log(vec1, U, lmbd),
new_U, new_lmbd)
return new_point, new_vec1
elif self.metric == 'log_cholesky':
v = self.retraction(u, vec2)
L = tf.linalg.cholesky(u)
inv_L = tf.linalg.inv(L)
inv_diag_L = tf.linalg.diag(1 / tf.linalg.diag_part(L))
X = _pull_back_chol(vec1, L, inv_L)
K = tf.linalg.cholesky(v)
L_transport = _lower(X) + tf.linalg.band_part(K, 0, 0) *\
inv_diag_L * tf.linalg.band_part(X, 0, 0)
return v, K @ adj(L_transport) + L_transport @ adj(K)
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., m, n, n),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., m, n, n),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., m, n, n),
the set of Reimannian gradients.
Note:
The complexity is O(mn^3)."""
n = tf.shape(u)[-1]
m = tf.shape(u)[-3]
shape = tf.shape(u)[:-3]
size = tf.size(shape)
idx = tf.range(size)
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.transpose(
egrad,
shape_conc(idx, (size + 2)[tf.newaxis], size[tf.newaxis],
(size + 1)[tf.newaxis]))
vec_mod = tf.reshape(
vec_mod, shape_conc(shape, (n * m)[tf.newaxis], n[tf.newaxis]))
u_mod = tf.transpose(
u,
shape_conc(idx, (size + 2)[tf.newaxis], size[tf.newaxis],
(size + 1)[tf.newaxis]))
u_mod = tf.reshape(
u_mod, shape_conc(shape, (n * m)[tf.newaxis], n[tf.newaxis]))
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.reshape(
vec_mod,
shape_conc(shape, n[tf.newaxis], m[tf.newaxis], n[tf.newaxis]))
vec_mod = tf.transpose(
vec_mod,
shape_conc(idx, (size + 1)[tf.newaxis], (size + 2)[tf.newaxis],
size[tf.newaxis]))
return vec_mod
def retraction(self, u, vec):
"""Transports a set of points from the manifold via a
retraction map.
Args:
u: complex valued tensor of shape (..., n ** 2, k), a set
of points to be transported.
vec: complex valued tensor of shape (..., n ** 2, k),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., n ** 2, k),
a set of transported points.
Note:
The complexity O(kn^3)"""
k = tf.shape(u)[-1]
n = tf.cast(tf.math.sqrt(tf.cast(tf.shape(u)[-2], dtype=tf.float32)),
dtype=tf.int32)
shape = tf.shape(u)[:-2]
# svd based retraction
u_new = u + vec
u_new = tf.reshape(
u_new, shape_conc(shape, n[tf.newaxis], (n * k)[tf.newaxis]))
_, U, V = tf.linalg.svd(u_new)
u_new = U @ adj(V)
u_new = tf.reshape(
u_new, shape_conc(shape, (n**2)[tf.newaxis], k[tf.newaxis]))
return u_new
def retraction(self, u, vec):
"""Transports a set of points from the manifold via a
retraction map.
Args:
u: complex valued tensor of shape (..., m, n, n), a set
of points to be transported.
vec: complex valued tensor of shape (..., m, n, n),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., m, n, n),
a set of transported points."""
n = u.shape[-1]
m = u.shape[-3]
shape = u.shape[:-3]
size = len(shape)
idx = tuple(range(size))
# svd based retraction
u_new = u + vec
u_new = tf.transpose(u_new, idx + (size + 1, size + 2, size))
u_new = tf.reshape(u_new, shape + (n, n * m))
_, U, V = tf.linalg.svd(u_new)
u_new = U @ adj(V)
u_new = tf.reshape(u_new, shape + (n, n, m))
u_new = tf.transpose(u_new, idx + (size + 2, size, size + 1))
return u_new
def random(self, shape, dtype=tf.complex64):
"""Returns a set of points from the manifold generated
randomly.
Args:
shape: tuple of integer numbers (..., n, n),
shape of a generated matrix.
dtype: type of an output tensor, can be
either tf.complex64 or tf.complex128.
Returns:
complex valued tensor of shape (..., n, n),
a generated matrix."""
list_of_dtypes = [tf.complex64, tf.complex128]
if dtype not in list_of_dtypes:
raise ValueError("Incorrect dtype")
real_dtype = tf.float64 if dtype == tf.complex128 else tf.float32
u = tf.complex(tf.random.normal(shape, dtype=real_dtype),
tf.random.normal(shape, dtype=real_dtype))
u = 0.5 * (u + adj(u))
return u
def is_in_manifold(self, u, tol=1e-5):
"""Checks if a point is in the manifold or not.
Args:
u: complex valued tensor of shape (..., m, n, n),
a point to be checked.
tol: small real value showing tolerance.
Returns:
bolean tensor of shape (...)."""
n = tf.shape(u)[-1]
m = tf.shape(u)[-3]
shape = tf.shape(u)[:-3]
u_resh = tf.linalg.matrix_transpose(u)
u_resh = tf.reshape(
u_resh, shape_conc(shape, (m * n)[tf.newaxis], n[tf.newaxis]))
u_resh = tf.linalg.matrix_transpose(u_resh)
uudag = u_resh @ adj(u_resh)
Id = tf.eye(tf.shape(uudag)[-1], dtype=u.dtype)
diff = tf.linalg.norm(uudag - Id, axis=(-2, -1))
uudag_norm = tf.linalg.norm(uudag, axis=(-2, -1))
Id_norm = tf.linalg.norm(Id, axis=(-2, -1))
rel_diff = tf.abs(diff / tf.math.sqrt(Id_norm * uudag_norm))
return tol > rel_diff
def retraction(self, u, vec):
"""Transports a set of points from the manifold via a
retraction map.
Args:
u: complex valued tensor of shape (..., n ** 2, k), a set
of points to be transported.
vec: complex valued tensor of shape (..., n ** 2, k),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., n ** 2, k),
a set of transported points."""
k = u.shape[-1]
n = int(math.sqrt(u.shape[-2]))
shape = u.shape[:-2]
# svd based retraction
u_new = u + vec
u_new = tf.reshape(u_new, shape + (n, n * k))
_, U, V = tf.linalg.svd(u_new)
u_new = U @ adj(V)
u_new = tf.reshape(u_new, shape + (n ** 2, k))
return u_new
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the complex Stiefel manifold.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
vec: complex valued tensor of shape (..., n, p),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n, p),
a set of projected vectors"""
return 0.5 * u @ (adj(u) @ vec - adj(vec) @ u) +\
(tf.eye(u.shape[-2], dtype=u.dtype) -\
u @ adj(u)) @ vec
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
egrad: complex valued tensor of shape (..., n, p),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n, p),
the set of Reimannian gradients."""
if self._metric == 'euclidean':
return 0.5 * u @ (adj(u) @ egrad - adj(egrad) @ u) +\
(tf.eye(u.shape[-2], dtype=u.dtype) -\
u @ adj(u)) @ egrad
elif self._metric == 'canonical':
return egrad - u @ adj(egrad) @ u
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the complex Stiefel manifold.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
vec: complex valued tensor of shape (..., n, p),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n, p),
a set of projected vectors
Note:
the complexity is O(np^2)"""
u_shape = tf.shape(u)
return 0.5 * u @ (adj(u) @ vec - adj(vec) @ u) +\
vec - u @ (adj(u) @ vec)
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., n ** 2, k),
a set of points from the manifold.
vec: complex valued tensor of shape (..., n ** 2, k),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n ** 2, k),
a set of projected vectors.
Note:
The complexity O(kn^3+k^2n^2)"""
k = tf.shape(u)[-1]
n = tf.cast(tf.math.sqrt(tf.cast(tf.shape(u)[-2], dtype=tf.float32)),
dtype=tf.int32)
shape = tf.shape(u)[:-2]
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.reshape(
vec, shape_conc(shape, n[tf.newaxis], (k * n)[tf.newaxis]))
vec_mod = tf.linalg.matrix_transpose(vec_mod)
u_mod = tf.reshape(
u, shape_conc(shape, n[tf.newaxis], (k * n)[tf.newaxis]))
u_mod = tf.linalg.matrix_transpose(u_mod)
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.linalg.matrix_transpose(vec_mod)
vec_mod = tf.reshape(
vec_mod, shape_conc(shape, (n**2)[tf.newaxis], k[tf.newaxis]))
# projection onto the horizontal space
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_mod - adj(vec_mod) @ u)
return vec_mod - u @ Omega
def inner(self, u, vec1, vec2):
"""Returns manifold wise inner product of vectors from
a tangent space.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
vec1: complex valued tensor of shape (..., n, p),
a set of tangent vectors from the complex
Stiefel manifold.
vec2: complex valued tensor of shape (..., n, p),
a set of tangent vectors from the complex
Stiefel manifold.
Returns:
complex valued tensor of shape (..., 1, 1),
manifold wise inner product
Note:
The complexity for the 'euclidean' metric is O(pn),
the complexity for the 'canonical' metric is O(np^2)"""
if self._metric == 'euclidean':
s_sq = tf.reduce_sum(tf.math.conj(vec1) * vec2,
axis=(-2, -1),
keepdims=True)
elif self._metric == 'canonical':
s_sq_1 = tf.reduce_sum(tf.math.conj(vec1) * vec2,
axis=(-2, -1),
keepdims=True)
vec1_dag_u = adj(vec1) @ u
u_dag_vec2 = adj(u) @ vec2
s_sq_2 = tf.reduce_sum(u_dag_vec2 *
tf.linalg.matrix_transpose(vec1_dag_u),
axis=(-2, -1),
keepdims=True)
s_sq = s_sq_1 - 0.5 * s_sq_2
return tf.cast(tf.math.real(s_sq), dtype=u.dtype)
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold.
vec: complex valued tensor of shape (..., n, n),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n, n),
a set of projected vectors"""
# projection onto the tangent space of ||u||_F = 1
vec_proj = vec - u * tf.linalg.trace(adj(u) @ vec)[..., tf.newaxis,
tf.newaxis]
# projection onto the horizontal space
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_proj - adj(vec_proj) @ u)
return vec_proj - u @ Omega
def retraction(self, u, vec):
"""Transports a set of points from the manifold via a
retraction map.
Args:
u: complex valued tensor of shape (..., n, n), a set
of points to be transported.
vec: complex valued tensor of shape (..., n, n),
a set of direction vectors.
Returns:
complex valued tensor of shape (..., n, n),
a set of transported points.
Note:
The complexity O(n^3)."""
if self._metric == 'log_euclidean':
lmbd, U = tf.linalg.eigh(u)
# geodesic in S
Su = U @ tf.linalg.diag(tf.math.log(lmbd)) @ adj(U)
Svec = _pull_back_log(vec, U, lmbd)
Sresult = Su + Svec
return tf.linalg.expm(Sresult)
elif self._metric == 'log_cholesky':
L = tf.linalg.cholesky(u)
inv_L = tf.linalg.inv(L)
X = _pull_back_chol(vec, L, inv_L)
inv_diag_L = tf.linalg.diag(1 / tf.linalg.diag_part(L))
cholesky_retraction = _lower(L) + _lower(X) +\
tf.linalg.band_part(L, 0, 0) *\
tf.exp(tf.linalg.band_part(X, 0, 0) * inv_diag_L)
return cholesky_retraction @ adj(cholesky_retraction)
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold.
vec: complex valued tensor of shape (..., n, n),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., n, n),
a set of projected vectors."""
return (vec + adj(vec)) / 2
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n, p),
a set of points from the complex Stiefel
manifold.
egrad: complex valued tensor of shape (..., n, p),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n, p),
the set of Reimannian gradients.
Note:
The complexity is O(np^2)"""
if self._metric == 'euclidean':
u_shape = tf.shape(u)
return 0.5 * u @ (adj(u) @ egrad - adj(egrad) @ u) +\
egrad - u @ (adj(u) @ egrad)
elif self._metric == 'canonical':
return egrad - u @ (adj(egrad) @ u)
def is_in_manifold(self, u, tol=1e-5):
"""Checks if a point is in the manifold or not.
Args:
u: complex valued tensor of shape (..., n, n),
a point to be checked.
tol: small real value showing tolerance.
Returns:
bolean tensor of shape (...)."""
diff_norm = tf.linalg.norm(u - adj(u), axis=(-2, -1))
u_norm = tf.linalg.norm(u, axis=(-2, -1))
rel_diff = tf.abs(diff_norm / u_norm)
return tol > rel_diff
def proj(self, u, vec):
"""Returns projection of vectors on a tangen space
of the manifold.
Args:
u: complex valued tensor of shape (..., m, n, n),
a set of points from the manifold.
vec: complex valued tensor of shape (..., m, n, n),
a set of vectors to be projected.
Returns:
complex valued tensor of shape (..., m, n, n),
a set of projected vectors."""
n = u.shape[-1]
m = u.shape[-3]
shape = u.shape[:-3]
size = len(shape)
idx = tuple(range(size))
# projection onto the tangent space of the Stiefel manifold
vec_mod = tf.transpose(vec, idx + (size + 2, size, size + 1))
vec_mod = tf.reshape(vec_mod, shape + (n * m, n))
u_mod = tf.transpose(u, idx + (size + 2, size, size + 1))
u_mod = tf.reshape(u_mod, shape + (n * m, n))
vec_mod = vec_mod - 0.5 * u_mod @ (adj(u_mod) @ vec_mod +\
adj(vec_mod) @ u_mod)
vec_mod = tf.reshape(vec_mod, shape + (n, m, n))
vec_mod = tf.transpose(vec_mod, idx + (size + 1, size + 2, size))
# projection onto the horizontal space (POVM element-wise)
uu = adj(u) @ u
Omega = lyap_symmetric(uu, adj(u) @ vec_mod - adj(vec_mod) @ u)
return vec_mod - u @ Omega
def egrad_to_rgrad(self, u, egrad):
"""Returns the Riemannian gradient from an Euclidean gradient.
Args:
u: complex valued tensor of shape (..., n, n),
a set of points from the manifold.
egrad: complex valued tensor of shape (..., n, n),
a set of Euclidean gradients.
Returns:
complex valued tensor of shape (..., n, n),
the set of Reimannian gradients."""
rgrad = egrad - u * tf.linalg.trace(adj(u) @ egrad)[..., tf.newaxis,
tf.newaxis]
return rgrad
