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