morphomatics 4.0__py3-none-any.whl

morphomatics/manifold/sphere.py

@@ -0,0 +1,241 @@
+################################################################################
+#                                                                              #
+#   This file is part of the Morphomatics library                              #
+#       see https://github.com/morphomatics/morphomatics                       #
+#                                                                              #
+#   Copyright (C) 2024 Zuse Institute Berlin                                   #
+#                                                                              #
+#   Morphomatics is distributed under the terms of the MIT License.            #
+#       see $MORPHOMATICS/LICENSE                                              #
+#                                                                              #
+################################################################################
+
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+
+from morphomatics.manifold import Manifold, Metric
+from morphomatics.manifold.euclidean import euclidean_inner
+from morphomatics.manifold.metric import _eval_adjJacobi_embed
+
+
+class Sphere(Manifold):
+    """The sphere of [... x k x m]-tensors embedded in R(n+1)
+    Elements are represented as normalized (row) vectors of length n + 1.
+    """
+
+    def __init__(self, point_shape=(3,), structure='Canonical'):
+        name = 'Points with unit Frobenius norm in ' +\
+               'x'.join(map(str, point_shape)) + '-dim. space.'
+        dimension = np.prod(point_shape)-1
+        super().__init__(name, dimension, point_shape)
+        if structure:
+            getattr(self, f'init{structure}Structure')()
+
+    def tree_flatten(self):
+        children, aux = super().tree_flatten()
+        return children, aux+(self.point_shape,)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Specifies an unflattening recipe for PyTree registration."""
+        *aux_data, shape = aux_data
+        obj = cls(shape, structure=None)
+        obj.tree_unflatten_instance(aux_data, children)
+        return obj
+
+    def initCanonicalStructure(self):
+        """
+        Instantiate Sphere with canonical structure.
+        """
+        structure = Sphere.CanonicalStructure(self)
+        self._metric = structure
+        self._connec = structure
+
+    def rand(self, key: jax.Array):
+        p = jax.random.normal(key, self.point_shape)
+        return p / jnp.linalg.norm(p)
+
+    def randvec(self, X, key: jax.Array):
+        H = jax.random.normal(key, self.point_shape)
+        return H - jnp.dot(X.reshape(-1), H.reshape(-1)) * X
+
+    def zerovec(self):
+        return jnp.zeros(self.point_shape)
+
+    @staticmethod
+    def antipode(p):
+        return -p
+
+    @staticmethod
+    def normalize(X):
+        """Return Frobenius-normalized version of X in ambient space."""
+        return X / jnp.sqrt((X**2).sum() + np.finfo(np.float64).eps)
+
+    def proj(self, p, X):
+        return X - euclidean_inner(p, X) * p
+
+    class CanonicalStructure(Metric):
+        """
+        The Riemannian metric used is the induced metric from the embedding space (R^nxn)^k, i.e., this manifold is a
+        Riemannian submanifold of (R^3x3)^k endowed with the usual trace inner product.
+        """
+
+        def __init__(self, M):
+            """
+            Constructor.
+            """
+            self._M = M
+
+        def __str__(self):
+            return "Canonical structure"
+
+        @property
+        def typicaldist(self):
+            return np.pi
+
+        def inner(self, p, X, Y):
+            return euclidean_inner(X, Y)
+
+        def norm(self, p, X):
+            return jnp.sqrt(self.inner(p, X, X))
+
+        def flat(self, p, X):
+            """Lower vector X at p with the metric"""
+            return X
+
+        def sharp(self, p, dX):
+            """Raise covector dX at p with the metric"""
+            return dX
+
+        def egrad2rgrad(self, p, X):
+            return self._M.proj(p, X)
+
+        def ehess2rhess(self, p, G, H, X):
+            """Converts the Euclidean gradient P_G and Hessian H of a function at
+            a point p along a tangent vector X to the Riemannian Hessian
+            along X on the manifold.
+            """
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def retr(self, p, X):
+            return self.exp(p, X)
+
+        def exp(self, p, X):
+            # numerical safeguard
+            p = Sphere.normalize(p)
+            X = self._M.proj(p, X)
+
+            def full_exp(sqn):
+                n = jnp.sqrt(sqn + jnp.finfo(jnp.float64).eps)
+                return jnp.cos(n) * p + jnp.sinc(n/jnp.pi) * X
+
+            def trunc_exp(sqn):
+                #return (1-sqn/2+sqn**2/24-sqn**3/720) * p + (1-sqn/6+sqn**2/120-sqn**3/5040) * X
+                # 4th-order approximation
+                return (1-sqn/2+sqn**2/24) * p + (1-sqn/6+sqn**2/120) * X
+
+            sq_norm = (X ** 2).sum()
+            q = jax.lax.cond(sq_norm < 1e-6, trunc_exp, full_exp, sq_norm)
+            return Sphere.normalize(q)
+
+        def log(self, p, q):
+
+            def full_log(a2):
+                a = jnp.sqrt(a2 + jnp.finfo(jnp.float64).eps)
+                return 1/jnp.sinc(a/jnp.pi) * q - a/jnp.tan(a) * p
+
+            def trunc_log(a2):
+                return (1 + a2/6 + 7*a2**2/360 + 31*a2**3/15120) * q - (1 - a2/3 - a2**2/45 - a2**3/945) * p
+                #return (1 + a**2/6 + 7*a**4/360) * q - (1 - a**2/3 - a**4/45) * p
+
+            sqd = self.squared_dist(p, q)
+            return jax.lax.cond(sqd < 1e-6, trunc_log, full_log, sqd)
+
+        def curvature_tensor(self, p, X, Y, Z):
+            """Evaluates the curvature tensor R of the connection at p on the vectors X, Y, Z. With nabla_X Y denoting the
+            covariant derivative of Y in direction X and [] being the Lie bracket, the convention
+                R(X,Y)Z = (nabla_X nabla_Y) Z - (nabla_Y nabla_X) Z - nabla_[X,Y] Z
+            is used.
+            """
+            return (Y*Z).sum() * X - (X*Z).sum() * Y
+
+        def geopoint(self, p, q, t):
+            return self.exp(p, t * self.log(p, q))
+
+        def transp(self, p, q, X):
+            d2 = self.squared_dist(p, q)
+            def do_transp(V):
+                log_p_q = self.log(p, q)
+                return V - self.inner(p, log_p_q, V)/d2 * (log_p_q + self.log(q, p))
+            return jax.lax.cond(d2 < 1e-6, lambda X: X, do_transp, X)
+
+        def pairmean(self, p, q):
+            return self.geopoint(p, q, .5)
+
+        def dist(self, p, q):
+            inner = (p * q).sum()
+            return jax.lax.cond(jnp.abs(inner) >= 1, lambda i: (i < 0)*jnp.pi, jnp.arccos, jnp.clip(inner, -1, 1))
+
+        def squared_dist(self, p, q):
+            inner = (p * q).sum()
+            return jax.lax.cond(inner > 1-1e-6, lambda _: jnp.sum((q-p)**2), lambda i: jnp.arccos(i)**2, jnp.clip(inner, None, 1-1e-6))
+
+        def jacobiField(self, p, q, t, X):
+            phi = self.dist(p, q)
+            v = self.log(p, q)
+            gamTS = self.exp(p, t*v)
+
+            v = v / phi
+            Xtan_norm = self.inner(p, X, v)
+            Xtan = Xtan_norm * v
+            Xorth = X - Xtan
+
+            # tangential component of J: (1-t) * transp(p, gamTS, Xtan)
+            Jtan = Xtan_norm / phi * self.log(gamTS, q)
+            return gamTS, (jnp.sin((1 - t) * phi) / jnp.sin(phi)) * Xorth + Jtan
+
+        def _adjJacobi(self, p, q, t, w):
+            # alternative version to adjJacobi relying on automatic differentiation
+            return self._M.proj(p, _eval_adjJacobi_embed(self, p, q, t, w))
+
+        def adjJacobi(self, p, q, t, w):
+            """Evaluate an adjoint Jacobi field.
+
+            The decomposition of the curvature operator and the fact that only two of its eigenvectors are necessary is
+            used in the algorithm.
+
+            :param p: element of the hyperboloid
+            :param q: element of the hyperboloid
+            :param t: scalar in [0,1]
+            :param w: tangent vector at gamma(t;p,q)
+            :returns: tangent vector at p
+            """
+            dist = self.dist(p, q)
+
+            def _eval(dW):
+                d, W = dW
+                # all computations can be done at p, so only w has to be parallel translated
+                W = self.transp(self.geopoint(p, q, t), p, W)
+
+                # first eigenvector is normalized tangent of the geodesic -> corresponding eigenvalue is 0
+                T = self.log(p, q) / jnp.clip(d, 1e-6)  # clipping only for NAN debugging
+
+                # second eigenvector is Gram Schmidt orthonormalization of W against T -> corresponding eigenvalue is 1
+                b1 = self.inner(p, W, T)
+                U = W - b1 * T
+                # Check whether W is (numerically) parallel to T;
+                # then, the adoint Jacobi field is only a scaling of the parallel transported tangent.
+                U = jax.lax.cond(jnp.linalg.norm(U) > 1e-3,
+                                 lambda v: v / jnp.clip(self.norm(p, v), 1e-3),  # clipping only for NAN debugging
+                                 lambda v: jnp.zeros_like(v), U)
+
+                b2 = self.inner(p, W, U)
+
+                a1 = 1 - t  # corresponds to the eigenvalue 0
+                a2 = jnp.sin((1 - t) * d) / jnp.sin(d)  # corresponds to the eigenvalue 1
+
+                return a1 * b1 * T + a2 * b2 * U
+
+            return jax.lax.cond(dist > 1e-6, _eval, lambda args: 1/jnp.clip(1-t, 1e-3) * args[-1], (dist, w))

morphomatics/manifold/tangent_bundle.py

@@ -0,0 +1,337 @@
+################################################################################
+#                                                                              #
+#   This file is part of the Morphomatics library                              #
+#       see https://github.com/morphomatics/morphomatics                       #
+#                                                                              #
+#   Copyright (C) 2024 Zuse Institute Berlin                                   #
+#                                                                              #
+#   Morphomatics is distributed under the terms of the MIT License.            #
+#       see $MORPHOMATICS/LICENSE                                              #
+#                                                                              #
+################################################################################
+
+import jax
+import jax.numpy as jnp
+
+from typing import Callable, List
+
+
+from morphomatics.manifold import Manifold, Metric
+
+
+class TangentBundle(Manifold):
+    """Tangent Bundle TM of a smooth manifold M
+
+        Elements of TM are modelled as arrays of the form [2, M.point_shape]
+    """
+
+    def __init__(self, M: Manifold, structure: str = 'Sasaki'):
+        point_shape = tuple([2, *M.point_shape])
+        name = 'Tangent bundle of ' + M.__str__() + '.'
+        dimension = 2 * M.dim
+        super().__init__(name, dimension, point_shape)
+        self._base_manifold = M
+        if structure:
+            getattr(self, f'init{structure}Structure')()
+
+    def tree_flatten(self):
+        children, aux = super().tree_flatten()
+        return children+(self.base_manifold,), aux
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Specifies an unflattening recipe for PyTree registration."""
+        *children, M = children
+        obj = cls(M, structure=None)
+        obj.tree_unflatten_instance(aux_data, children)
+        return obj
+
+    @property
+    def base_manifold(self) -> Manifold:
+        """Return the base manifold of whose tangent bundle is TM """
+        return self._base_manifold
+
+    def bundle_projection(self, pu: jnp.array) -> jnp.array:
+        """Canonical projection of the tangent bundle onto its base manifold"""
+        return pu[0]
+
+    def initSasakiStructure(self):
+        """
+        Instantiate the tangent bundle with Sasaki structure.
+        """
+        structure = TangentBundle.SasakiStructure(self)
+        self._metric = structure
+        self._connec = structure
+
+    def rand(self, key: jax.Array) -> jnp.array:
+        """Random element of TM"""
+        k1, k2 = jax.random.split(key, 2)
+        p = self._base_manifold.rand(k1)
+        u = self._base_manifold.randvec(k2)
+        return jnp.stack((p, u))
+
+    def randvec(self, pu: jnp.array, key: jax.Array) -> jnp.array:
+        """Random vector in the tangent space of the point pu
+
+        :param pu: element of TM
+        :return: tangent vector at pu
+        """
+        k1, k2 = jax.random.split(key, 2)
+        p = self.bundle_projection(pu)
+        return jnp.stack((self._base_manifold.randvec(p, k1), self._base_manifold.randvec(p, k2)))
+
+    def zerovec(self) -> jnp.array:
+        """Zero vector in any tangen space
+        """
+        return jnp.zeros(self.point_shape)
+
+    def proj(self, pu: jnp.array, vw: jnp.array) -> jnp.array:
+        raise NotImplementedError('This function has not been implemented yet.')
+
+    class SasakiStructure(Metric):
+        """
+        This class implements the Sasaki metric: The natural metric on the tangent bundle TM of a Riemannian manifold M.
+
+        The Sasaki metric is characterized by the following three properties:
+         * the canonical projection of TM becomes a Riemannian submersion,
+         * parallel vector fields along curves are orthogonal to their fibres, and
+         * its restriction to any tangent space is Euclidean.
+
+        Geodesic computations are realized via a discrete formulation of the geodesic equation on TM that involve
+        geodesics, parallel translation, and the curvature tensor on the base manifold M; for details see
+            Muralidharan, P., & Fletcher, P. T. (2012, June).
+            Sasaki metrics for analysis of longitudinal data on manifolds.
+            In 2012 IEEE conference on computer vision and pattern recognition (pp. 1027-1034). IEEE.
+            https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4270017/
+        """
+
+        def __init__(self, TM: Manifold, Ns: int = 10):
+            """
+            Constructor.
+
+            :param TM: TangentBundle object
+            :param Ns: scalar that determines the number of discretization steps used in the approximation of the
+            exponential and logarithm maps
+            """
+            self._TM = TM
+            self.Ns = Ns
+
+        def __str__(self) -> str:
+            return "Sasaki structure"
+
+        @property
+        def typicaldist(self):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def inner(self, pu: jnp.array, vw: jnp.array, xy: jnp.array) -> float:
+            """Inner product between two tangent vectors at point in TM.
+
+            :param pu: element of _TM
+            :param vw: tangent vector at pv
+            :param xy: tangent vector at pv
+            :return: inner product (scalar) of vw and xy
+            """
+            p = self._TM.bundle_projection(pu)
+            # compute Sasaki inner product via metric of the underlying manifold
+            base_metric = self._TM.base_manifold.metric.inner
+            return base_metric(p, vw[0], xy[0]) + base_metric(p, vw[1], xy[1])
+
+        def flat(self, pu: jnp.array, xy: jnp.array) -> jnp.array:
+            """Lower vector xy at pu with the metric"""
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def sharp(self, pu: jnp.array, d_xy: jnp.array) -> jnp.array:
+            """Raise covector d_xy at pu with the metric"""
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def egrad2rgrad(self, pu: jnp.array, vw: jnp.array) -> jnp.array:
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def ehess2rhess(self, pu, G, H, vw):
+            """Converts the Euclidean gradient G and Hessian H of a function at
+            a point pv along a tangent vector uw to the Riemannian Hessian
+            along X on the manifold.
+            """
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def geodesic_discrete(self, pu: jnp.array, qr: jnp.array) -> List[jnp.array]:
+            """
+            Compute Sasaki geodesic employing a variational time discretization.
+
+            :param pu: element of TM
+            :param qr: element of TM
+            :return: array-like, shape=[n_steps + 1, 2, M.shape]
+                Discrete geodesic x(s)=(p(s), u(s)) in Sasaki metric connecting
+                pu = x(0) and qr = x(1).
+            """
+
+            Ns = self.Ns
+            base_connection = self._TM.base_manifold.connec
+            par_trans = base_connection.transp
+            p0, u0 = pu[0], pu[1]
+            pL, uL = qr[0], qr[1]
+
+            def grad(c):
+                """ Gradient of path energy for discrete geodesic c """
+                def grad_i(puP, pu, puN):
+                    v, w = base_connection.log(pu[0], puN[0]), par_trans(puN[0], pu[0], puN[1]) - pu[1]
+                    gp = .5 * (v + base_connection.log(pu[0], puP[0])) \
+                              + base_connection.curvature_tensor(pu[0], pu[1], w, v)
+                    gu = w + par_trans(puP[0], pu[0], puP[1]) - pu[1]
+                    return jnp.array([gp, gu])
+                return -Ns * jax.vmap(grad_i)(c[:-2], c[1:-1], c[2:])
+
+            # Initial guess for gradient_descent
+            vw = base_connection.log(p0, pL)
+            s = jnp.linspace(0., 1., Ns + 1)
+            def init(t):
+                p_ini = base_connection.exp(p0, t * vw)
+                u_ini = (1 - t) * par_trans(p0, p_ini, u0) + t * par_trans(pL, p_ini, uL)
+                return jnp.array([p_ini, u_ini])
+            pu_ini = jax.vmap(init)(s[1:-1])
+            pu_ini = jnp.vstack([pu[None], pu_ini, qr[None]])
+
+            # Minimization by gradient descent
+            x = _gradient_descent(self._TM, pu_ini, grad)
+            # x = _gradient_descent(pu_ini, grad, self.exp)
+
+            return x
+
+        def exp(self, pu: jnp.array, vw: jnp.array) -> jnp.array:
+            """Compute the exponential of the Levi-Civita connection of the Sasaki metric.
+
+            Exponential map at pv of uw computed by
+            shooting a Sasaki geodesic using an Euler integration in TM.
+
+            :param pu: element of TM
+            :param vw: tangent vector at pv
+            :return: point at time 1 of the geodesic that starts in pu with initial velocity vw
+            """
+
+            base_connection = self._TM.base_manifold.connec
+            par_trans = self._TM.base_manifold.connec.transp
+            Ns = self.Ns
+            eps = 1 / Ns
+
+            def body(carry, _):
+                p, u, v, w = carry
+                p_ = base_connection.exp(p, eps * v)
+                u_ = par_trans(p, p_, u + eps * w)
+                v_ = par_trans(p, p_, v - eps * base_connection.curvature_tensor(p, u, w, v))
+                w_ = par_trans(p, p_, w)
+                return (p_, u_, v_, w_), None
+
+            (p, u, *_), _ = jax.lax.scan(body, (*pu, *vw), jnp.empty(Ns))
+
+            return jnp.stack([p, u])
+
+        def log(self, pu: jnp.array, qr: jnp.array) -> jnp.array:
+            """Compute the logarithm of the Levi-Civita connection of the Sasaki metric.
+
+            Logarithmic map at base_point p of pu computed by
+            iteratively relaxing a discretized geodesic between pu and qw.
+
+            For a derivation of the algorithm see https://opus4.kobv.de/opus4-zib/frontdoor/index/index/docId/8717.
+
+            :param pu: element of TM
+            :param qr: element of TM
+            :return: tangent vector at pu (inverse of exp)
+            """
+
+            base_connection = self._TM.base_manifold.connec
+            par_trans = base_connection.transp
+            Ns = self.Ns
+
+            def do_log(bs_pt, pt):
+                pu = self.geodesic_discrete(bs_pt, pt)
+                p1, u1 = pu[1][0], pu[1][1]
+                p0, u0 = bs_pt[0], bs_pt[1]
+                w = par_trans(p1, p0, u1) - u0
+                v = base_connection.log(p0, p1)
+                return Ns * jnp.array([v, w])
+
+            return do_log(pu, qr)
+
+        def curvature_tensor(self, pu, X, Y, Z):
+            """Evaluates the curvature tensor R of the connection at pu on the vectors X, Y, Z. With nabla_X Y denoting
+            the covariant derivative of Y in direction X and [] being the Lie bracket, the convention
+                R(X,Y)Z = (nabla_X nabla_Y) Z - (nabla_Y nabla_X) Z - nabla_[X,Y] Z
+            is used.
+            """
+            # TODO: Implement from "Curvature of the Induced Riemannian Metric on the Tangent Bundle of a Riemannian
+            #  Manifold" (1971) by Kowalski
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def geopoint(self, pu: jnp.array, qr: jnp.array, t: float) -> jnp.array:
+            """Evaluate geodesic in TM
+
+            :param pu: element of TM
+            :param qr: element of TM
+            :param t: scalar between 0 and 1
+            :return: element of the geodesic between pv and pu evaluated at time t
+            """
+            return self.exp(pu, t * self.log(pu, qr))
+
+        def retr(self, pu: jnp.array, vw: jnp.array) -> jnp.array:
+            return self.exp(pu, vw)
+
+        def transp(self, pu, qr, vw):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def pairmean(self, pu: jnp.array, qr: jnp.array) -> jnp.array:
+            """Fréchet mean of 2 point in TM
+
+            :param pu: element of TM
+            :param qr: element of TM
+            :return: Fréchet mean of pv and qw
+            """
+            return self.geopoint(pu, qr, .5)
+
+        def dist(self, pu: jnp.array, qr: jnp.array) -> float:
+            """Distance function that is induced on TM by the Sasaki metric
+
+            :param pu: element of TM
+            :param qr: element of TM
+            :return: distance between pu and qr in TM
+            """
+            vw = self.log(pu, qr)
+            return jnp.sqrt(self.inner(pu, vw, vw))
+
+        def jacobiField(self, p_A: jnp.array, p_B: jnp.array, t: float, X: jnp.array) -> jnp.array:
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def adjJacobi(self, p_A: jnp.array, p_B: jnp.array, t: float, X: jnp.array) -> jnp.array:
+            raise NotImplementedError('This function has not been implemented yet.')
+
+
+def _gradient_descent(TM: Manifold, x_ini: jnp.array, grad: Callable, step_size: float = 0.1, max_iter: int = 100,
+                      tol: float = 1e-6) -> jnp.array:
+    """Apply a gradient descent to compute the discrete geodesic in TM with the Sasaki structure.
+
+    :param TM: Tangent bundle
+    :param x_ini: initial discrete curve as list of points in TM
+    :param grad: gradient function
+    :param step_size: step size
+    :param max_iter: maximum number of iterations
+    :param tol: tolerance for convergence
+    :return: discrete geodesic as list of points in _TM
+    """
+
+    def body(args):
+        x, grad_norm, tol, i = args
+        grad_x = grad(x)
+
+        x.at[1:-1].set(jax.vmap(TM.connec.exp)(x[1:-1], -step_size * grad_x))
+        grad_norm = jnp.linalg.norm(grad_x)
+
+        return x, grad_norm, tol, i + 1
+
+    def cond(args):
+        _, g_norm, tol, i = args
+
+        c = jnp.array([g_norm > tol, i < max_iter])
+        return jnp.all(c)
+
+    x, g_norm, _, i = jax.lax.while_loop(cond, body, (x_ini, jnp.array(1.), tol, jnp.array(0)))
+
+    return x