def projx(self, x): warnings.warn(("{}.projx performs a projection onto the open set of" + " PSD matrices").format(self.__class__.__name__)) x_sym = (utils.transposem(x) + x) / 2.0 s, u, v = tf.linalg.svd(x_sym) sigma = tf.linalg.diag(tf.maximum(s, 0.0)) return v @ sigma @ utils.transposem(v)
def retr(self, x, u): xut = x @ utils.transposem(u) xut_sym = (xut - utils.transposem(xut)) / 2.0 eye = tf.eye(tf.shape(xut)[-1], batch_shape=tf.shape(xut)[:-2], dtype=x.dtype) return tf.linalg.solve(xut_sym + eye, x - xut_sym @ x)
def diff_from_cholesky(self, x, u): """Inverse of the differential of diffeomorphism to the Cholesky space""" assert_x = tf.debugging.Assert( self._cholesky.check_point_on_manifold(x), [x]) assert_u = tf.debugging.Assert( self._cholesky.check_vector_on_tangent(x, u), [u]) with tf.control_dependencies([assert_x, assert_u]): return x @ utils.transposem(u) + u @ utils.transposem(x)
def retr(self, x, u): xtu = utils.transposem(x) @ u w_ = (u - x @ xtu) @ utils.transposem(x) w = w_ - utils.transposem(w_) y = x + u for _ in range(self.num_iter): y = x + w @ ((x + y) / 2.0) return y
def inner(self, x, u, v, keepdims=False): xtu = utils.transposem(x) @ u xtv = utils.transposem(x) @ v u_v_inner = tf.reduce_sum(u * v, axis=[-2, -1], keepdims=keepdims) xtu_xtv_inner = tf.reduce_sum(xtu * xtv, axis=[-2, -1], keepdims=keepdims) return u_v_inner - 0.5 * xtu_xtv_inner
def geodesic(self, x, u, t): xtu = utils.transposem(x) @ u w_ = (u - x @ xtu) @ utils.transposem(x) w = w_ - utils.transposem(w_) eye = tf.linalg.eye(tf.shape(w)[-1], batch_shape=tf.shape(w)[:-2], dtype=x.dtype) cayley_t = tf.linalg.inv(eye - t * w / 2.0) @ (eye + t * w / 2.0) return cayley_t @ x
def ptransp(self, x, y, v): log_xy = self.log(x, y) s, u, vt = tf.linalg.svd(log_xy, full_matrices=False) cos_s = tf.linalg.diag(tf.math.cos(s)) sin_s = tf.linalg.diag(tf.math.sin(s)) geod = ((-x @ utils.transposem(vt) @ sin_s + u @ cos_s) @ utils.transposem(u) @ v) proj = v - u @ utils.transposem(u) @ v return geod + proj
def geodesic(self, x, u, t): xtu = utils.transposem(x) @ u utu = utils.transposem(u) @ u eye = tf.eye(tf.shape(utu)[-1], batch_shape=tf.shape(utu)[:-2], dtype=x.dtype) logw = blockm(xtu, -utu, eye, xtu) w = tf.linalg.expm(t * logw) z = tf.concat([tf.linalg.expm(-xtu * t), tf.zeros_like(utu)], axis=-2) y = tf.concat([x, u], axis=-1) @ w @ z return y
def _check_point_on_manifold(self, x, atol, rtol): x_t = utils.transposem(x) eigvals, _ = tf.linalg.eigh(x) is_symmetric = utils.allclose(x, x_t, atol, rtol) is_pos_vals = utils.allclose(eigvals, tf.abs(eigvals), atol, rtol) is_zero_vals = utils.allclose(eigvals, tf.zeros_like(eigvals), atol, rtol) return is_symmetric & is_pos_vals & tf.logical_not(is_zero_vals)
def _check_point_on_manifold(self, x, atol, rtol): xtx = utils.transposem(x) @ x eye = tf.eye( tf.shape(xtx)[-1], batch_shape=tf.shape(xtx)[:-2], dtype=x.dtype ) is_orth = utils.allclose(xtx, eye, atol, rtol) det = tf.linalg.det(x) is_unit_det = utils.allclose(det, tf.ones_like(det), atol, rtol) return is_orth & is_unit_det
def diff_to_cholesky(self, x, u): """Differential of the diffeomorphism to the Cholesky space""" assert_x = tf.debugging.Assert(self.check_point_on_manifold(x), [x]) assert_u = tf.debugging.Assert(self.check_vector_on_tangent(x, u), [u]) with tf.control_dependencies([assert_x, assert_u]): y = self.to_cholesky(x) y_inv = tf.linalg.inv(y) p = y_inv @ u @ utils.transposem(y_inv) p_diag, p_lower = self._cholesky._diag_and_strictly_lower(p) return y @ (p_lower + 0.5 * tf.linalg.diag(p_diag))
def _check_point_on_manifold(self, x, atol, rtol): xtx = utils.transposem(x) @ x shape = xtx.shape.as_list() eye = tf.eye(shape[-1], batch_shape=shape[:-2]) is_idempotent = utils.allclose(xtx, tf.cast(eye, x.dtype), atol, rtol) s = tf.linalg.svd(x, compute_uv=False) rank = tf.math.count_nonzero(s, axis=-1, dtype=tf.float32) k = tf.ones_like(rank) * int(x.shape[-1]) is_col_rank = utils.allclose(rank, k, atol, rtol) return is_idempotent & is_col_rank
def _diff_power(self, x, d, power): e, v = tf.linalg.eigh(x) v_t = utils.transposem(v) e = tf.expand_dims(e, -2) if power == "log": pow_e = tf.math.log(e) elif power == "exp": pow_e = tf.math.exp(e) s = utils.transposem(tf.ones_like(e)) @ e pow_s = utils.transposem(tf.ones_like(pow_e)) @ pow_e denom = utils.transposem(s) - s numer = utils.transposem(pow_s) - pow_s abs_denom = tf.math.abs(denom) eps = utils.get_eps(x) if power == "log": numer = tf.where(abs_denom < eps, tf.ones_like(numer), numer) denom = tf.where(abs_denom < eps, utils.transposem(s), denom) elif power == "exp": numer = tf.where(abs_denom < eps, utils.transposem(pow_s), numer) denom = tf.where(abs_denom < eps, tf.ones_like(denom), denom) t = v_t @ d @ v * numer / denom return v @ t @ v_t
def proju(self, x, u): xtu = utils.transposem(x) @ u return u - x @ xtu
def _check_vector_on_tangent(self, x, u, atol, rtol): xtu = utils.transposem(x) @ u return utils.allclose(xtu, tf.zeros_like(xtu), atol, rtol)
def ptransp(self, x, y, v): e = tf.linalg.sqrtm(y @ tf.linalg.inv(x)) return e @ v @ utils.transposem(e)
def proju(self, x, u): return 0.5 * (utils.transposem(u) + u)
def log(self, x, y): xty = utils.transposem(x) @ y u = utils.logm(xty) return (u - utils.transposem(u)) / 2.0
def _check_vector_on_tangent(self, x, u, atol, rtol): u_t = utils.transposem(u) return utils.allclose(u, u_t, atol, rtol)
def proju(self, x, u): xtu = utils.transposem(x) @ u return (xtu - utils.transposem(xtu)) / 2.0
def geodesic(self, x, u, t): s, u, vt = tf.linalg.svd(u, full_matrices=False) cos_ts = tf.linalg.diag(tf.math.cos(t * s)) sin_ts = tf.linalg.diag(tf.math.sin(t * s)) return (x @ utils.transposem(vt) @ cos_ts + u @ sin_ts) @ vt
def dist(self, x, y, keepdims=False): s = tf.linalg.svd(utils.transposem(x) @ y, compute_uv=False) theta = tf.math.acos(tf.clip_by_value(s, -1.0, 1.0)) norm = tf.linalg.norm(theta, axis=-1, keepdims=False) return norm[..., tf.newaxis, tf.newaxis] if keepdims else norm
def call(self, inputs): if self._expand: inputs = tf.expand_dims(inputs, -3) return utils.transposem(self.w) @ inputs
def proju(self, x, u): u_sym = (utils.transposem(u) + u) / 2.0 u_diag, u_lower = self._diag_and_strictly_lower(u_sym) x_diag = tf.linalg.diag_part(x) return u_lower + tf.linalg.diag(u_diag * x_diag**2)
def projx(self, x): x_sym = (utils.transposem(x) + x) / 2.0 s, _u, v = tf.linalg.svd(x_sym) sigma = tf.linalg.diag(tf.maximum(s, 0.0)) spd = v @ sigma @ utils.transposem(v) return tf.linalg.cholesky(spd)
def call(self, inputs): return inputs @ utils.transposem(inputs)
def exp(self, x, u): s, u, vt = tf.linalg.svd(u, full_matrices=False) cos_s = tf.linalg.diag(tf.math.cos(s)) sin_s = tf.linalg.diag(tf.math.sin(s)) return (x @ utils.transposem(vt) @ cos_s + u @ sin_s) @ vt
def call(self, inputs): s, u, v = tf.linalg.svd(inputs) log_s = tf.math.log(s) return u @ tf.linalg.diag(log_s) @ utils.transposem(v)
def from_cholesky(self, x): """Inverse of the diffeomorphism to the Cholesky space""" assert_x = tf.debugging.Assert( self._cholesky.check_point_on_manifold(x), [x]) with tf.control_dependencies([assert_x]): return x @ utils.transposem(x)
def _check_vector_on_tangent(self, x, u, atol, rtol): diff = utils.transposem(u) + u return utils.allclose(diff, tf.zeros_like(diff), atol, rtol)