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
