morphomatics 4.0__py3-none-any.whl

morphomatics/manifold/gl_p_coords.py

@@ -0,0 +1,149 @@
+################################################################################
+#                                                                              #
+#   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.random
+import numpy as np
+import jax.numpy as jnp
+
+from scipy import sparse
+
+try:
+    from sksparse.cholmod import cholesky as direct_solve
+except:
+    from scipy.sparse.linalg import factorized as direct_solve
+
+from ..geom import Surface
+from . import GLpn
+from . import ShapeSpace
+from .util import align
+
+
+class GLpCoords(ShapeSpace):
+    """
+    Shape space based the group of matrices with positive determinant.
+
+    See:
+    Felix Ambellan, Stefan Zachow and Christoph von Tycowicz.
+    An as-invariant-as-possible GL+(3)-based statistical shape model.
+    Proc. 7th MICCAI workshop on Mathematical Foundations of Computational Anatomy, pp. 219--228, 2019.
+    """
+
+    def __init__(self, reference: Surface):
+        """
+        :arg reference: Reference surface (shapes will be encoded as deformations thereof)
+        """
+        assert reference is not None
+        self.ref = reference
+        k = len(self.ref.f)
+
+        self.update_ref_geom(self.ref.v)
+
+        self.GLp = GLpn(k, structure='AffineGroup')
+
+        name = 'Shape Space based on the orientation preserving component of the general linear group'
+        super().__init__(name, self.GLp.dim, self.GLp.point_shape, self.GLp.connec, None, self.GLp.group)
+
+    def tree_flatten(self):
+        return tuple(), (self.ref.v.tolist(), self.ref.f.tolist())
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Specifies an unflattening recipe for PyTree registration."""
+        return cls(Surface(*aux_data))
+
+    @property
+    def M(self):
+        return self.GLp
+
+    @property
+    def n_triangles(self):
+        return len(self.ref.f)
+
+    def update_ref_geom(self, v):
+        self.ref.v = v
+
+        # center of gravity
+        self.CoG = self.ref.v.mean(axis=0)
+
+        # setup Poisson system
+        S = self.ref.div @ self.ref.grad
+        # add soft-constraint fixing translational DoF
+        S += sparse.coo_matrix(([1.0], ([0], [0])), S.shape)  # make pos-def
+        self.poisson = direct_solve(S.tocsc())
+
+    def to_coords(self, v):
+        """
+        :arg v: #v-by-3 array of vertex coordinates
+        :return: GLp coords.
+        """
+
+        # align
+        v = align(v, self.ref.v)
+
+        # compute gradients
+        D = self.ref.grad @ v
+
+        # D holds transpose of def. grads.
+        # decompose...
+        U, S, Vt = np.linalg.svd(D.reshape(-1, 3, 3))
+
+        # ...rotation
+        R = np.einsum('...ij,...jk', U, Vt)
+        W = np.ones_like(S)
+        W[:, -1] = np.linalg.det(R)
+        R = np.einsum('...ij,...j,...jk', U, W, Vt)
+
+        # ...stretch
+        S[:, -1] = 1  # no stretch (=1) in normal direction
+        # for degenerate triangles
+        # TODO: check which direction is normal in degenerate case
+        S[S < 1e-6] = 1e-6
+        U = np.einsum('...ij,...j,...kj', U, S, U)
+        return jnp.einsum('...ij,...jl', R, U)
+
+    def from_coords(self, D):
+        """
+        :arg D: GLp coords.
+        :returns: #v-by-3 array of vertex coordinates
+        """
+        # solve Poisson system
+        rhs = self.ref.div @ D.reshape(-1, 3)
+        v = self.poisson(rhs)
+        # move to CoG
+        v += self.CoG - v.mean(axis=0)
+
+        return v
+
+    @property
+    def ref_coords(self):
+        """ Identity coordinates (i.e., the reference shape).
+        """
+        return self.group.identity
+
+    def rand(self, key: jax.Array):
+        """Random set of coordinates () won't represent a 'nice' shape).
+        """
+        return self.GLp.rand(key)
+
+    def randvec(self, A, key: jax.Array):
+        return self.GLp.randvec(A, key)
+
+    def zerovec(self):
+        """Zero tangent vector in any tangent space.
+        """
+        return self.GLp.zerovec()
+
+    def proj(self, p, X):
+        return X
+
+    def geopoint(self, A, B, t):
+        return self.GLp.connec.geopoint(A, B, t)

morphomatics/manifold/gl_p_n.py

@@ -0,0 +1,201 @@
+################################################################################
+#                                                                              #
+#   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 jax.scipy.linalg import expm, funm
+
+from morphomatics.manifold import Manifold, LieGroup, Connection
+
+
+class GLpn(Manifold):
+    """Returns the Lie group GL^+(n), i.e., the group of n-by-n matrices each with positive determinant.
+
+     manifold = GLpn(n)
+
+     Elements of GL^+(n) are represented as arrays of size nxn.
+
+     # NOTE: Tangent vectors are represented as left translations in the Lie algebra, i.e., a tangent vector X at g is
+     represented as d_gL_{g^(-1)}(X)
+     """
+
+    def __init__(self, n=3, structure='AffineGroup'):
+        self._n = n
+
+        name = 'Orientation preserving maps of R^n'
+
+        super().__init__(name, n**2, point_shape=(n, n))
+
+        if structure:
+            getattr(self, f'init{structure}Structure')()
+
+    def tree_flatten(self):
+        children, aux = super().tree_flatten()
+        return children, aux+(self._n,)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Specifies an unflattening recipe for PyTree registration."""
+        *aux_data, n = aux_data
+        obj = cls(n, structure=None)
+        obj.tree_unflatten_instance(aux_data, children)
+        return obj
+
+    @property
+    def n(self):
+        return self._n
+
+    def rand(self, key: jax.Array):
+        """Returns a random point in the Lie group. This does not
+        follow a specific distribution."""
+        A = jax.random.normal(key, self.point_shape)
+        return expm(A)
+
+    def randvec(self, A, key: jax.Array):
+        """Returns a random vector in the tangent space at A.
+        """
+        return jax.random.normal(key, self.point_shape)
+
+    def zerovec(self):
+        """Returns the zero vector in the tangent space at g."""
+        return jnp.zeros((self.n, self.n))
+
+    def proj(self, p, X):
+        return X
+
+    def initAffineGroupStructure(self):
+        """
+        Standard group structure with canonical Cartan Shouten connection.
+        """
+        structure = GLpn.AffineGroupStructure(self)
+        self._connec = structure
+        self._group = structure
+
+    class AffineGroupStructure(Connection, LieGroup):
+        """
+        Standard group structure on GL+(n) where the composition of two elements is given by component-wise matrix
+        multiplication. The connection is the corresponding canonical Cartan Shouten (CCS) connection. No Riemannian
+        metric is used.
+        """
+
+        def __init__(self, M):
+            """
+            Constructor.
+            """
+            self._M = M
+
+        def __str__(self):
+            return 'standard group structure on GL+(n) with CCS connection'
+
+        # Group
+
+        @property
+        def identity(self):
+            """Returns the identity element e of the Lie group."""
+            return jnp.eye(self._M.n)
+
+        def bracket(self, X, Y):
+            """Lie bracket in Lie algebra."""
+            return jnp.einsum('ij,jl->il', X, Y) - jnp.einsum('ij,jl->il', Y, X)
+
+        def coords(self, X):
+            """Coordinate map for the tangent space at the identity."""
+            return jnp.reshape(X, (self._M.n ** 2, 1))
+
+        def coords_inverse(self, X):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def lefttrans(self, g, f):
+            """Left translation of g by f.
+            """
+            return jnp.einsum('ij,jl->il', f, g)
+
+        def righttrans(self, g, f):
+            """Right translation of g by f.
+            """
+            return jnp.einsum('ij,jl->il', g, f)
+
+        def inverse(self, g):
+            """Inverse map of the Lie group.
+            """
+            return jnp.linalg.inv(g)
+
+        def exp(self, *argv):
+            """Computes the Lie-theoretic and connection exponential map
+            (depending on signature, i.e. whether footpoint is given as well)
+            """
+            return jax.lax.cond(len(argv) == 1,
+                                lambda A: A[-1],  # group exp
+                                lambda A: jnp.einsum('ij,jk', A[-1], A[0]),  # exp of CCS connection
+                                (argv[0], expm(argv[-1])))
+
+        retr = exp
+
+        def log(self, *argv):
+            """Computes the Lie-theoretic and connection logarithm map
+            (depending on signature, i.e. whether footpoint is given as well)
+            """
+            # NOTE: as logm() is not available in jax we apply log via funm() (so far this is CPU only; not as stable as
+            # logm in numpy)
+            logm = lambda m: jnp.real(funm(m, jnp.log))
+            return logm(jax.lax.cond(len(argv) == 1,
+                                     lambda A: A[-1],
+                                     lambda A: jnp.einsum('ij,kj', A[-1], A[0]),
+                                     argv))
+
+        def curvature_tensor(self, f, X, Y, Z):
+            """Evaluates the curvature tensor R of the connection at f 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 / 4 * self.bracket(self.bracket(X, Y), Z)
+
+        def dleft(self, f, X):
+            """Derivative of the left translation by f applied to the tangent vector X at the identity.
+            """
+            return jnp.einsum('ij,jl->il', f, X)
+
+        def dright(self, f, X):
+            """Derivative of the right translation by f at g applied to the tangent vector X.
+            """
+            return jnp.einsum('ij,jl->il', X, f)
+
+        def dleft_inv(self, f, X):
+            """Derivative of the left translation by f^{-1} at f applied to the tangent vector X.
+            """
+            return jnp.einsum('ij,jl->il', self.inverse(f), X)
+
+        def dright_inv(self, f, X):
+            """Derivative of the right translation by f^{-1} at f applied to the tangent vector X.
+            """
+            return jnp.einsum('ij,jl->il', X, self.inverse(f))
+
+        def adjrep(self, g, X):
+            """Adjoint representation of g applied to the tangent vector X at the identity.
+            """
+            return jnp.einsum('ij,jl,lm->im', g, X, self.inverse(g))
+
+        def transp(self, f, g, X):
+            """
+            Parallel transport of the CCS connection along one-parameter subgroups; see Sec. 5.3.3 of
+            X. Pennec and M. Lorenzi,
+            "Beyond Riemannian geometry: The affine connection setting for transformation groups."
+
+            """
+            f_invg = self.lefttrans(g, self.inverse(f))
+            h = self.geopoint(self.identity, f_invg, .5)
+
+            return self.dleft_inv(f_invg, self.dleft(h, self.dright(h, X)))
+
+        def jacobiField(self, R, Q, t, X):
+            raise NotImplementedError('This function has not been implemented yet.')

morphomatics/manifold/grassmann.py

@@ -0,0 +1,174 @@
+################################################################################
+#                                                                              #
+#   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                                              #
+#                                                                              #
+################################################################################
+
+from morphomatics.manifold import Manifold, Metric
+from morphomatics.manifold.discrete_ops import pole_ladder
+
+import jax
+import jax.numpy as jnp
+from jax.numpy.linalg import svd
+
+
+class Grassmann(Manifold):
+    """ The Grassmannian.
+    Manifold of m-dimensional subspaces of n dimensional real vector space
+    Elements are represented as n x m matrices.
+    """
+    def __init__(self, n=3, m=1, structure='Canonical'):
+        if n < m or m < 1:
+            raise ValueError(
+                "Need n >= p >= 1. Values supplied were n = {n} and m = {m}"
+            )
+
+        name = 'Grassmann manifold Gr({n},{m})'.format(n=n, m=m)
+        dimension = int(n * m - m ** 2)
+        self._n = n
+        self._m = m
+        super().__init__(name, dimension, point_shape=(n, m))
+        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 Grassmannian with canonical structure.
+        """
+        structure = Grassmann.CanonicalStructure(self)
+        self._metric = structure
+        self._connec = structure
+
+    # Generate random Grassmann point using qr of random normally distributed matrix.
+    def rand(self, key: jax.Array):
+        Q, _ = jnp.linalg.qr(jax.random.normal(key, self.point_shape))
+        return Q
+
+    def randvec(self, p, key: jax.Array):
+        U = jax.random.normal(key, p.shape)
+        U = U - jnp.einsum('ij,kj,kl', p, p, U)
+        U = U / jnp.linalg.norm(U)
+        return U
+
+    def zerovec(self):
+        return jnp.zeros(self.point_shape)
+
+    @staticmethod
+    def project(p):
+        """Project arbitrary matrix to the manifold."""
+        return jnp.linalg.qr(p)[0]
+
+    def proj(self, p, U):
+        """Project ambient tangent vector to tangent space at p."""
+        return U - jnp.einsum('ij,kj,kl', p, p, U)
+
+    class CanonicalStructure(Metric):
+        """
+        The Riemannian metric used is the induced metric from the embedding space (R^nxm)^k, i.e., this manifold is a
+        Riemannian submanifold of (R^nxm)^k endowed with the usual trace inner product.
+        """
+        def __init__(self, M):
+            """
+            Constructor.
+            """
+            self._M = M
+
+        def __str__(self):
+            return "Canonical structure of Grassmannian"
+
+        @property
+        def typicaldist(self):
+            return jnp.sqrt(jnp.prod(self._M.point_shape[-2:]))
+
+        # Geodesic distance for Grassmann
+        def dist(self, p, q):
+            s = svd(jnp.einsum('ji,jk', p, q), compute_uv=False)
+            s = jnp.arccos(jnp.min(s))
+            return jnp.linalg.norm(s)
+
+        def inner(self, p, G, H):
+            # Inner product (Riemannian metric) on the tangent space
+            # For the Grassmann this is the Frobenius inner product.
+            return jnp.einsum('ij,ij', G, H)
+
+        def flat(self, p, G):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def sharp(self, p, dG):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def egrad2rgrad(self, p, X):
+            return self._M.proj(p, X)
+
+        def ehess2rhess(self, p, G, H, X):
+            # Convert Euclidean into Riemannian Hessian.
+            xpG = jnp.einsum('ij,kj,kl', X, p, G)
+            return self._M.proj(p, H) - xpG
+
+        def retr(self, X, G):
+            # We do not need to worry about flipping signs of columns here,
+            # since only the column space is important, not the actual
+            # columns. Compare this with the Stiefel manifold.
+
+            # Compute the polar factorization of Y = X+P_G
+            u, _, vt = svd(X + G, full_matrices=False)
+            return u @ vt
+
+        def norm(self, p, G):
+            # Norm on the tangent space is simply the Euclidean norm.
+            return jnp.linalg.norm(G)
+
+        def transp(self, p1, p2, d):
+
+            return pole_ladder(self._M, p1, p2, d)
+
+        def exp(self, p, U):
+            u, s, vt = svd(U, full_matrices=False)
+
+            Y = jnp.einsum('ij,kj,k,kl', p, vt, jnp.cos(s), vt) + \
+                jnp.einsum('ij,j,jk', u, jnp.sin(s), vt)
+
+            # From numerical experiments, it seems necessary to
+            # re-orthonormalize. This is overall quite expensive.
+            Y, _ = jnp.linalg.qr(Y)
+            return Y
+
+        def log(self, p, q):
+            qtp = q.T @ p
+            At = q.T - qtp @ p.T
+            Bt = jnp.linalg.solve(qtp, At)
+            u, s, vt = svd(Bt.T, full_matrices=False)
+
+            return jnp.einsum('ij,j,jk', u, jnp.arctan(s), vt)
+
+        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 jacobiField(self, p, q, t, X):
+            raise NotImplementedError('This function has not been implemented yet.')
+
+        def adjJacobi(self, p, q, t, X):
+            raise NotImplementedError('This function has not been implemented yet.')