morphomatics 4.0__py3-none-any.whl

morphomatics/manifold/hyperbolic_space.py

@@ -0,0 +1,271 @@
+################################################################################
+#                                                                              #
+#   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.metric import _eval_adjJacobi_embed
+
+
+class HyperbolicSpace(Manifold):
+    """n-dimensional hyperbolic space represented by the hyperboloid model. Elements are represented as (row) vectors of
+    length n + 1.
+    """
+
+    def __init__(self, point_shape=(3,), structure='Canonical'):
+        name = 'Points in ' +\
+               'x'.join(map(str, point_shape)) + '-dim. space ' + 'for which the Minkowski quadratic form is -1.'
+        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 = HyperbolicSpace.CanonicalStructure(self)
+        self._metric = structure
+        self._connec = structure
+
+    @staticmethod
+    def minkowski_inner(x, y):
+        """Minkowski semi-Riemannian metric on R^(n+1). The hyperboloid model of n-dimensional
+        hyperbolic space consists of all points x with minkowski_inner(x, x) = -1. The tangent space at x consists of
+        all vectors v in R^(n+1) with minkowski_inner(x, v) = 0."""
+        return -x[-1] * y[-1] + (x[:-1] * y[:-1]).sum()
+
+    @staticmethod
+    def project_to_manifold(x):
+        """Projection onto the hyperboloid."""
+        return x.at[-1].set(jnp.sqrt(1 + jnp.sum(x[:-1] ** 2)))
+
+    @staticmethod
+    def regularize(p):
+        """Regularize/project points that are slightly off the hyperboloid to avoid numerical problems. Only apply when
+        minkowski_inner(p, p) < 0."""
+        return p / jnp.sqrt(jnp.abs(HyperbolicSpace.minkowski_inner(p, p)) + np.finfo(np.float64).eps)
+
+    def rand(self, key: jax.Array):
+        p = jax.random.normal(key, self.point_shape)
+        return self.project_to_manifold(p)
+
+    def randvec(self, p, key: jax.Array):
+        H = jax.random.normal(key, self.point_shape)
+        return H + HyperbolicSpace.minkowski_inner(p, H) * p
+
+    def zerovec(self):
+        return jnp.zeros(self.point_shape)
+
+    def pole(self):
+        o = jnp.zeros(self.point_shape)
+        o = o.at[-1].set(1)
+        return o
+
+    def proj(self, p, H):
+        return H + HyperbolicSpace.minkowski_inner(p, H) * p
+
+    class CanonicalStructure(Metric):
+        """
+        The Riemannian metric used is the induced metric from the Minkowski sub-Riemannian metric on the embedding space
+        R^n+1.
+        """
+
+        def __init__(self, M):
+            """
+            Constructor.
+            """
+            self._M = M
+
+        def __str__(self):
+            return "Canonical structure"
+
+        @property
+        def typicaldist(self):
+            return jnp.sqrt(self._M.dim)
+
+        def inner(self, p, X, Y):
+            return HyperbolicSpace.minkowski_inner(X, Y)
+
+        @property
+        def metric_matrix(self):
+            """Matrix representation of the Minkowski metric"""
+            G = jnp.eye(self._M.dim + 1)
+            G = G.at[-1, -1].set(-1)
+            return G
+
+        def norm(self, p, X):
+            # inner can be smaller than 0 due to cancelling of digits when 2 similar numbers are subtracted
+            return jnp.sqrt(jax.nn.relu(self.inner(p, X, X)))
+
+        def dist(self, p, q):
+            return jnp.sqrt(self.squared_dist(p, q))
+
+        def squared_dist(self, p, q):
+            # in theory minkowski_inner(p, q) < -1 always holds, but for numerical reasons we clip
+            mink_inner_neg = jnp.clip(-HyperbolicSpace.minkowski_inner(p, q), a_min=1)
+
+            def trunc_dist_sq(x):
+                # 4th order approximation of arccosh**2 around 1
+                return 2*(x-1) - 1/3*(x-1)**2 + 4/45*(x-1)**3
+
+            d2 = jax.lax.cond(mink_inner_neg > 1+1e-6, lambda i: jnp.acosh(i)**2, trunc_dist_sq, mink_inner_neg)
+
+            return jax.nn.relu(d2)
+
+        def egrad2rgrad(self, p, H):
+            H = H.at[-1].set(-H[-1])
+            return self._M.proj(p, H)
+
+        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 exp(self, p, X):
+            # numerical safeguard
+            X = self._M.proj(p, X)
+
+            def full_exp(n2):
+                n = jnp.sqrt(n2)
+                return jnp.cosh(n) * p + jnp.sinh(n) * X/n
+            
+            def trunc_exp(n2):
+                # 6th order approximation of cosh(n)*p + sinh(n)*X/n around 0
+                return p+X + 1/6 * n2 * (3*p+X) + 1/120 * n2**2 * (5*p+X)
+
+            sqnorm_X = self.inner(p, X, X)
+            sqnorm_X = jnp.clip(sqnorm_X, a_min=0.)
+            
+            p = jax.lax.cond(sqnorm_X < 1e-6, trunc_exp, full_exp, sqnorm_X)
+            
+            return HyperbolicSpace.project_to_manifold(p)
+
+        def retr(self, p, X):
+            return self.exp(p, X)
+
+        def log(self, p, q):
+            sqd = self.squared_dist(p, q)
+            
+            def full_log(d2):
+                # For the formula see Eq. (13) in Pennec, "Hessian of the Riemannian squared distance", HAL INRIA, 2017.
+                d = jnp.sqrt(d2)
+                return d/jnp.sinh(d) * (q - jnp.cosh(d) * p)
+            
+            def trunc_log(d2):
+                # see also Pennec, "Hessian of the Riemannian squared distance", HAL INRIA, 2017.
+                return (1 - d2/6) * (q - (1 + d2/2 + d2**2/24) * p)
+            
+            v = jax.lax.cond(sqd < 1e-6, trunc_log, full_log, sqd)
+            
+            return self._M.proj(p, v)
+
+        def geopoint(self, p, q, t):
+            """Evaluate gam(t;p,q) = exp_p(t*log_p(q))."""
+            dist = self.dist(p, q)
+
+            def full_geopoint(d):
+                r = q / jnp.sinh(d) - 1/jnp.tanh(d) * p  # log_p(q) / d
+                return jnp.cosh(t*d) * p + jnp.sinh(t*d) * r
+
+            return jax.lax.cond(dist > 1e-6, full_geopoint, lambda _: (1-t)*p + t*q, dist)
+
+        def flat(self, p, X):
+            """Lower vector with the metric"""
+            dX = X.at[-1].set(-X[-1])
+            return dX
+
+        def sharp(self, p, dX):
+            """Raise covector with the metric"""
+            X = dX.at[-1].set(-dX[-1])
+            return X
+
+        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 -1 * (HyperbolicSpace.minkowski_inner(Z, Y) * X - HyperbolicSpace.minkowski_inner(Z, X) * Y)
+
+        def transp(self, p, q, X):
+            sqd = self.squared_dist(p, q)
+
+            def do_transp(X):
+                sum_l = self.log(p, q) + self.log(q, p)
+                return X - self.inner(p, self.log(p, q), X)/sqd * sum_l
+
+            return jax.lax.cond(sqd < 1e-6, lambda X: X, do_transp, X)
+
+        def jacobiField(self, p, q, t, X):
+            return NotImplementedError('This function has not been implemented yet.')
+
+        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 V: 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) / d
+
+                # 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-6,
+                                 lambda v: v / self.norm(p, v),
+                                 lambda v: jnp.zeros_like(v), U)
+
+                b2 = self.inner(p, W, U)
+
+                a1 = 1-t  # corresponds to the eigenvalue 0
+                a2 = jnp.sinh((1-t)*d) / jnp.sinh(d)  # corresponds to the eigenvalue -1
+
+                return a1 * b1 * T + a2 * b2 * U
+
+            return jax.lax.cond(dist > 1e-6, _eval, lambda args: 1/(1-t) * args[-1], (dist, w))

morphomatics/manifold/kendall.py

@@ -0,0 +1,269 @@
+################################################################################
+#                                                                              #
+#   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
+from typing import Sequence
+
+import jax
+import jax.numpy as jnp
+
+from morphomatics.manifold import ShapeSpace, Metric, Sphere
+from morphomatics.manifold.discrete_ops import pole_ladder
+
+
+class Kendall(ShapeSpace):
+    """
+    Kendall's shape space: (SO_m)-equivalence classes of preshape points (projection of centered landmarks onto the sphere)
+    """
+
+    def __init__(self, shape: Sequence[int], structure='Canonical'):
+        if len(shape) == 0:
+            raise TypeError("Need shape parameters.")
+
+        # Pre-Shape space (sphere)
+        self._S = Sphere(shape)
+        dimension = int(self._S.dim - shape[-1] * (shape[-1] - 1) / 2)
+
+        self.ref = None
+
+        name = 'Kendall shape space of ' + 'x'.join(map(str, shape[:-1])) + ' Landmarks in R^' + str(shape[-1])
+        super().__init__(name, dimension, 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 update_ref_geom(self, v):
+        self.ref = self.to_coords(v)
+
+    def to_coords(self, v):
+        '''
+        :arg v: array of landmark coordinates
+        :return: manifold coordinates
+        '''
+        return Kendall.project(v)
+
+    def from_coords(self, c):
+        '''
+        :arg c: manifold coords.
+        :returns: array of landmark coordinates
+        '''
+        return c
+
+    @property
+    def ref_coords(self):
+        """ :returns: Coordinates of reference shape """
+        return self.ref
+
+    def rand(self, key: jax.Array):
+        p = jax.random.normal(key, self.point_shape)
+        return Kendall.project(p)
+
+    def randvec(self, p, key: jax.Array):
+        v = jax.random.normal(key, self.point_shape)
+        return Kendall.horizontal(p, self._S.proj(p, v))
+
+    def zerovec(self):
+        return jnp.zeros(self.point_shape)
+
+    @staticmethod
+    def wellpos(x, y):
+        """
+        Rotate y such that it aligns to x.
+        :param x: (centered) reference landmark configuration.
+        :param y: (centered) landmarks to be aligned.
+        :returns: y well-positioned to x.
+        """
+        return y @ Kendall.opt_rot(x, y)
+
+    @staticmethod
+    def opt_rot(x, y):
+        """
+        Rotate y such that it aligns to x.
+        :param x: (centered) reference landmark configuration.
+        :param y: (centered) landmarks to be aligned.
+        :returns: rotation R such that y*R is well-positioned to x.
+        """
+        m = x.shape[-1]
+        sigma = jnp.ones(m)
+        # full_matrices=False equals full_matrices=True for quadratic input but allows for auto diff
+        u, _, v = jnp.linalg.svd(x.reshape(-1, m).T @ y.reshape(-1, m), full_matrices=False)
+        sigma = sigma.at[-1].set(jnp.sign(jnp.linalg.det(u @ v)))
+        return jnp.einsum('ji,j,kj', v, sigma, u)
+
+    @staticmethod
+    def center(x):
+        """
+        Remove mean from x.
+        """
+        mean = x.reshape(-1, x.shape[-1]).mean(axis=0)
+        return x - mean
+
+    @staticmethod
+    def project(x):
+        """
+        Project to pre-shape space.
+        : param x: Point to project.
+        :returns: Projected x.
+        """
+        x = Kendall.center(x)
+        return x / jnp.linalg.norm(x)
+
+    def proj(self, p, X):
+        """ Project a vector X from the ambient Euclidean space onto the tangent space at p. """
+        # TODO: think about naming convention.
+        return Kendall.horizontal(p, self._S.proj(p, X))
+
+    @staticmethod
+    def vertical(p, X):
+        """
+        Compute vertical component of X at base point p by solving the sylvester equation
+        App^T+pp^TA = Xp^T-pX^T for A. If p has full rank (det(pp^T) > 0), then there exists a unique solution
+        A, it is skew-symmetric and Ap is the vertical component of X
+        """
+        d = p.shape[-1]
+        S = p.reshape(-1, d).T @ p.reshape(-1, d)
+        rhs = X.reshape(-1, d).T @ p.reshape(-1, d)
+        rhs = rhs.T - rhs
+        S = jnp.kron(jnp.eye(d), S) + jnp.kron(S, jnp.eye(d))
+        A, *_ = jnp.linalg.lstsq(S, rhs.reshape(-1))
+        return jnp.einsum('...i,ij', p, A.reshape(d, d))
+
+    @staticmethod
+    def horizontal(p, X):
+        """
+        compute horizontal component of X.
+        """
+        X = Kendall.center(X)
+        return X - Kendall.vertical(p, X)
+
+    def initCanonicalStructure(self):
+        """
+        Instantiate the preshape sphere with canonical structure.
+        """
+        structure = Kendall.CanonicalStructure(self)
+        self._metric = structure
+        self._connec = structure
+
+    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
+            self._S = M._S
+
+        def __str__(self):
+            return "canonical structure"
+
+        @property
+        def typicaldist(self):
+            return np.pi/2
+
+        def inner(self, p, X, Y):
+            return self._S.metric.inner(p, X, Y)
+
+        def norm(self, p, X):
+            return self._S.metric.norm(p, X)
+
+        def flat(self, p, X):
+            """Lower vector X at p with the metric"""
+            return self._S.metric.flat(p, X)
+
+        def sharp(self, p, dX):
+            """Raise covector dX at p with the metric"""
+            return self._S.metric.sharp(p, dX)
+
+        def egrad2rgrad(self, p, X):
+            return self._M.proj(p, X)
+
+        def ehess2rhess(self, p, G, H, X):
+            """ Convert 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.
+            """
+            Y = self._S.metric.ehess2rhess(p, G, H, X)
+            return Kendall.horizontal(p, Y)
+
+        def exp(self, p, X):
+            return self._S.connec.exp(p, X)
+
+        retr = exp
+
+        def log(self, p, q):
+            q = Kendall.wellpos(p, q)
+            return self._S.connec.log(p, q)
+
+        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.
+            """
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def geopoint(self, p, q, t):
+            return self.exp(p, t * self.log(p, q))
+
+        def transp(self, p, q, X):
+            # pole ladder transports along horizontal geodesic, thus, map (p, X) to the well-positioned representative
+            R = Kendall.opt_rot(q, p)
+            return pole_ladder(self._M, p @ R, q, X @ R, 10)
+
+        def pairmean(self, p, q):
+            return self.geopoint(p, q, .5)
+
+        def dist(self, p, q):
+            q = Kendall.wellpos(p, q)
+            return self._S.metric.dist(p, q)
+
+        def squared_dist(self, p, q):
+            q = Kendall.wellpos(p, q)
+            return self._S.metric.squared_dist(p, q)
+
+        def jacobiField(self, p, q, t, X):
+            # return self.proj(*super().jacobiField(p, q, t, X))
+
+            # q = Kendall.wellpos(p, q)
+            # phi = self._S.metric.dist(p, q)
+            # v = self._S.connec.log(p, q)
+            # gamTS = self._S.connec.exp(p, t * v)
+            #
+            # v = v / self._S.metric.norm(p, v)
+            # Xtan = self._S.metric.inner(p, X, v) * v
+            # Xorth = X - Xtan
+
+            # # tangential component of J: (1-t) * transp(p, gamTS, Xtan)
+            # Jtan = Xtan_norm / phi * self._S.connec.log(gamTS, q)
+            # return gamTS, (np.sin((1 - t) * phi) / np.sin(phi)) * Kendall.horizontal(gamTS, Xorth) + Jtan
+
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def adjJacobi(self, p, q, t, X):
+            # return self.proj(p, super().adjJacobi(p, q, t, X))
+
+            raise NotImplementedError('This function has not been implemented yet.')

morphomatics/manifold/lie_group.py

@@ -0,0 +1,102 @@
+################################################################################
+#                                                                              #
+#   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                                              #
+#                                                                              #
+################################################################################
+
+# postponed evaluation of annotations to circumvent cyclic dependencies (will be default behavior in Python 4.0)
+from __future__ import annotations
+
+import abc
+
+class LieGroup(metaclass=abc.ABCMeta):
+    """
+    Interface setting out a template for Lie group classes.
+    """
+
+    def __init__(self, M: Manifold):
+        """ Construct connection.
+        :param M: underlying manifold
+        """
+        self._M = M
+
+    @abc.abstractmethod
+    def __str__(self):
+        """Returns a string representation of the particular group."""
+
+#    @abc.abstractmethod
+#    def compose(self, g, f):
+#        """Group operation"""
+
+    @property
+    @abc.abstractmethod
+    def identity(self):
+        """Returns the identity element e of the Lie group."""
+
+    @abc.abstractmethod
+    def coords(self, X):
+        """Coordinate map for the tangent space at the identity."""
+
+    @abc.abstractmethod
+    def coords_inverse(self, X):
+        """Inverse coordinate map for the tangent space at the identity."""
+
+    @abc.abstractmethod
+    def bracket(self, X, Y):
+        """Lie bracket in Lie algebra."""
+
+    @abc.abstractmethod
+    def lefttrans(self, g, f):
+        """Left translation of g by f.
+        """
+
+    @abc.abstractmethod
+    def righttrans(self, g, f):
+        """Right translation of g by f.
+        """
+
+    @abc.abstractmethod
+    def inverse(self, g):
+        """Inverse map of the Lie group.
+        """
+
+    @abc.abstractmethod
+    def exp(self, X):
+        """Computes the Lie-theoretic exponential map of a tangent vector X at e.
+        """
+
+    @abc.abstractmethod
+    def log(self, g):
+        """Computes the Lie-theoretic logarithm of g. This is the inverse of `exp`.
+        """
+
+    @abc.abstractmethod
+    def dleft(self, f, X):
+        """Derivative of the left translation by f at e applied to the tangent vector X.
+        """
+
+    @abc.abstractmethod
+    def dright(self, f, X):
+        """Derivative of the right translation by f at e applied to the tangent vector X.
+        """
+
+    @abc.abstractmethod
+    def dleft_inv(self, f, X):
+        """Derivative of the left translation by f^{-1} at f applied to the tangent vector X.
+        """
+
+    @abc.abstractmethod
+    def dright_inv(self, f, X):
+        """Derivative of the right translation by f^{-1} at f applied to the tangent vector X.
+        """
+
+    @abc.abstractmethod
+    def adjrep(self, g, X):
+        """Adjoint representation of g applied to the tangent vector X at the identity.
+        """