morphomatics 4.0__py3-none-any.whl

morphomatics/manifold/util.py

@@ -0,0 +1,126 @@
+################################################################################
+#                                                                              #
+#   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 morphomatics.geom.surface import Surface
+
+
+def align(src, ref):
+    """ (Constrained) Procrustes alignment of src to ref using Kabsch algorithm
+    :arg src: n-by-3 array of vertex coordinates of source object
+    :arg ref: n-by-3 array of vertex coordinates of reference
+    :returns: aligned coords.
+    """
+    n = len(ref)
+    # cross-covariance matrix
+    c_s = src.mean(axis=0)
+    c_r = ref.mean(axis=0)
+    xCov = (ref.T @ src) / n - jnp.outer(c_r, c_s)
+    # optimal rotation
+    U, S, Vt = jnp.linalg.svd(xCov)
+    R = Vt.T @ U.T
+    if jnp.linalg.det(R) < 0:
+        R = Vt.T @ np.diag([1, 1, -1]) @ U.T
+    # return aligned coords.
+    return src @ R + (c_r - c_s @ R)
+
+
+def generalized_procrustes(surf):
+    """ Generalized Procrustes analysis.
+    :arg surf: list of surfaces to be aligned. The meshes must be in correspondence.
+    """
+    ref = Surface(jnp.copy(surf[0].v), jnp.copy(surf[0].f))
+    old_ref = Surface(jnp.copy(ref.v), jnp.copy(ref.f))
+
+    n_steps = 0
+    # do until convergence
+    while (jnp.linalg.norm(ref.v - old_ref.v) > 1e-11 and 1000 > n_steps) or n_steps == 0:
+        n_steps = n_steps + 1
+        old_ref = Surface(jnp.copy(ref.v), jnp.copy(ref.f))
+        # align meshes to reference
+        for i, s in enumerate(surf):
+            s.v = align(s.v, ref.v)
+
+        # compute new reference
+        v_ref = jnp.mean(jnp.array([s.v for s in surf]), axis=0)
+        ref.v = v_ref
+
+
+def preshape(v):
+    """ Center point cloud at origin and normalize its size
+    :arg v: n-by-3 array of vertex coordinates
+    :returns: n-by-3 array of adjusted vertex coordinates
+    """
+    # center
+    v = v - 1 / v.shape[0] * jnp.tile(jnp.sum(v, axis=0), (v.shape[0], 1))
+    # normalize
+    v /= jnp.linalg.norm(v)
+    return v
+
+
+def multiprod(A: jnp.ndarray, B: jnp.ndarray) -> jnp.ndarray:
+    # vectorized matrix - matrix multiplication
+    return jnp.einsum('...ij,...jk', A, B)
+
+
+def multitransp(A):
+    # vectorized matrix transpose
+    return jnp.einsum('...ij->...ji', A)
+
+
+def multiskew(A):
+    return 0.5 * (A - jnp.einsum('...ij->...ji', A))
+
+
+def multisym(A):
+    return 0.5 * (A + jnp.einsum('...ij->...ji', A))
+
+
+def vectime3d(x, A):
+    """
+    :param x: vector of length k
+    :param A: array of size k x n x m
+    :return: k x n x m array such that the j-th n x m slice of A is multiplied with the j-th element of x
+
+    In case of k=1, x * A is returned.
+    """
+    if jnp.isscalar(x) and A.ndim == 2:
+        return x * A
+
+    x = jnp.atleast_2d(x)
+    assert x.ndim <= 2 and jnp.size(A.shape) == 3
+    assert x.shape[0] == 1 or x.shape[1] == 1
+    assert x.shape[0] == A.shape[0] or x.shape[1] == A.shape[0]
+
+    if x.shape[1] == 1:
+        x = x.T
+
+    A = jnp.einsum('kij->ijk', A)
+    return jnp.einsum('ijk->kij', x * A)
+
+
+def gram_schmidt(A):
+    """Orthogonalize a set of vectors stored as the columns of matrix A."""
+    # Get the number of vectors.
+    n = A.shape[1]
+    for j in range(n):
+        # To orthogonalize the vector in column j with respect to the
+        # previous vectors, subtract from it its projection onto
+        # each of the previous vectors.
+        for k in range(j):
+            A = A.at[:, j].set(A[:, j] - jnp.dot(A[:, k], A[:, j]) * A[:, k])
+        A = A.at[:, j].set(A[:, j] / jnp.linalg.norm(A[:, j]))
+
+    return A

morphomatics/nn/__init__.py

@@ -0,0 +1,15 @@
+################################################################################
+#                                                                              #
+#   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 .tangent_layers import TangentMLP, TangentInvariant
+from .wFM_layers import MfdFC, MfdInvariant
+from .flow_layers import FlowLayer, MfdGcnBlock

morphomatics/nn/flow_layers.py

@@ -0,0 +1,219 @@
+################################################################################
+#                                                                              #
+#   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 typing import Iterable, Callable
+
+import jax
+import jax.numpy as jnp
+from jax.scipy.sparse.linalg import bicgstab
+
+import jraph
+from flax import linen as nn
+
+from morphomatics.manifold import Manifold, PowerManifold
+from morphomatics.opt import RiemannianNewtonRaphson
+from morphomatics.graph.operators import mfdg_laplace
+from morphomatics.nn.tangent_layers import TangentMLP
+
+
+class FlowLayer(nn.Module):
+    """
+    Graph flow layer for graphs with manifold-valued features. The flow equation is integrated explicitly by default,
+    but an implicit scheme is also available.
+
+    See
+
+    M. Hanik, G, Steidl, C. v. Tycowicz. "Manifold GCN: Diffusion-based Convolutional Neural Network for
+    Manifold-valued Graphs" (https://arxiv.org/abs/2401.14381)
+
+    for a detailed description.
+
+    Inputs:
+    :param M: manifold in which the features lie
+    :param n_steps: number of explicit steps to approximate the flow with explicit Euler
+    :param implicit: boolean indicating whether to use implicit or explicit Euler integration
+    :param max_step_length: maximum step size for Euler integration
+
+    Note: Too long Euler steps can lead to numerical instabilities with some manifolds, e.g., the hyperbolic space.
+    In this case, a maximal step length should be ued.
+
+    """
+
+    M: Manifold
+    n_steps: int = 1
+    implicit: bool = False
+    max_step_length: float = jnp.inf
+    t_init: Callable = lambda *args: nn.initializers.truncated_normal(stddev=1.)(*args) + 1.
+    delta_init: Callable = lambda *args: nn.initializers.truncated_normal(stddev=1.)(*args) + 1.
+
+    def _single_euler_step(self, G: jraph.GraphsTuple, time: jnp.ndarray, delta: jnp.ndarray) -> jraph.GraphsTuple:
+        """Single step of the explicit Euler method for diffusion
+
+        :param G: graph with manifold valued vectors as features; length of vector must equal the flow layer width
+        :param time: vector of time parameters (same length as number of features)
+        :param delta: vector of "minimal step sizes"
+        :return: updated graph
+        """
+
+        def _multi_laplace(channel):
+            return mfdg_laplace(self.M, G._replace(nodes=channel))
+
+        def _activation(feature, vector, d):
+            nrm = jnp.sqrt(self.M.metric.inner(feature, vector, vector) + jnp.finfo(jnp.float64).eps)
+            alp = jax.nn.sigmoid(nrm - d)
+
+            # make sure that the step size is not larger than max_step_length
+            return jax.lax.cond(nrm * alp <= self.max_step_length,
+                                lambda w: alp * w,
+                                lambda w: w * self.max_step_length / nrm, vector)
+
+        v = jax.vmap(_multi_laplace, in_axes=1, out_axes=1)(G.nodes)
+
+        # ReLU-type activation
+        delta = jnp.stack([delta, ] * v.shape[0])
+        v = jax.vmap(jax.vmap(_activation))(G.nodes, v, delta)
+
+        v = -v * time.reshape((1, -1) + (1,) * (v.ndim - 2))
+        x = jax.vmap(jax.vmap(self.M.connec.exp))(G.nodes, v)
+        return G._replace(nodes=x)
+
+    def _implicit_euler_step(self, G: jraph.GraphsTuple, time: jnp.ndarray, delta=None) -> jraph.GraphsTuple:
+        """Single step of the implicit Euler method for diffusion
+
+        :param G: graph with manifold valued vectors as features; length of vector must equal the flow layer width
+        :param time: vector of time parameters (same length as number of features)
+        :param delta: only needed for out of syntax reasons
+        :return: updated graph
+        """
+        # n_nodes x n_channels x point_shape
+        n, c, *shape = G.nodes.shape
+
+        # power manifold
+        P = PowerManifold(self.M, n * c)
+
+        # current state
+        x_cur = G.nodes.reshape(-1, *shape)
+
+        # root of F characterizes solution to implicit Euler step
+        def F(x: jnp.array):
+            L = lambda a: mfdg_laplace(self.M, G._replace(nodes=a))
+            Lx = jax.vmap(L, in_axes=1, out_axes=1)(x.reshape(n, c, *shape))
+            tLx = Lx * time.reshape((1, -1) + (1,) * len(shape))
+            diff = P.connec.log(x, x_cur)
+            return diff - tLx.reshape(-1, *shape)
+
+        # x_next = RiemannianNewtonRaphson.solve(P, F, x_cur, maxiter=1)
+        ###############################
+        # unroll single interation
+        ###############################
+        # solve for update direction: v = -J⁻¹F(x)
+        J = lambda v: jax.jvp(F, (x_cur,), (v,))[1]
+        v, _ = bicgstab(J, -F(x_cur))
+        # step
+        x_next = P.connec.exp(x_cur, v)
+
+        return G._replace(nodes=x_next.reshape(n, c, *shape))
+
+    @nn.compact
+    def __call__(self, G: jraph.GraphsTuple) -> jraph.GraphsTuple:
+        """
+        :param G: graphs tuple with features of shape: num_nodes * num_channels * point_shape
+        :return: graphs tuple with features of shape: num_nodes * num_channels * point_shape
+
+        Apply discretized diffusion flow (with final activation) to each channel.
+        """
+        step_method = self._implicit_euler_step if self.implicit else self._single_euler_step
+
+        width = G.nodes.shape[1]  # number of channels
+        ####################
+        t = self.param("t_sqrt", self.t_init, (width,), G.nodes.dtype)
+        delta = self.param("delta_sqrt", self.delta_init, (width,), G.nodes.dtype)
+
+        # map to non-negative weights
+        t = t ** 2
+        delta = delta ** 2
+
+        def step(graph, _):
+            graph = step_method(graph, t / self.n_steps, delta)
+            return graph, None
+
+        G, _ = jax.lax.scan(step, G, None, self.n_steps, unroll=self.n_steps)
+
+        return G
+
+
+class MfdGcnBlock(nn.Module):
+    """
+    Manifold convolution network block as proposed in
+    M. Hanik, G, Steidl, C. v. Tycowicz. "Manifold GCN: Diffusion-based Convolutional Neural Network for
+    Manifold-valued Graphs" (https://arxiv.org/abs/2401.14381)
+
+    An explicit skip connection can be added that attaches the input to the output as an additional channels.
+    Note: If skip connections between each flow-tMLP unit are wanted, use several blocks with only one layer.
+
+    :param M: manifold constituting the signal domain
+    :param channel_sizes: sequence of channel sizes
+    :param n_steps: number of Euler steps that are performed in the flow layer
+    :param implicit: boolean indicating whether to use implicit or explicit Euler integration
+    :param max_step_length: maximum step size for Euler integration (see the flow layer)
+    :param explicit_skip: boolean indicating whether additionally to perform an explicit skip connection
+    :param inputs_are_copies: when true, only the first input channel is passed through by the skip connection
+    """
+
+    M: Manifold
+    channel_sizes: Iterable[int]
+    n_steps: int = 1
+    implicit: bool = False
+    max_step_length: float = jnp.inf
+    explicit_skip: bool = False
+    inputs_are_copies: bool = False
+
+    def setup(self):
+        layers = []
+        channel_sizes = tuple(self.channel_sizes)
+        for i, channel_size in enumerate(channel_sizes):
+            layers.append(
+                (
+                    FlowLayer(self.M, self.n_steps, self.implicit, self.max_step_length),
+                    TangentMLP(self.M, (channel_size,))
+                )
+            )
+        self.layers = tuple(layers)
+
+    def __call__(self, G: jraph.GraphsTuple) -> jraph.GraphsTuple:
+        """
+        :param G: graphs tuple with features of shape: num_nodes * num_channels * point_shape
+        :return: graphs tuple with features of shape: num_nodes * out_channels * point_shape
+
+        We use jraph pooling; hence, the batches are combined in the same graph and thus "hidden" in num_nodes.
+        The number of output channels is the number of output channels of the last (tangentMLP) layer plus, if
+        activated, the number of channels that are passed through by the skip connection (either num_channels or 1).
+        """
+
+        # save input for the skip connection
+        if self.explicit_skip:
+            if self.inputs_are_copies:
+                z = G.nodes[:, 1, None]
+            else:
+                z = G.nodes
+
+        for layer_unit in self.layers:
+            # flow layer
+            G = layer_unit[0](G)
+            # tangent MLP
+            G = G._replace(nodes=layer_unit[1](G.nodes[None])[0])
+
+        # skip connection
+        if self.explicit_skip:
+            G = G._replace(nodes=jax.lax.concatenate([z, G.nodes], 1))
+
+        return G

morphomatics/nn/tangent_layers.py

@@ -0,0 +1,176 @@
+################################################################################
+#                                                                              #
+#   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 typing import Sequence
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+import flax.linen as nn
+
+from morphomatics.manifold import Manifold
+
+
+class TangentMLP(nn.Module):
+    """
+    Generalized multi-layer perceptron for manifold-valued features as proposed in
+    M. Hanik, G, Steidl, C. v. Tycowicz. "Manifold GCN: Diffusion-based Convolutional Neural Network for
+    Manifold-valued Graphs" (https://arxiv.org/abs/2401.14381).
+
+    :param M: Manifold input signal takes values in
+    :param out_sizes: number of output feature channels (sequence thereof)
+    """
+
+    M: Manifold
+    out_sizes: Sequence[int]
+
+    @nn.compact
+    def __call__(self, x):
+        """
+        Apply tangent MLP layer.
+        :param x: input sequence with shape: batch * sequence_length * in_channel * M.point_shape
+            (M being the underlying manifold)
+        :return: output with shape: batch * sequence_length * out_channel * M.point_shape
+        """
+
+        n_batch, n_seq, n_in, *pt_shape = x.shape
+
+        # flatten first two axes -> shape: (batch * sequence_length) * in_channel * M.point_shape
+        x = x.reshape(-1, n_in, *pt_shape)
+
+        # compute logs -> tangent vectors
+        pernode = jax.vmap(self.M.connec.log, in_axes=(None, 0))
+        v = jax.vmap(pernode)(x[:, 0], x[:, 1:])
+
+        # shape into vectors
+        v = v.reshape(n_batch, n_seq, n_in-1, -1)
+
+        # apply vector neuron MLP
+        v = VectorNeuronMLP(self.out_sizes)(v)
+
+        # shape back into tangent vectors
+        v = v.reshape(n_batch * n_seq, -1, *pt_shape)
+
+        # map back to manifold
+        pernode = jax.vmap(self.M.connec.exp, in_axes=(None, 0))
+        y = jax.vmap(pernode)(x[:, 0], v)
+
+        return y.reshape(n_batch, n_seq, -1, *pt_shape)
+
+
+class TangentInvariant(nn.Module):
+    """
+    Invariant layer for manifold-valued features extending the TangentMLP layer.
+    Specifically, computes inner products of the input (linearized via Log) with
+     a set of tangent vectors obtained from them (using TangentMLP).
+     Finally, the products are passed through a fully connected layer to match desired output size.
+
+    :param M: Manifold input signal takes values in
+    :param out_channel: number of output feature channels
+    :param vec_sizes: sequence of widths for TangentMLP
+    :param use_bias: whether to use bias in the final fully connected layer
+    """
+
+    M: Manifold
+    out_channel: int
+    vec_sizes: Sequence[int] = (3,)
+    use_bias: bool = True
+
+    @nn.compact
+    def __call__(self, x):
+        """
+        Apply tangent MLP layer.
+        :param x: input sequence with shape: batch * sequence_length * in_channel * M.point_shape
+            (M being the underlying manifold)
+        :return: output with shape: batch * sequence_length * out_channel
+        """
+
+        n_batch, n_seq, n_in, *pt_shape = x.shape
+
+        # flatten first two axes -> shape: (batch * sequence_length) * in_channel * M.point_shape
+        x = x.reshape(-1, n_in, *pt_shape)
+
+        # compute logs -> tangent vectors
+        log_batched = jax.vmap(self.M.connec.log, in_axes=(None, 0))
+        v = jax.vmap(log_batched)(x[:, 0], x[:, 1:])
+
+        # shape into vectors
+        w = v.reshape(n_batch, n_seq, n_in-1, -1)
+
+        # apply vector neuron MLP
+        w = VectorNeuronMLP(self.vec_sizes)(w)
+
+        # shape back into tangent vectors
+        w = w.reshape(n_batch * n_seq, -1, *pt_shape)
+
+        # lower indices of tangent vectors (either v or w)
+        flat_batched = jax.vmap(self.M.metric.flat, in_axes=(None, 0))
+        if w.shape[1] > n_in:
+            v = jax.vmap(flat_batched)(x[:, 0], v)
+        else:
+            w = jax.vmap(flat_batched)(x[:, 0], w)
+        # compute inner products
+        y = jnp.einsum('...ij,...kj', v.reshape(*v.shape[:2], -1), w.reshape(*w.shape[:2], -1))
+
+        f = nn.Dense(self.out_channel, use_bias=self.use_bias)
+        return f(y.reshape(n_batch, n_seq, -1))
+
+
+class VectorNeuronMLP(nn.Module):
+    """
+    Vector Neuron MLP layer as described in
+    Deng, C., Litany, O., Duan, Y., Poulenard, A., Tagliasacchi, A., & Guibas, L. J. (2021).
+    Vector neurons: A general framework for SO(3)-equivariant networks.
+    In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 12200-12209).
+
+    :param output_sizes: sequence (length: m+1) of layer widths
+    """
+
+    output_sizes: Sequence[int]
+    negative_slope: float = 0.2
+
+    @nn.compact
+    def __call__(self, x: jnp.array):
+        """
+        Apply layer.
+        :param x: input sequence with shape: batch * sequence_length * in_channel * vector_dim
+        :return: output with shape: batch * sequence_length * out_channel * vector_shape
+        """
+
+        n_in = x.shape[2]
+
+        # apply layers
+        for i, n_out in enumerate(self.output_sizes):
+            # initialize weights
+            u_init = nn.initializers.truncated_normal(stddev=np.sqrt(2/n_in))
+            w_init = lambda *args: nn.initializers.truncated_normal(stddev=1. / np.sqrt(n_in))(*args) + 1. / n_in
+
+            U = self.param(f"U_{i}", u_init, (n_out, n_in), x.dtype)
+            W = self.param(f"W_{i}", w_init, (n_out, n_in), x.dtype)
+
+            n_in = n_out
+
+            # compute direction k (and its squared norm)
+            k = jnp.einsum('ij,...jk', U, x)
+            sqnrm_k = jnp.sum(k**2, axis=-1, keepdims=True) + np.finfo(np.float64).eps
+
+            # compute feature q
+            q = jnp.einsum('ij,...jk', W, x)
+
+            # Rotation-equivariant ReLU
+            dot_qk = jnp.sum(q * k, axis=-1, keepdims=True)
+            x = q + k * jax.nn.relu(-dot_qk) / sqnrm_k
+            # Leaky ReLU (weighted average of q and ReLU(q))
+            x = self.negative_slope * q + (1 - self.negative_slope) * x
+
+        return x