Source code for pns.base

"""Basic functions to handle data on a hypersphere.

Conventions for intrinsic coordiates are:
    - xi[..., 0]: azimuthal angle in [-pi, pi]
    - xi[..., 1:]: centered elevation angles in [-pi/2, pi/2]
"""

import numpy as np

__all__ = [
    "intrinsic_to_extrinsic",
    "extrinsic_to_intrinsic",
    "rotation_matrix",
    "exp_map",
    "log_map",
    "circle_mean",
]


[docs] def intrinsic_to_extrinsic(xi): """Convert hypersphere intrinsic coordinates to extrinsic. Parameters ---------- xi : array of shape (N, d) Hyperspherical coordinates. Returns ------- X : array of shape (N, d+1) Cartesian coordinates. Notes ----- Intrinsic coordiates are: - xi[..., 0]: azimuthal angle in [-pi, pi] - xi[..., 1:]: centered elevation angles in [-pi/2, pi/2] Examples -------- >>> import numpy as np >>> from pns.base import intrinsic_to_extrinsic >>> xi = np.array([[0, np.pi/2]]) >>> X = intrinsic_to_extrinsic(xi) >>> np.isclose(X, [[0, 0, 1]]) array([[ True, True, True]]) """ N, d = xi.shape X = np.zeros((N, d + 1)) for n in range(N): angles = xi[n] cos_elev = np.cumprod(np.cos(angles[1:][::-1]))[::-1] X[n, 0] = np.cos(angles[0]) * (cos_elev[0] if d > 1 else 1) X[n, 1] = np.sin(angles[0]) * (cos_elev[0] if d > 1 else 1) for k in range(1, d): X[n, k + 1] = np.sin(angles[k]) * (cos_elev[k] if k < d - 1 else 1) return X
[docs] def extrinsic_to_intrinsic(X): """Convert hypersphere extrinsic coordinates to intrinsic. Parameters ---------- X : array of shape (N, d+1) Cartesian coordinates. Returns ------- xi : array of shape (N, d) Hyperspherical coordinates. Examples -------- >>> from pns.base import extrinsic_to_intrinsic >>> from pns.util import circular_data >>> X = circular_data() >>> X.shape (100, 3) >>> extrinsic_to_intrinsic(X).shape (100, 2) """ N, D = X.shape d = D - 1 xi = np.zeros((N, d)) xi[:, 0] = np.arctan2(X[:, 1], X[:, 0]) for i in range(1, d - 1): denom = np.linalg.norm(X[:, i:], axis=1) xi[:, i] = np.arctan(X[:, i + 1] / denom) xi[:, -1] = np.arctan2(X[:, -1], np.linalg.norm(X[:, :-1], axis=1)) return xi
[docs] def rotation_matrix(v): r"""Rotation matrix. Moves :math:`v` to the north pole. Parameters ---------- v : (m+1) real array Unit-norm direction vector. Returns ------- (m+1, m+1) array of float64 Rotational matrix. Notes ----- This is the function :math:`R(v_k)` in the original paper. Examples -------- Moving ``[0, 1, 0]`` to the north pole, in other words, moving the north pole to ``[0, -1, 0]``. >>> from pns.base import rotation_matrix >>> from pns.util import unit_sphere, circular_data >>> X = circular_data() >>> R = rotation_matrix([0, 1, 0]) >>> X_rotated = X @ R.T >>> import matplotlib.pyplot as plt # doctest: +SKIP ... ax = plt.figure().add_subplot(projection='3d', computed_zorder=False) ... ax.plot_surface(*unit_sphere(), color='skyblue', alpha=0.6, edgecolor='gray') ... ax.scatter(*X.T) ... ax.scatter(*X_rotated.T) """ a = np.zeros((len(v),)) a[-1] = 1.0 b = v c = b - a * (a @ b) c_norm = np.linalg.norm(c) if c_norm == 0: R = np.eye(len(c)) else: c /= np.linalg.norm(c) A = np.outer(a, c) - np.outer(c, a) theta = np.arccos(v[-1]) Id = np.eye(len(A)) R = ( Id + np.sin(theta) * A + (np.cos(theta) - 1) * (np.outer(a, a) + np.outer(c, c)) ) return R
[docs] def exp_map(z): """Exponential map of hypersphere at (0, 0, ..., 0, 1). Parameters ---------- z : (N, d) real array Vectors on tangent space. Returns ------- (N, d+1) real array Points on d-sphere. """ norm = np.linalg.norm(z, axis=1)[..., np.newaxis] return np.hstack([np.sin(norm) / norm * z, np.cos(norm)])
[docs] def log_map(x): """Log map of hypersphere at (0, 0, ..., 0, 1). Parameters ---------- x : (N, d+1) real array Points on d-sphere. Returns ------- (N, d) real array Vectors on tangent space. """ thetas = np.arccos(x[:, -1:]) return thetas / np.sin(thetas) * x[:, :-1]
[docs] def circle_mean(X): """Frechet mean of data on a circle. Parameters ---------- X : real array of shape (N, 2) Returns ------- mean : array of shape (2,) Notes ----- Uses the algorithm [1]_ implemented by the Geomstats [2]_ package. Copyright (c) 2018 Nina Miolane References ---------- .. [1] Hotz, T. and S. F. Huckemann (2015), "Intrinsic means on the circle: Uniqueness, locus and asymptotics", Annals of the Institute of Statistical Mathematics 67 (1), 177-193. https://arxiv.org/abs/1108.2141 .. [2] https://github.com/geomstats/geomstats Examples -------- >>> import numpy as np >>> from pns.base import circle_mean >>> t = np.linspace(-np.pi, np.pi / 2, 10) >>> X = np.stack([np.cos(t), np.sin(t)]).T >>> mean = circle_mean(X) >>> import matplotlib.pyplot as plt # doctest: +SKIP ... plt.scatter(*X.T, label="data") ... plt.scatter(*mean, label="mean") ... plt.legend() ... plt.gca().set_aspect("equal") """ # Convert X to angles in [-pi, pi) X = np.arctan2(X[..., 1], X[..., 0])[..., np.newaxis] # CircleMean._circle_mean from geomstats mean0 = np.mean(X) var0 = np.sum((X - mean0) ** 2) N, _ = X.shape X = np.sort(X, axis=0) # CircleMean._circle_variances from geomstats means = (mean0 + np.linspace(0.0, 2 * np.pi, N + 1)[:-1]) % (2 * np.pi) means = np.where(means >= np.pi, means - 2 * np.pi, means) parts = np.array([np.sum(X) / N if means[0] < 0 else 0]) m_plus = means >= 0 left_sums = np.cumsum(X) right_sums = left_sums[-1] - left_sums i = np.arange(N, dtype=right_sums.dtype) j = i[1:] parts2 = right_sums[:-1] / (N - j) first_term = parts2[:1] parts2 = np.where(m_plus[1:], left_sums[:-1] / j, parts2) parts = np.concatenate([parts, first_term, parts2[1:]]) plus_vec = (4 * np.pi * i / N) * (np.pi + parts - mean0) - (2 * np.pi * i / N) ** 2 minus_vec = (4 * np.pi * (N - i) / N) * (np.pi - parts + mean0) - ( 2 * np.pi * (N - i) / N ) ** 2 minus_vec = np.where(m_plus, plus_vec, minus_vec) means = np.transpose(np.vstack([means, var0 + minus_vec])) frechet_mean = means[np.argmin(means[:, 1]), 0] cos = np.cos(frechet_mean) sin = np.sin(frechet_mean) return np.hstack([cos, sin])