def find_normalization_factor(variances, variances_range,
normalization_factor_var):
"""Find the normalization factor given some variances.
Parameters
----------
variances : array-like, shape=[n_gaussians,]
Array of standard deviations for each component
of some GMM.
variances_range : array-like, shape=[n_variances,]
Array of standard deviations.
normalization_factor_var : array-like, shape=[n_variances,]
Array of computed normalization factor.
Returns
-------
norm_factor : array-like, shape=[n_gaussians,]
Array of normalization factors for the given
variances.
"""
n_gaussians, precision = variances.shape[0], variances_range.shape[0]
ref = gs.expand_dims(variances_range, 0)
ref = gs.repeat(ref, n_gaussians, axis=0)
val = gs.expand_dims(variances, 1)
val = gs.repeat(val, precision, axis=1)
difference = gs.abs(ref - val)
index = gs.argmin(difference, axis=-1)
norm_factor = normalization_factor_var[index]
return norm_factor
def permute_vectorization_test_data(self):
space = self._PointSet(2)
points = space.random_point(3)
permutation = gs.array([0, 1])
smoke_data = [
dict(space=space, graph=points[0], id_permutation=permutation),
dict(
space=space,
graph=points[0],
id_permutation=gs.repeat(gs.expand_dims(permutation, 0),
2,
axis=0),
),
dict(space=space, graph=points, id_permutation=permutation),
dict(
space=space,
graph=points,
id_permutation=gs.repeat(gs.expand_dims(permutation, 0),
points.shape[0],
axis=0),
),
]
return self.generate_tests(smoke_data)
def load_cities():
"""Load data from data/cities/cities.json.
Returns
-------
data : array-like, shape=[50, 2]
Array with each row representing one sample,
i. e. latitude and longitude of a city.
Angles are in radians.
name : list
List of city names.
"""
with open(CITIES_PATH, encoding='utf-8') as json_file:
data_file = json.load(json_file)
names = [row['city'] for row in data_file]
data = list(
map(
lambda row:
[row[col_name] / 180 * gs.pi for col_name in ['lat', 'lng']],
data_file,
))
data = gs.array(data)
colat = gs.pi / 2 - data[:, 0]
colat = gs.expand_dims(colat, axis=1)
lng = gs.expand_dims(data[:, 1] + gs.pi, axis=1)
data = gs.concatenate([colat, lng], axis=1)
sphere = Hypersphere(dim=2)
data = sphere.spherical_to_extrinsic(data)
return data, names
def normalization_factor(self, variances):
"""Return normalization factor.
Parameters
----------
variances : array-like, shape=[n,]
Array of equally distant values of the
variance precision time.
Returns
-------
norm_func : array-like, shape=[n,]
Normalisation factor for all given variances.
"""
binomial_coefficient = None
n_samples = variances.shape[0]
expand_variances = gs.expand_dims(variances, axis=0)
expand_variances = gs.repeat(expand_variances, self.dim, axis=0)
if binomial_coefficient is None:
dim_range = gs.arange(self.dim)
dim_range[0] = 1
n_fact = dim_range.prod()
k_fact = gs.concatenate([
gs.expand_dims(dim_range[:i].prod(), 0)
for i in range(1, dim_range.shape[0] + 1)
], 0)
nmk_fact = gs.flip(k_fact, 0)
binomial_coefficient = n_fact / (k_fact * nmk_fact)
binomial_coefficient = gs.expand_dims(binomial_coefficient, -1)
binomial_coefficient = gs.repeat(binomial_coefficient,
n_samples,
axis=1)
range_ = gs.expand_dims(gs.arange(self.dim), -1)
range_ = gs.repeat(range_, n_samples, axis=1)
ones_ = gs.expand_dims(gs.ones(self.dim), -1)
ones_ = gs.repeat(ones_, n_samples, axis=1)
alternate_neg = (-ones_)**(range_)
erf_arg = ((
(self.dim - 1) - 2 * range_) * expand_variances) / gs.sqrt(2)
exp_arg = ((((self.dim - 1) - 2 * range_) * expand_variances) /
gs.sqrt(2))**2
norm_func_1 = (1 + gs.erf(erf_arg)) * gs.exp(exp_arg)
norm_func_2 = binomial_coefficient * norm_func_1
norm_func_3 = alternate_neg * norm_func_2
norm_func = NORMALIZATION_FACTOR_CST * variances * \
norm_func_3.sum(0) * (1 / (2 ** (self.dim - 1)))
return norm_func
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
"""Closed-form for the gradient of pose_loss.
Parameters
----------
y_pred : array-like
Prediction on SO(3).
y_true : array-like
Ground-truth on SO(3).
metric : RiemannianMetric
Metric used to compute the loss and gradient.
representation : str, {'vector', 'matrix'}
Representation chosen for points in SE(3).
Returns
-------
lie_grad : array-like
Tangent vector at point y_pred.
"""
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
lie_grad = lie_group.grad(y_pred, y_true, SO3, metric)
if representation == 'quaternion':
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = - 2 * quat_vec / (quat_sq_norm)
differential_scalar = gs.to_ndarray(differential_scalar, to_ndim=2)
differential_scalar = gs.transpose(differential_scalar)
differential_vec = (2 * (quat_scalar / quat_sq_norm
- 2 * quat_arctan2 / quat_vec_norm)
* (gs.einsum('ni,nj->nij', quat_vec, quat_vec)
/ quat_vec_norm ** 2)
+ 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_vec = gs.squeeze(differential_vec)
differential = gs.concatenate(
[differential_scalar, differential_vec],
axis=1)
y_pred = SO3.rotation_vector_from_quaternion(y_pred)
y_true = SO3.rotation_vector_from_quaternion(y_true)
lie_grad = lie_group.grad(y_pred, y_true, SO3, metric)
lie_grad = gs.matmul(lie_grad, differential)
lie_grad = gs.squeeze(lie_grad, axis=0)
return lie_grad
def grad(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""
Closed-form for the gradient of pose_loss.
:return: tangent vector at point y_pred.
"""
if gs.ndim(y_pred) == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if gs.ndim(y_true) == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SE3, metric)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred_pose = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:4]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm**2 + quat_scalar**2
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = -2 * quat_vec / (quat_sq_norm)
differential_vec = (
2 *
(quat_scalar / quat_sq_norm - 2 * quat_arctan2 / quat_vec_norm) *
(gs.einsum('ni,nj->nij', quat_vec, quat_vec) / quat_vec_norm *
quat_vec_norm) + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_scalar_t = gs.transpose(differential_scalar, axes=(1, 0))
upper_left_block = gs.hstack(
(differential_scalar_t, differential_vec[0]))
upper_right_block = gs.zeros((3, 3))
lower_right_block = gs.eye(3)
lower_left_block = gs.zeros((3, 4))
top = gs.hstack((upper_left_block, upper_right_block))
bottom = gs.hstack((lower_left_block, lower_right_block))
differential = gs.vstack((top, bottom))
differential = gs.expand_dims(differential, axis=0)
grad = gs.einsum('ni,nij->ni', grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def grad(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""
Closed-form for the gradient of pose_loss.
:return: tangent vector at point y_pred.
"""
if y_pred.ndim == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if y_true.ndim == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SE3, metric)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred_pose = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)
differential = gs.zeros((1, 6, 7))
upper_left_block = gs.zeros((1, 3, 4))
lower_right_block = gs.zeros((1, 3, 3))
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:4]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm**2 + quat_scalar**2
# TODO(nina): check that this sq norm is 1?
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = -2 * quat_vec / (quat_sq_norm)
differential_vec = (
2 *
(quat_scalar / quat_sq_norm - 2 * quat_arctan2 / quat_vec_norm) *
gs.outer(quat_vec, quat_vec) / quat_vec_norm**2 +
2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
upper_left_block[0, :, :1] = differential_scalar.transpose()
upper_left_block[0, :, 1:] = differential_vec
lower_right_block[0, :, :] = gs.eye(3)
differential[0, :3, :4] = upper_left_block
differential[0, 3:, 4:] = lower_right_block
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def convert_to_half_plane_coordinates(points):
disk_coords = points[:, 1:] / (1 + points[:, :1])
disk_x = disk_coords[:, 0]
disk_y = disk_coords[:, 1]
denominator = disk_x ** 2 + (1 - disk_y) ** 2
coords_0 = gs.expand_dims(2 * disk_x / denominator, axis=1)
coords_1 = gs.expand_dims((1 - disk_x ** 2 - disk_y ** 2) / denominator, axis=1)
half_plane_coords = gs.concatenate([coords_0, coords_1], axis=1)
return half_plane_coords
def convert_to_klein_coordinates(points):
poincare_coords = points[:, 1:] / (1 + points[:, :1])
poincare_radius = gs.linalg.norm(poincare_coords, axis=1)
poincare_angle = gs.arctan2(poincare_coords[:, 1], poincare_coords[:, 0])
klein_radius = 2 * poincare_radius / (1 + poincare_radius ** 2)
klein_angle = poincare_angle
coords_0 = gs.expand_dims(klein_radius * gs.cos(klein_angle), axis=1)
coords_1 = gs.expand_dims(klein_radius * gs.sin(klein_angle), axis=1)
klein_coords = gs.concatenate([coords_0, coords_1], axis=1)
return klein_coords
def exp(self, tangent_vec, base_point):
"""Compute the Riemannian exponential of a tangent vector.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, dim]
Tangent vector at a base point.
base_point : array-like, shape=[n_samples, dim]
Point in hyperbolic space.
Returns
-------
exp : array-like, shape=[n_samples, dim]
Point in hyperbolic space equal to the Riemannian exponential
of tangent_vec at the base point.
"""
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
norm_base_point =\
gs.expand_dims(gs.linalg.norm(base_point, axis=-1),
axis=-1)
den = 1 - norm_base_point**2
norm_tan =\
gs.expand_dims(gs.linalg.norm(tangent_vec, axis=-1),
axis=-1)
lambda_base_point = 1 / den
zero_tan =\
gs.isclose(gs.sum(tangent_vec * tangent_vec, axis=-1), 0.)
if gs.any(zero_tan):
if norm_tan[zero_tan].shape[0] != 0:
norm_tan[zero_tan] = EPSILON
direction = gs.einsum('...i,...k->...i', tangent_vec, 1 / norm_tan)
factor = gs.tanh(lambda_base_point * norm_tan)
exp = self.mobius_add(base_point, direction * factor)
zero_tan =\
gs.isclose(gs.sum(tangent_vec * tangent_vec, axis=-1), 0.)
if gs.any(zero_tan):
if exp[zero_tan].shape[0] != 0:
exp[zero_tan] = base_point[zero_tan]
return exp
def gmm_pdf(
data, means, variances, norm_func,
metric, variances_range, norm_func_var):
"""Return the separate probability density function of GMM.
The probability density function is computed for
each component of the GMM separately (i.e., mixture coefficients
are not taken into account).
Parameters
----------
data : array-like, shape=[n_samples, dim]
Points at which the GMM probability density is computed.
means : array-like, shape=[n_gaussians, dim]
Means of each component of the GMM.
variances : array-like, shape=[n_gaussians,]
Variances of each component of the GMM.
norm_func : function
Normalisation factor function.
metric : function
Distance function associated with the used metric.
Returns
-------
pdf : array-like, shape=[n_samples, n_gaussians,]
Probability density function computed at each data
sample and for each component of the GMM.
"""
data_length, _, _ = data.shape + (means.shape[0],)
variances_expanded = gs.expand_dims(variances, 0)
variances_expanded = gs.repeat(variances_expanded, data_length, 0)
variances_flatten = variances_expanded.flatten()
distances = -(metric.dist_broadcast(data, means) ** 2)
distances = gs.reshape(distances, (data.shape[0] * variances.shape[0]))
num = gs.exp(
distances / (2 * variances_flatten ** 2))
den = norm_func(variances, variances_range, norm_func_var)
den = gs.expand_dims(den, 0)
den = gs.repeat(den, data_length, axis=0).flatten()
pdf = num / den
pdf = gs.reshape(
pdf, (data.shape[0], means.shape[0]))
return pdf
def weighted_gmm_pdf(mixture_coefficients,
mesh_data,
means,
variances,
metric):
"""Return the probability density function of a GMM.
Parameters
----------
mixture_coefficients : array-like, shape=[n_gaussians,]
Coefficients of the Gaussian mixture model.
mesh_data : array-like, shape=[n_precision, dim]
Points at which the GMM probability density is computed.
means : array-like, shape=[n_gaussians, dim]
Means of each component of the GMM.
variances : array-like, shape=[n_gaussians,]
Variances of each component of the GMM.
metric : function
Distance function associated with the used metric.
Returns
-------
weighted_pdf : array-like, shape=[n_precision, n_gaussians,]
Probability density function computed for each point of
the mesh data, for each component of the GMM.
"""
distance_to_mean = metric.dist_broadcast(mesh_data, means)
variances_units = gs.expand_dims(variances, 0)
variances_units = gs.repeat(
variances_units, distance_to_mean.shape[0], axis=0)
distribution_normal = gs.exp(
-(distance_to_mean ** 2) / (2 * variances_units ** 2))
zeta_sigma = PI_2_3 * variances
zeta_sigma = zeta_sigma * gs.exp(
(variances ** 2 / 2) * gs.erf(variances / gs.sqrt(2)))
result_num = gs.expand_dims(mixture_coefficients, 0)
result_num = gs.repeat(
result_num, len(distribution_normal), axis=0)
result_num = result_num * distribution_normal
result_denum = gs.expand_dims(zeta_sigma, 0)
result_denum = gs.repeat(
result_denum, len(distribution_normal), axis=0)
weighted_pdf = result_num / result_denum
return weighted_pdf
def loss(example_embedding, context_embedding, negative_embedding, manifold):
"""Compute loss and grad.
Compute loss and grad given embedding of the current example,
embedding of the context and negative sampling embedding.
"""
n_edges, dim =\
negative_embedding.shape[0], example_embedding.shape[-1]
example_embedding = gs.expand_dims(example_embedding, 0)
context_embedding = gs.expand_dims(context_embedding, 0)
positive_distance =\
manifold.metric.squared_dist(
example_embedding, context_embedding)
positive_loss =\
log_sigmoid(-positive_distance)
reshaped_example_embedding =\
gs.repeat(example_embedding, n_edges, axis=0)
negative_distance =\
manifold.metric.squared_dist(
reshaped_example_embedding, negative_embedding)
negative_loss = log_sigmoid(negative_distance)
total_loss = -(positive_loss + negative_loss.sum())
positive_log_sigmoid_grad =\
-grad_log_sigmoid(-positive_distance)
positive_distance_grad =\
grad_squared_distance(example_embedding, context_embedding)
positive_grad =\
gs.repeat(positive_log_sigmoid_grad, dim, axis=-1)\
* positive_distance_grad
negative_distance_grad =\
grad_squared_distance(reshaped_example_embedding, negative_embedding)
negative_distance = gs.to_ndarray(negative_distance, to_ndim=2, axis=-1)
negative_log_sigmoid_grad =\
grad_log_sigmoid(negative_distance)
negative_grad = negative_log_sigmoid_grad\
* negative_distance_grad
example_grad = -(positive_grad + negative_grad.sum(axis=0))
return total_loss, example_grad
def _ball_to_extrinsic_coordinates(point):
"""Convert ball to extrinsic coordinates.
Convert the parameterization of a point in hyperbolic space
from its poincare ball model coordinates, to the extrinsic
coordinates.
Parameters
----------
point : array-like, shape=[n_samples, dimension]
Point in hyperbolic space in Poincare ball coordinates.
Returns
-------
extrinsic : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space in extrinsic coordinates.
"""
squared_norm = gs.sum(point**2, -1)
denominator = 1 - squared_norm
t = gs.to_ndarray((1 + squared_norm) / denominator, to_ndim=2, axis=1)
expanded_denominator = gs.expand_dims(denominator, -1)
expanded_denominator = gs.repeat(expanded_denominator, point.shape[-1],
-1)
intrinsic = (2 * point) / expanded_denominator
return gs.concatenate([t, intrinsic], -1)
def _update_medoid_indexes(self, distances, labels, medoid_indices):
for cluster in range(self.n_clusters):
cluster_index = gs.where(labels == cluster)[0]
if len(cluster_index) == 0:
logging.warning('One cluster is empty.')
continue
in_cluster_distances = distances[
cluster_index,
gs.expand_dims(cluster_index, axis=-1)]
in_cluster_all_costs = gs.sum(in_cluster_distances, axis=1)
min_cost_index = gs.argmin(in_cluster_all_costs)
min_cost = in_cluster_all_costs[min_cost_index]
current_cost = in_cluster_all_costs[gs.argmax(
cluster_index == medoid_indices[cluster])]
if min_cost < current_cost:
medoid_indices[cluster] = cluster_index[min_cost_index]
def vector_from_skew_matrix(self, skew_mat):
"""Derive a vector from the skew-symmetric matrix.
In 3D, compute the vector defining the cross product
associated to the skew-symmetric matrix skew mat.
Parameters
----------
skew_mat : array-like, shape=[..., n, n]
Returns
-------
vec : array-like, shape=[..., dim]
"""
n_skew_mats, _, _ = skew_mat.shape
vec_dim = self.dim
vec = gs.zeros((n_skew_mats, vec_dim))
if self.n == 2: # SO(2)
vec = skew_mat[:, 0, 1]
vec = gs.expand_dims(vec, axis=1)
elif self.n == 3: # SO(3)
vec_1 = gs.to_ndarray(skew_mat[:, 2, 1], to_ndim=2, axis=1)
vec_2 = gs.to_ndarray(skew_mat[:, 0, 2], to_ndim=2, axis=1)
vec_3 = gs.to_ndarray(skew_mat[:, 1, 0], to_ndim=2, axis=1)
vec = gs.concatenate([vec_1, vec_2, vec_3], axis=1)
return vec
def norm_factor_gradient(self, variances):
"""Compute normalization factor and its gradient.
Compute normalization factor given current variance
and dimensionality.
Parameters
----------
variances : array-like, shape=[n]
Value of variance.
Returns
-------
norm_factor : array-like, shape=[n]
Normalisation factor.
norm_factor_gradient : array-like, shape=[n]
Gradient of the normalization factor.
"""
variances = gs.transpose(gs.to_ndarray(variances, to_ndim=2))
dim_range = gs.arange(0, self.dim, 1.0)
alpha = self._compute_alpha(dim_range)
binomial_coefficient = gs.ones(self.dim)
binomial_coefficient[1:] = (self.dim - 1 + 1 - dim_range[1:]) / dim_range[1:]
binomial_coefficient = gs.cumprod(binomial_coefficient)
beta = ((-gs.ones(self.dim)) ** dim_range) * binomial_coefficient
sigma_repeated = gs.repeat(variances, self.dim, -1)
prod_alpha_sigma = gs.einsum("ij,j->ij", sigma_repeated, alpha)
term_2 = gs.exp((prod_alpha_sigma) ** 2) * (1 + gs.erf(prod_alpha_sigma))
term_1 = gs.sqrt(gs.pi / 2.0) * (1.0 / (2 ** (self.dim - 1)))
term_2 = gs.einsum("ij,j->ij", term_2, beta)
norm_factor = term_1 * variances * gs.sum(term_2, axis=-1, keepdims=True)
grad_term_1 = 1 / variances
grad_term_21 = 1 / gs.sum(term_2, axis=-1, keepdims=True)
grad_term_211 = (
gs.exp((prod_alpha_sigma) ** 2)
* (1 + gs.erf(prod_alpha_sigma))
* gs.einsum("ij,j->ij", sigma_repeated, alpha**2)
* 2
)
grad_term_212 = gs.repeat(
gs.expand_dims((2 / gs.sqrt(gs.pi)) * alpha, axis=0),
variances.shape[0],
axis=0,
)
grad_term_22 = grad_term_211 + grad_term_212
grad_term_22 = gs.einsum("ij, j->ij", grad_term_22, beta)
grad_term_22 = gs.sum(grad_term_22, axis=-1, keepdims=True)
norm_factor_gradient = grad_term_1 + (grad_term_21 * grad_term_22)
return gs.squeeze(norm_factor), gs.squeeze(norm_factor_gradient)
def sample_y(self, X, n_samples=1, random_state=0):
"""Draw samples from Wrapped Gaussian process and evaluate at X.
A fitted Wrapped Gaussian process can be use to sample
values through the following steps:
- Use the stored Gaussian process regression on the dataset
to sample tangent values
- Compute the base-points using the prior
- Flatten (and repeat if needed) both the base-points and the
tangent samples to benefit from vectorized computation.
- Map the tangent samples on the manifold via the metric's exp with the
flattened and repeated base-points yielded by the prior
Parameters
----------
X : array-like of shape (n_samples_X, n_features) or list of object
Query points where the WGP is evaluated.
n_samples : int, default=1
Number of samples drawn from the Wrapped Gaussian process per query point.
random_state : int, RandomState instance or None, default=0
Determines random number generation to randomly draw samples.
Pass an int for reproducible results across multiple function
calls.
Returns
-------
y_samples : ndarray of shape (n_samples_X, n_samples), or \
(n_samples_X, n_targets, n_samples)
Values of n_samples samples drawn from wrapped Gaussian process and
evaluated at query points.
"""
tangent_samples = self._euclidean_gpr.sample_y(X, n_samples,
random_state)
tangent_samples = gs.cast(tangent_samples, dtype=X.dtype)
# flatten the samples
tangent_samples = gs.reshape(gs.transpose(tangent_samples, [0, 2, 1]),
(-1, *self.y_train_shape_))
# generate the base_points
base_points = self.prior(X)
# repeat the base points in order to match the tangent samples
base_points = gs.repeat(gs.expand_dims(base_points, 2),
n_samples,
axis=2)
# flatten the base_points
base_points = gs.reshape(gs.transpose(base_points, [0, 2, 1]),
(-1, *self.y_train_shape_))
# get the flattened samples
y_samples = self.metric.exp(tangent_samples, base_point=base_points)
y_samples = gs.transpose(
gs.reshape(y_samples,
(X.shape[0], n_samples, *self.y_train_shape_)),
[0, 2, 1],
)
return y_samples
def log(self, point, base_point):
"""Compute Riemannian logarithm of a point wrt a base point.
If point_type = 'poincare' then base_point belongs
to the Poincare ball and point is a vector in the Euclidean
space of the same dimension as the ball.
Parameters
----------
point : array-like, shape=[..., dim]
Point in hyperbolic space.
base_point : array-like, shape=[..., dim]
Point in hyperbolic space.
Returns
-------
log : array-like, shape=[..., dim]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
add_base_point = self.mobius_add(-base_point, point)
norm_add =\
gs.expand_dims(gs.linalg.norm(
add_base_point, axis=-1), axis=-1)
norm_base_point =\
gs.expand_dims(gs.linalg.norm(
base_point, axis=-1), axis=-1)
log = (1 - norm_base_point**2) * gs.arctanh(norm_add)
mask_0 = gs.isclose(gs.squeeze(norm_add, axis=-1), 0.)
mask_non0 = ~mask_0
add_base_point = gs.assignment(
add_base_point,
gs.zeros_like(add_base_point[mask_0]),
mask_0)
add_base_point = gs.assignment(
add_base_point,
add_base_point[mask_non0] / norm_add[mask_non0],
mask_non0)
log = gs.einsum(
'...i,...j->...j', log, add_base_point)
return log
def find_variance_from_index(weighted_distances, variances_range,
phi_inv_var):
r"""Return the variance given weighted distances.
Parameters
----------
weighted_distances : array-like, shape=[n_gaussians,]
Mean of the weighted distances between training data
and current barycentres. The weights of each data sample
corresponds to the probability of belonging to a component
of the Gaussian mixture model.
variances_range : array-like, shape=[n_variances,]
Array of standard deviations.
phi_inv_var : array-like, shape=[n_variances,]
Array of the computed inverse of a function phi
whose expression is closed-form
:math:`\sigma\mapsto \sigma^3 \times \frac{d }
{\mathstrut d\sigma}\log \zeta_m(\sigma)'
where :math:'\sigma' denotes the variance
and :math:'\zeta' the normalization coefficient
and :math:'m' the dimension.
Returns
-------
var : array-like, shape=[n_gaussians,]
Estimated variances for each component of the GMM.
"""
n_gaussians, precision = \
weighted_distances.shape[0], variances_range.shape[0]
ref = gs.expand_dims(phi_inv_var, 0)
ref = gs.repeat(ref, n_gaussians, axis=0)
val = gs.expand_dims(weighted_distances, 1)
val = gs.repeat(val, precision, axis=1)
abs_difference = gs.abs(ref - val)
index = gs.argmin(abs_difference, -1)
var = variances_range[index]
return var
def regularize_tangent_vec_at_identity(self,
tangent_vec,
metric=None,
point_type=None):
"""
In 3D, regularize a tangent_vector by getting its norm at the identity,
determined by the metric, to be less than pi.
"""
if point_type is None:
point_type = self.default_point_type
if point_type == 'vector':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
if self.n == 3:
if metric is None:
metric = self.left_canonical_metric
tangent_vec_metric_norm = metric.norm(tangent_vec)
tangent_vec_canonical_norm = gs.linalg.norm(tangent_vec,
axis=1)
if gs.ndim(tangent_vec_canonical_norm) == 1:
tangent_vec_canonical_norm = gs.expand_dims(
tangent_vec_canonical_norm, axis=1)
mask_norm_0 = gs.isclose(tangent_vec_metric_norm, 0)
mask_canonical_norm_0 = gs.isclose(tangent_vec_canonical_norm,
0)
mask_0 = mask_norm_0 | mask_canonical_norm_0
mask_else = ~mask_0
mask_0 = gs.squeeze(mask_0, axis=1)
mask_else = gs.squeeze(mask_else, axis=1)
coef = gs.empty_like(tangent_vec_metric_norm)
regularized_vec = tangent_vec
regularized_vec[mask_0] = tangent_vec[mask_0]
coef[mask_else] = (tangent_vec_metric_norm[mask_else] /
tangent_vec_canonical_norm[mask_else])
regularized_vec[mask_else] = self.regularize(
coef[mask_else] * tangent_vec[mask_else])
regularized_vec[mask_else] = (regularized_vec[mask_else] /
coef[mask_else])
else:
# TODO(nina): regularization needed in nD?
regularized_vec = tangent_vec
elif point_type == 'matrix':
# TODO(nina): regularization in terms
# of skew-symmetric matrices?
regularized_vec = tangent_vec
return regularized_vec
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SO3, metric)
if representation == 'quaternion':
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = - 2 * quat_vec / (quat_sq_norm)
differential_scalar = gs.to_ndarray(differential_scalar, to_ndim=2)
differential_scalar = gs.transpose(differential_scalar)
differential_vec = (2 * (quat_scalar / quat_sq_norm
- 2 * quat_arctan2 / quat_vec_norm)
* (gs.einsum('ni,nj->nij', quat_vec, quat_vec)
/ quat_vec_norm ** 2)
+ 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_vec = gs.squeeze(differential_vec)
differential = gs.concatenate(
[differential_scalar, differential_vec],
axis=1)
y_pred = SO3.rotation_vector_from_quaternion(y_pred)
y_true = SO3.rotation_vector_from_quaternion(y_true)
grad = lie_group.grad(y_pred, y_true, SO3, metric)
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def loss(y_pred, y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""
Loss function given by a riemannian metric on a Lie group,
by default the left-invariant canonical metric.
"""
if gs.ndim(y_pred) == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if gs.ndim(y_true) == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
loss = lie_group.loss(y_pred, y_true, SE3, metric)
return loss
def test_unary_op_vec(self, func_name, a):
gs_fnc = get_backend_fnc(func_name)
res = gs_fnc(a)
a_expanded = gs.expand_dims(a, 0)
a_rep = gs.repeat(a_expanded, 2, axis=0)
res_a_rep = gs_fnc(a_rep)
for res_ in res_a_rep:
self.assertAllClose(res_, res)
def fit(self, X, max_iter=100):
"""Predict for each data point the closest center in terms of
riemannian_metric distance
Parameters
----------
X : array-like, shape=[n_samples, n_features]
Training data, where n_samples is the number of samples
and n_features is the number of features.
max_iter : Maximum number of iterations
Returns
-------
self : object
Return centroids array
"""
n_samples = X.shape[0]
belongs = gs.zeros(n_samples)
self.centroids = [
gs.expand_dims(X[randint(0, n_samples - 1)], 0)
for i in range(self.n_clusters)
]
self.centroids = gs.concatenate(self.centroids)
index = 0
while index < max_iter:
index += 1
dists = [
gs.to_ndarray(
self.riemannian_metric.dist(self.centroids[i], X), 2, 1)
for i in range(self.n_clusters)
]
dists = gs.hstack(dists)
belongs = gs.argmin(dists, 1)
old_centroids = gs.copy(self.centroids)
for i in range(self.n_clusters):
fold = gs.squeeze(X[belongs == i])
if len(fold) > 0:
self.centroids[i] = self.riemannian_metric.mean(fold)
else:
self.centroids[i] = X[randint(0, n_samples - 1)]
centroids_distances = self.riemannian_metric.dist(
old_centroids, self.centroids)
if gs.mean(centroids_distances) < self.tol:
if self.verbose > 0:
print("Convergence Reached after ", index, " iterations")
return gs.copy(self.centroids)
return gs.copy(self.centroids)
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 loss(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""Loss function given by a Riemannian metric on a Lie group.
Parameters
----------
y_pred : array-like
Prediction on SE(3).
y_true : array-like
Ground-truth on SE(3).
metric : RiemannianMetric
Metric used to compute the loss and gradient.
representation : str, {'vector', 'matrix'}
Representation chosen for points in SE(3).
Returns
-------
lie_loss : array-like
Loss using the Riemannian metric.
"""
if gs.ndim(y_pred) == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if gs.ndim(y_true) == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
lie_loss = lie_group.loss(y_pred, y_true, SE3, metric)
if gs.ndim(lie_loss) == 2:
lie_loss = gs.squeeze(lie_loss, axis=1)
if gs.ndim(lie_loss) == 1 and gs.shape(lie_loss)[0] == 1:
lie_loss = gs.squeeze(lie_loss, axis=0)
return lie_loss
def log(self, point, base_point):
"""Compute Riemannian logarithm of a point wrt a base point.
If point_type = 'poincare' then base_point belongs
to the Poincare ball and point is a vector in the Euclidean
space of the same dimension as the ball.
Parameters
----------
point : array-like, shape=[n_samples, dim]
Point in hyperbolic space.
base_point : array-like, shape=[n_samples, dim]
Point in hyperbolic space.
Returns
-------
log : array-like, shape=[n_samples, dim]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
base_point = gs.to_ndarray(base_point, to_ndim=2)
point = gs.to_ndarray(point, to_ndim=2)
add_base_point = self.mobius_add(-base_point, point)
norm_add =\
gs.expand_dims(gs.linalg.norm(
add_base_point, axis=-1), axis=-1)
norm_base_point =\
gs.expand_dims(gs.linalg.norm(
base_point, axis=-1), axis=-1)
log = (1 - norm_base_point**2) * gs.arctanh(norm_add)
log = gs.einsum('...i,...j->...j', log, (add_base_point / norm_add))
mask_0 = gs.isclose(gs.squeeze(norm_add, axis=-1), 0.)
if gs.any(mask_0):
log[mask_0] = 0
return log
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SO3, metric)
if representation == 'quaternion':
differential = gs.zeros((1, 6, 7))
differential = gs.zeros((1, 3, 4))
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = - 2 * quat_vec / (quat_sq_norm)
differential_vec = (2 * (quat_scalar / quat_sq_norm
- 2 * quat_arctan2 / quat_vec_norm)
* gs.outer(quat_vec, quat_vec) / quat_vec_norm ** 2
+ 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential[0, :, :1] = differential_scalar.transpose()
differential[0, :, 1:] = differential_vec
y_pred = SO3.rotation_vector_from_quaternion(y_pred)
y_true = SO3.rotation_vector_from_quaternion(y_true)
grad = lie_group.grad(y_pred, y_true, SO3, metric)
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def mobius_add(self, point_a, point_b):
r"""Compute the Mobius addition of two points.
Mobius addition operation that is a necessary operation
to compute the log and exp using the 'ball' representation.
.. math::
a\oplus b=\frac{(1+2\langle a,b\rangle + ||b||^2)a+
(1-||a||^2)b}{1+2\langle a,b\rangle + ||a||^2||b||^2}
Parameters
----------
point_a : array-like, shape=[n_samples, dim]
Point in hyperbolic space.
point_b : array-like, shape=[n_samples, dim]
Point in hyperbolic space.
Returns
-------
mobius_add : array-like, shape=[n_samples, 1]
Result of the Mobius addition.
"""
point_a = gs.to_ndarray(point_a, to_ndim=2)
point_b = gs.to_ndarray(point_b, to_ndim=2)
ball_manifold = PoincareBall(self.dim, scale=self.scale)
point_a_belong = ball_manifold.belongs(point_a)
point_b_belong = ball_manifold.belongs(point_b)
if (not gs.all(point_a_belong) or not gs.all(point_b_belong)):
raise NameError("Points do not belong to the Poincare ball")
norm_point_a = gs.sum(point_a**2, axis=-1, keepdims=True)
norm_point_b = gs.sum(point_b**2, axis=-1, keepdims=True)
sum_prod_a_b = gs.einsum('...i,...i->...', point_a, point_b)
sum_prod_a_b = gs.expand_dims(sum_prod_a_b, axis=-1)
add_num_1 = 1 + 2 * sum_prod_a_b + norm_point_b
add_num_1 = gs.einsum('...i,...k->...k', add_num_1, point_a)
add_num_2 = gs.einsum('...i,...k->...k', (1 - norm_point_a), point_b)
add_nominator = add_num_1 + add_num_2
add_denominator = (1 + 2 * sum_prod_a_b + norm_point_a * norm_point_b)
mobius_add =\
gs.einsum('...i,...k->...i', add_nominator, 1 / add_denominator)
return mobius_add
