def test_broadcast_arrays(self): array_1 = gs.array([[1, 2, 3]]) array_2 = gs.array([[4], [5]]) result = gs.broadcast_arrays(array_1, array_2) result_verdict = [gs.array([[1, 2, 3], [1, 2, 3]]), gs.array([[4, 4, 4], [5, 5, 5]])] self.assertAllClose(result[0], result_verdict[0]) self.assertAllClose(result[1], result_verdict[1]) with self.assertRaises((ValueError, RuntimeError)): gs.broadcast_arrays(gs.array([1, 2]), gs.array([3, 4, 5]))
def is_tangent(self, vector, base_point, atol=gs.atol): """Check whether the vector is tangent at base_point. The tangent space of the product manifold is the direct sum of tangent spaces. Parameters ---------- vector : array-like, shape=[..., n_copies, *base_shape] Vector. base_point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. atol : float Absolute tolerance. Optional, default: backend atol. Returns ------- is_tangent : bool Boolean denoting if vector is a tangent vector at the base point. """ vector_, point_ = gs.broadcast_arrays(vector, base_point) point_ = gs.reshape(point_, (-1, *self.base_shape)) vector_ = gs.reshape(vector_, (-1, *self.base_shape)) each_tangent = self.base_manifold.is_tangent(vector_, point_) reshaped = gs.reshape(each_tangent, (-1, self.n_copies)) return gs.all(reshaped, axis=1)
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point): """Compute the inner-product of two tangent vectors at a base point. Inner product defined by the Riemannian metric at point `base_point` between tangent vectors `tangent_vec_a` and `tangent_vec_b`. Parameters ---------- tangent_vec_a : array-like, shape=[..., n_copies, *base_shape] First tangent vector at base point. tangent_vec_b : array-like, shape=[..., n_copies, *base_shape] Second tangent vector at base point. base_point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. Optional, default: None. Returns ------- inner_prod : array-like, shape=[...,] Inner-product of the two tangent vectors. """ tangent_vec_a_, tangent_vec_b_, point_ = gs.broadcast_arrays( tangent_vec_a, tangent_vec_b, base_point ) point_ = gs.reshape(point_, (-1, *self.base_shape)) vector_a = gs.reshape(tangent_vec_a_, (-1, *self.base_shape)) vector_b = gs.reshape(tangent_vec_b_, (-1, *self.base_shape)) inner_each = self.base_metric.inner_product(vector_a, vector_b, point_) reshaped = gs.reshape(inner_each, (-1, self.n_copies)) return gs.squeeze(gs.sum(reshaped, axis=-1))
def to_tangent(self, vector, base_point): """Project a vector to a tangent space of the manifold. The tangent space of the product manifold is the direct sum of tangent spaces. Parameters ---------- vector : array-like, shape=[..., n_copies, *base_shape] Vector. base_point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. Returns ------- tangent_vec : array-like, shape=[..., n_copies, *base_shape] Tangent vector at base point. """ vector_, point_ = gs.broadcast_arrays(vector, base_point) point_ = gs.reshape(point_, (-1, *self.base_shape)) vector_ = gs.reshape(vector_, (-1, *self.base_shape)) each_tangent = self.base_manifold.to_tangent(vector_, point_) reshaped = gs.reshape(each_tangent, (-1, self.n_copies) + self.base_shape) return gs.squeeze(reshaped)
def dist_broadcast(self, point_a, point_b): """Compute the geodesic distance between points. If n_samples_a == n_samples_b then dist is the element-wise distance result of a point in points_a with the point from points_b of the same index. If n_samples_a not equal to n_samples_b then dist is the result of applying geodesic distance for each point from points_a to all points from points_b. Parameters ---------- point_a : array-like, shape=[n_samples_a, dim] Set of points in hyperbolic space. point_b : array-like, shape=[n_samples_b, dim] Second set of points in hyperbolic space. Returns ------- dist : array-like, shape=[n_samples_a, dim] or [n_samples_a, n_samples_b, dim] Geodesic distance between the two points. """ if point_a.shape[-1] != point_b.shape[-1]: raise ValueError('Manifold dimensions not equal') if point_a.shape[0] != point_b.shape[0]: point_a_broadcast, point_b_broadcast = gs.broadcast_arrays( point_a[:, None], point_b[None, ...]) point_a_flatten = gs.reshape(point_a_broadcast, (-1, point_a_broadcast.shape[-1])) point_b_flatten = gs.reshape(point_b_broadcast, (-1, point_b_broadcast.shape[-1])) point_a_norm = gs.clip(gs.sum(point_a_flatten**2, -1), 0., 1 - EPSILON) point_b_norm = gs.clip(gs.sum(point_b_flatten**2, -1), 0., 1 - EPSILON) square_diff = (point_a_flatten - point_b_flatten)**2 diff_norm = gs.sum(square_diff, -1) norm_function = 1 + 2 * \ diff_norm / ((1 - point_a_norm) * (1 - point_b_norm)) dist = gs.log(norm_function + gs.sqrt(norm_function**2 - 1)) dist *= self.scale dist = gs.reshape(dist, (point_a.shape[0], point_b.shape[0])) dist = gs.squeeze(dist) elif point_a.shape == point_b.shape: dist = self.dist(point_a, point_b) return dist
def dist_broadcast(self, point_a, point_b): """Compute the geodesic distance between points. If n_samples_a == n_samples_b then dist is the element-wise distance result of a point in points_a with the point from points_b of the same index. If n_samples_a not equal to n_samples_b then dist is the result of applying geodesic distance for each point from points_a to all points from points_b. Parameters ---------- point_a : array-like, shape=[n_samples_a, dim] Set of points in the Poincare ball. point_b : array-like, shape=[n_samples_b, dim] Second set of points in the Poincare ball. Returns ------- dist : array-like, shape=[n_samples_a, dim] or [n_samples_a, n_samples_b, dim] Geodesic distance between the two points. """ ndim = len(self.shape) if point_a.shape[-ndim:] != point_b.shape[-ndim:]: raise ValueError("Manifold dimensions not equal") if ndim in (point_a.ndim, point_b.ndim) or (point_a.shape == point_b.shape): return self.dist(point_a, point_b) n_samples = point_a.shape[0] * point_b.shape[0] point_a_broadcast, point_b_broadcast = gs.broadcast_arrays( point_a[:, None], point_b[None, ...] ) point_a_flatten = gs.reshape( point_a_broadcast, (n_samples,) + point_a.shape[-ndim:] ) point_b_flatten = gs.reshape( point_b_broadcast, (n_samples,) + point_a.shape[-ndim:] ) dist = self.dist(point_a_flatten, point_b_flatten) dist = gs.reshape(dist, (point_a.shape[0], point_b.shape[0])) dist = gs.squeeze(dist) return dist
def matching(self, base_graph, graph_to_permute): """Match graphs. Parameters ---------- base_graph : list of Graph or array-like, shape=[..., n, n]. Base graph. graph_to_permute : list of Graph or array-like, shape=[..., n, n]. Graph to align. """ base_graph, graph_to_permute = gs.broadcast_arrays( base_graph, graph_to_permute) is_single = gs.ndim(base_graph) == 2 if is_single: base_graph = gs.expand_dims(base_graph, 0) graph_to_permute = gs.expand_dims(graph_to_permute, 0) perm = self.matcher.match(base_graph, graph_to_permute) self.perm_ = gs.array(perm[0]) if is_single else gs.array(perm) return self.perm_
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None): """Compute the inner-product of two tangent vectors at a base point. Parameters ---------- tangent_vec_a : array-like, shape=[..., n_samples] First tangent vector at base point. tangent_vec_b : array-like, shape=[..., n_samples] Second tangent vector at base point. base_point : array-like, shape=[..., n_samples], optional Point on the hypersphere. Returns ------- inner_prod : array-like, shape=[...,] Inner-product of the two tangent vectors. """ tangent_vec_a, tangent_vec_b = gs.broadcast_arrays(tangent_vec_a, tangent_vec_b) x = gs.broadcast_to(self.x, tangent_vec_a.shape) l2_norm = gs.trapz(tangent_vec_a * tangent_vec_b, x=x, axis=-1) return l2_norm
def log(self, point, base_point, **kwargs): """Compute the Riemannian logarithm of a point. Parameters ---------- point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. base_point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. Optional, default: None. Returns ------- log : array-like, shape=[..., n_copies, *base_shape] Tangent vector at the base point equal to the Riemannian logarithm of point at the base point. """ point_, base_point_ = gs.broadcast_arrays(point, base_point) base_point_ = gs.reshape(base_point_, (-1, *self.base_shape)) point_ = gs.reshape(point_, (-1, *self.base_shape)) each_log = self.base_metric.log(point_, base_point_) reshaped = gs.reshape(each_log, (-1, self.n_copies) + self.base_shape) return gs.squeeze(reshaped)
def exp(self, tangent_vec, base_point, **kwargs): """Compute the Riemannian exponential of a tangent vector. Parameters ---------- tangent_vec : array-like, shape=[..., n_copies, *base_shape] Tangent vector at a base point. base_point : array-like, shape=[..., n_copies, *base_shape] Point on the manifold. Optional, default: None. Returns ------- exp : array-like, shape=[..., n_copies, *base_shape] Point on the manifold equal to the Riemannian exponential of tangent_vec at the base point. """ tangent_vec, point_ = gs.broadcast_arrays(tangent_vec, base_point) point_ = gs.reshape(point_, (-1, *self.base_shape)) vector_ = gs.reshape(tangent_vec, (-1, *self.base_shape)) each_exp = self.base_metric.exp(vector_, point_) reshaped = gs.reshape(each_exp, (-1, self.n_copies) + self.base_shape) return gs.squeeze(reshaped)
def _batch_gradient_descent( points, metric, weights=None, max_iter=32, lr=1e-3, epsilon=5e-3, point_type="vector", verbose=False, ): """Perform batch gradient descent.""" if point_type == "vector": if points.ndim < 3: return _default_gradient_descent( points, metric, weights, max_iter, point_type, epsilon, lr, verbose ) einsum_str = "ni,nij->ij" ndim = 1 else: if points.ndim < 4: return _default_gradient_descent( points, metric, weights, max_iter, point_type, epsilon, lr, verbose ) einsum_str = "nk,nkij->kij" ndim = 2 shape = points.shape n_points = shape[0] n_batch = shape[1] if n_points == 1: return points[0] if weights is None: weights = gs.ones((n_points, n_batch)) flat_shape = (n_batch * n_points,) + shape[-ndim:] estimates = points[0] points_flattened = gs.reshape(points, (n_points * n_batch,) + shape[-ndim:]) convergence = math.inf iteration = 0 convergence_old = convergence while convergence > epsilon and max_iter > iteration: iteration += 1 estimates_broadcast, _ = gs.broadcast_arrays(estimates, points) estimates_flattened = gs.reshape(estimates_broadcast, flat_shape) tangent_grad = metric.log(points_flattened, estimates_flattened) tangent_grad = gs.reshape(tangent_grad, shape) tangent_mean = gs.einsum(einsum_str, weights, tangent_grad) / n_points next_estimates = metric.exp(lr * tangent_mean, estimates) convergence = gs.sum(metric.squared_norm(tangent_mean, estimates)) estimates = next_estimates if convergence < convergence_old: convergence_old = convergence elif convergence > convergence_old: lr = lr / 2.0 if iteration == max_iter: logging.warning( "Maximum number of iterations {} reached. The " "mean may be inaccurate".format(max_iter) ) if verbose: logging.info( "n_iter: {}, final dist: {}," "final step size: {}".format(iteration, convergence, lr) ) return estimates