morphomatics 4.0__py3-none-any.whl

morphomatics/manifold/manifold.py

@@ -0,0 +1,162 @@
+################################################################################
+#                                                                              #
+#   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 abc
+import jax
+from jax.tree_util import register_pytree_node
+
+import operator as op
+
+from morphomatics.manifold import Metric, Connection, LieGroup
+
+class ManifoldMeta(abc.ABCMeta):
+    """Metaclass for abstract base class for which subclasses are to be registered as jax PyTree"""
+    def __new__(mcls, name, bases, namespace, **kwargs):
+        cls = super().__new__(mcls, name, bases, namespace, **kwargs)
+        register_pytree_node(cls, op.methodcaller('tree_flatten'), cls.tree_unflatten)
+        return cls
+
+class Manifold(metaclass=ManifoldMeta):
+    """
+    Abstract base class setting out a template for manifold classes.
+    Morphomatics's Lie group and Riemannian manifold classes inherit from Manifold.
+    """
+
+    def __init__(self, name, dimension: int, point_shape,
+                 connec: Connection = None, metric: Metric = None, group: LieGroup = None):
+        self._name = name
+        self._dimension = dimension
+        self._point_shape = point_shape
+        # (possibly) define a connection on the tangent bundle
+        self._connec = connec
+        # (possibly) define a metric on the tangent bundle
+        self._metric = metric
+        # (possibly) define a group operation turning the manifold into a Lie group
+        self._group = group
+
+    @classmethod
+    @abc.abstractmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Specifies an unflattening recipe for PyTree registration."""
+
+    def tree_unflatten_instance(self, aux_data, children):
+        """Specifies an unflattening recipe given an instance (possibly of a subclass)."""
+        C, M, G = aux_data
+        Cdict, Mdict, Gdict = children
+
+        def setup(cls, dict):
+            if cls is None.__class__:
+                return None
+            obj = cls(self)
+            obj.__dict__.update(dict)
+            return obj
+
+        self._connec = setup(C, Cdict)
+        self._metric = setup(M, Mdict)
+        self._group = setup(G, Gdict)
+
+    def tree_flatten(self):
+        """Specifies a flattening recipe for PyTree registration."""
+        aux_data = (self._connec.__class__, self._metric.__class__, self._group.__class__)
+        wo_ = lambda o: {} if o is None else {k: v for k, v in o.__dict__.items() if k[0] != '_'}
+        children = (wo_(self._connec), wo_(self._metric), wo_(self._group))
+        return (children, aux_data)
+
+    def __str__(self):
+        return self._name
+
+    def __repr__(self):
+        """Returns a string representation of the particular manifold."""
+        conf = 'metric='+str(self._metric) if self._metric else ''
+        conf += ' connection='+str(self._connec) if self._connec else ''
+        conf += ' group='+str(self._group) if self._group else ''
+        if not conf:
+            return self._name
+        return f'{self._name} ({conf.strip()})'
+
+    @property
+    def dim(self):
+        """The dimension of the manifold"""
+        return self._dimension
+
+    @property
+    def point_shape(self):
+        """Dimensions of elements of the manifold.
+
+        Tuple of dimension, e.g., if an element is given by a 3-by-3 matrix, then its point shape is [3, 3].
+        """
+        return self._point_shape
+
+    @property
+    def metric(self) -> Metric:
+        return self._metric
+
+    @property
+    def connec(self) -> Connection:
+        return self._connec
+
+    @property
+    def group(self) -> LieGroup:
+        return self._group
+
+    @abc.abstractmethod
+    def rand(self, key: jax.Array):
+        """Returns a random point of the manifold."""
+
+    @abc.abstractmethod
+    def randvec(self, p, key: jax.Array):
+        """Returns a random vector in the tangent space at p."""
+
+    @abc.abstractmethod
+    def zerovec(self):
+        """Returns the zero vector in any tangent space."""
+
+    @abc.abstractmethod
+    def proj(self, p, X):
+        """Projects a vector X in the ambient space on the tangent space at
+        p.
+        """
+
+    def projToGeodesic(self, X, Y, P, max_iter=10):
+        '''
+        Project P onto geodesic from X to Y.
+
+        See:
+        Felix Ambellan, Stefan Zachow, Christoph von Tycowicz.
+        Geodesic B-Score for Improved Assessment of Knee Osteoarthritis.
+        Proc. Information Processing in Medical Imaging (IPMI), LNCS, 2021.
+
+        :arg X, Y: manifold coords defining geodesic X->Y.
+        :arg P: manifold coords to be projected to X->Y.
+        :returns: manifold coords of projection of P to X->Y
+        '''
+        assert X.shape == Y.shape
+        assert Y.shape == P.shape
+        assert self.connec
+        assert self.metric
+
+        # initial guess
+        Pi = X.copy()
+
+        # solver loop
+        for _ in range(max_iter):
+            v = self.connec.log(Pi, Y)
+            v = v / self.metric.norm(Pi, v)
+            w = self.connec.log(Pi, P)
+            d = self.metric.inner(Pi, v, w)
+
+            # print(f'|<v, w>|={d}')
+            if abs(d) < 1e-6: break
+
+            Pi = self.connec.exp(Pi, d * v)
+
+        return Pi

morphomatics/manifold/manopt_wrapper.py

@@ -0,0 +1,185 @@
+################################################################################
+#                                                                              #
+#   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 functools
+
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+import jax.tree_util as tree
+
+from .manifold import Manifold
+
+try:
+    import pymanopt
+
+    from pymanopt.manifolds.manifold import Manifold as ManoptManifold
+    from pymanopt.autodiff.backends._backend import Backend
+    from pymanopt.autodiff import backend_decorator_factory
+
+
+    def conjugate_result(function):
+        @functools.wraps(function)
+        def wrapper(*args, **kwargs):
+            return tree.tree_map(jnp.conj, function(*args, **kwargs))
+
+        return wrapper
+
+    class JaxBackend(Backend):
+        def __init__(self):
+            super().__init__("Jax")
+
+        @staticmethod
+        def is_available():
+            return jax is not None
+
+        @Backend._assert_backend_available
+        def prepare_function(self, function):
+            return function
+
+        @Backend._assert_backend_available
+        def generate_gradient_operator(self, function, num_arguments):
+            gradient = conjugate_result(
+                jax.grad(function) if num_arguments == 1 else
+                jax.grad(function, argnums=list(range(num_arguments)))
+            )
+            return gradient
+
+        @staticmethod
+        def _hessian_vector_product(function, argnum):
+            raise NotImplementedError()
+
+        @Backend._assert_backend_available
+        def generate_hessian_operator(self, function, num_arguments):
+            raise NotImplementedError()
+
+
+    if "jax" not in pymanopt.autodiff.backends.__all__:
+        factory = backend_decorator_factory(JaxBackend)
+        pymanopt.autodiff.backends.jax = factory
+        pymanopt.function.jax = factory
+        pymanopt.autodiff.backends.__all__.append("jax")
+        pymanopt.function.__all__.append("jax")
+
+except ImportError:
+    _has_manopt = False
+    ManoptManifold = object
+else:
+    _has_manopt = True
+
+
+class ManoptWrap(ManoptManifold):
+    """
+    Wraper for pymanopt to make manifolds from morphomatics compatible.
+    """
+
+    def __init__(self, M: Manifold):
+        self._M = M
+        super().__init__(str(M), M.dim) # as of v0.2.6rc1
+
+    # Manifold properties that subclasses can define
+
+    @property
+    def typicaldist(self):
+        """Returns the "scale" of the manifold. This is used by the
+        trust-regions solver to determine default initial and maximal
+        trust-region radii.
+        """
+        return self._M.metric.typicaldist
+
+    def inner_product(self, X, G, H):
+        """Returns the inner product (i.e., the Riemannian metric) between two
+        tangent vectors `P_G` and `H` in the tangent space at `X`.
+        """
+        return self._M.metric.inner(X, G, H)
+
+    def dist(self, X, Y):
+        """
+        Geodesic distance on the manifold
+        """
+        return self._M.metric.dist(X, Y)
+
+    def projection(self, X, G):
+        """Projects a vector `P_G` in the ambient space on the tangent space at
+        `X`.
+        """
+        return self._M.metric.proj(X, G)
+
+    def euclidean_to_riemannian_gradient(self, X, G):
+        """Maps the Euclidean gradient P_G in the ambient space on the tangent
+        space of the manifold at X.
+        """
+        return self._M.metric.egrad2rgrad(X, G)
+
+    def norm(self, X, G):
+        """Computes the norm of a tangent vector `P_G` in the tangent space at
+        `X`.
+        """
+        return self._M.metric.norm(X, G)
+
+    def exp(self, X, U):
+        """
+        The exponential (in the sense of Lie group theory) of a tangent
+        vector U at X.
+        """
+        return self._M.connec.exp(X, U)
+
+    def retraction(self, X, G):
+        """
+        A retraction mapping from the tangent space at X to the manifold.
+        See Absil for definition of retraction.
+        """
+        return self.exp(X, G)
+
+    def log(self, X, Y):
+        """
+        The logarithm (in the sense of Lie group theory) of Y. This is the
+        inverse of exp.
+        """
+        return self._M.connec.log(X, Y)
+
+    def transport(self, x1, x2, d):
+        """
+        Transports d, which is a tangent vector at x1, into the tangent
+        space at x2.
+        """
+        return self._M.connec.transp(x1, x2, d)
+
+    def random_point(self):
+        """Returns a random point on the manifold."""
+        key = jax.random.PRNGKey(np.random.randint(1 << 32))
+        return self._M.rand(key)
+
+    def random_tangent_vector(self, X):
+        """Returns a random vector in the tangent space at `X`. This does not
+        follow a specific distribution.
+        """
+        key = jax.random.PRNGKey(np.random.randint(1 << 32))
+        return self._M.randvec(X, key)
+
+    def zero_vector(self, X):
+        """Returns the zero vector in the tangent space at X."""
+        return self._M.zerovec(X)
+
+    def euclidean_to_riemannian_hessian(self, p, grad, Hess, X):
+        """
+        Convert Euclidean into Riemannian Hessian.
+        """
+        return
+
+    def pair_mean(self, X, Y):
+        """
+        Computes the intrinsic mean of X and Y, that is, a point that lies
+        mid-way between X and Y on the geodesic arc joining them.
+        """
+        return self.exp(X, 0.5 * self.log(X, Y))

morphomatics/manifold/metric.py

@@ -0,0 +1,110 @@
+################################################################################
+#                                                                              #
+#   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 abc
+
+import jax
+import jax.numpy as jnp
+
+from morphomatics.manifold.connection import Connection
+
+
+class Metric(Connection):
+    """
+    Interface setting out a template for a metric on the tangent bundle of a manifold. It is modelled as a subclass of
+    its Levi-Civita connection.
+    """
+
+    @property
+    @abc.abstractmethod
+    def typicaldist(self):
+        """Returns the "scale" of the manifold. This is used by the
+        trust-regions solver to determine default initial and maximal
+        trust-region radii.
+        """
+
+    @abc.abstractmethod
+    def dist(self, p, q):
+        """Returns the geodesic distance between two points p and q on the
+        manifold."""
+
+    def squared_dist(self, p, q):
+        return self.dist(p, q)**2
+
+    @abc.abstractmethod
+    def inner(self, p, X, Y):
+        """Returns the inner product (i.e., the Riemannian metric) between two
+        tangent vectors X and Y from the tangent space at p.
+        """
+
+    @abc.abstractmethod
+    def egrad2rgrad(self, p, X):
+        """Maps the Euclidean gradient X in the ambient space on the tangent
+        space of the manifold at p.
+        """
+
+    @abc.abstractmethod
+    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.
+        """
+
+    def norm(self, p, X):
+        """Computes the norm of a tangent vector X in the tangent space at
+        p.
+        """
+        return jnp.sqrt(self.inner(p, X, X))
+
+    @abc.abstractmethod
+    def flat(self, p, X):
+        """Lower vector X at p with the metric"""
+
+    @abc.abstractmethod
+    def sharp(self, p, dX):
+        """Raise covector dX at p with the metric"""
+
+    @abc.abstractmethod
+    def adjJacobi(self, p, q, t, X):
+        """Evaluates an adjoint Jacobi field for the geodesic gam from p to q at p.
+        :param p: element of the Riemannian manifold
+        :param q: element of the Riemannian manifold
+        :param t: scalar in [0,1]
+        :param X: tangent vector at gam(t)
+        :return: tangent vector at p
+        """
+
+    def adjDxgeo(self, p, q, t, X):
+        """Evaluates the adjoint of the differential of the geodesic gamma from p to q w.r.t. the starting point p at X,
+        i.e, the adjoint  of d_p gamma(t; ., q) applied to X, which is an element of the tangent space at p.
+        """
+
+        return self.adjJacobi(p, q, t, X)
+
+    def adjDygeo(self, p, q, t, X):
+        """Evaluates the adjoint of the differential of the geodesic gamma from p to q w.r.t. the endpoint q at X,
+        i.e, the adjoint  of d_q gamma(t; p, .) applied to X, which is en element of the tangent space at q.
+        """
+
+        return self.adjJacobi(q, p, 1 - t, X)
+
+
+def _eval_adjJacobi_embed(g: Metric, p, q, t, X):
+    """ Implementation of eval_adjJacobi for isometrically embedded manifolds using (forward-mode) automatic
+    differentiation of geopoint(..).
+
+    ATTENTION: the result must be projected to the tangent space!
+    """
+    f = lambda O: g.geopoint(O, q, t)
+    gam, Jt = jax.vjp(f, p)
+    co_X = g.flat(gam, X)
+    return g.sharp(p, Jt(co_X)[0])

morphomatics/manifold/point_distribution_model.py

@@ -0,0 +1,143 @@
+################################################################################
+#                                                                              #
+#   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
+import numpy as np
+
+from ..geom import Surface
+from . import ShapeSpace, Metric
+from .util import align
+
+
+class PointDistributionModel(ShapeSpace, Metric):
+    """ Linear manifold space model. """
+
+    def __init__(self, reference: Surface):
+        """
+        :arg reference: Reference surface (shapes will be encoded as deformations thereof)
+        """
+        assert reference is not None
+        self.ref = reference
+
+        name = 'Point Distribution Model'
+        dimension = reference.v.size
+        point_shape = reference.v.shape
+        super().__init__(name, dimension, point_shape, self, self, None)
+
+    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 typicaldist(self):
+        return self.ref.v.std()
+
+    def update_ref_geom(self, v):
+        self.ref.v=v
+
+    def to_coords(self, v):
+        # align
+        return align(v, self.ref.v)
+
+    def from_coords(self, c):
+        return np.asarray(c)
+
+    @property
+    def ref_coords(self):
+        return self.ref.v
+
+    def dist(self, p, q):
+        return jnp.sqrt(self.squared_dist(p, q))
+
+    def squared_dist(self, p, q):
+        return jnp.sum((p-q)**2)
+
+    def inner(self, p, X, Y):
+        """
+        :arg X: (list of) tangent vector(s) at p
+        :arg Y: (list of) tangent vector(s) at p
+        :returns: inner product at p between P_G and H, i.e. <X,Y>_p
+        """
+        return X.reshape(-1) @ Y.reshape(-1)
+
+    def flat(self, p, X):
+        return X
+
+    def sharp(self, p, dX):
+        return dX
+
+    def proj(self, p, X):
+        return X
+
+    def egrad2rgrad(self, p, X):
+        return X
+
+    def ehess2rhess(self, p, egrad, ehess, X):
+        return ehess
+
+    def exp(self, p, X):
+        return p + X
+
+    retr = exp
+
+    def log(self, p, q):
+        return q - p
+
+    def curvature_tensor(self, p, X, Y, Z):
+        return self.zerovec()
+
+    def rand(self, key: jax.Array):
+        v = jax.random.normal(key, self.ref.v.shape)
+        return self.to_coords(v)
+
+    def zerovec(self):
+        jnp.zeros(self.ref.v.shape)
+
+    def transp(self, p, q, X):
+        return X
+
+    def jacobiField(self, p, q, t, X):
+        """Evaluates a Jacobi field (with boundary conditions gam(0) = X, gam(1) = 0) along the geodesic gam from p to q.
+        :param p: point
+        :param q: point
+        :param t: scalar in [0,1]
+        :param X: tangent vector at p
+        :return: tangent vector at gam(t)
+        """
+        return (1-t) * X
+
+    def adjJacobi(self, p, q, t, X):
+        """
+        :param p: point
+        :param q: point
+        :param t: scalar in [0, 1]
+        :param X: vector at gam(p,q,t)
+        :return: vector at p
+        """
+        return X / (1.0 - t)
+
+    def projToGeodesic(self, p, q, m):
+        '''
+        :arg X, Y: manifold coords defining geodesic X->Y.
+        :arg P: manifold coords to be projected to X->Y.
+        :returns: manifold coords of projection of P to X->Y
+        '''
+        return super().projToGeodesic(p, q, m, max_iter=1)
+
+    def coords(self, X):
+        """Coordinate map of the tangent space at the identity"""
+        return X.reshape(-1)