source: framspy/dissimilarity/alignmodel.py @ 1325

import math
import numpy as np


def wcentre(matrix, weights):
        sw = weights.sum()
        swx = (matrix * weights).sum(axis=1)
        swx /= sw
        return (matrix.transpose() - swx).transpose() * np.sqrt(weights)


def weightedMDS(distances, weights):
        n = len(weights)
        distances = distances ** 2
        for i in range(2):
                distances = wcentre(distances, weights)
                distances = distances.T
        distances *= -0.5
        _, eigenvalues, vh = np.linalg.svd(distances)
        W = (vh / np.sqrt(weights)).T
        S = np.zeros((n, n))
        np.fill_diagonal(S, eigenvalues)
        S = S ** 0.5
        dcoords = W.dot(S)
        coords = np.zeros((n, 3))
        coords[:, 0] = dcoords[:, 0]
        for i in range(1, 3):
                if n > i:
                        coords[:, i] = dcoords[:, i]
        return coords


def align(model, fixedZaxis=False):
        numparts = model.numparts._value()
        distmatrix = np.zeros((numparts, numparts), dtype=float)
        for p1 in range(numparts):
                for p2 in range(p1 + 1, numparts):  # only calculate a triangle since Euclidean distance is symmetrical
                        P1 = model.getPart(p1)
                        P2 = model.getPart(p2)
                        if fixedZaxis:
                                # fixed vertical axis, so pretend all points are on the xy plane
                                z_dist = 0
                        else:
                                z_dist = (P1.z._value() - P2.z._value()) ** 2
                        distmatrix[p1, p2] = distmatrix[p2, p1] = math.sqrt((P1.x._value() - P2.x._value()) ** 2 + (P1.y._value() - P2.y._value()) ** 2 + z_dist)

        if model.numjoints._value() > 0:
                weightvector = np.zeros((numparts), dtype=int)
        else:
                weightvector = np.ones((numparts), dtype=int)

        for j in range(model.numjoints._value()):
                J = model.getJoint(j)
                weightvector[J.p1._value()] += 1
                weightvector[J.p2._value()] += 1
        weightvector = weightvector.astype(float)  # convert to float once, since later it would be promoted to float so many times anyway...
        coords = weightedMDS(distmatrix, weightvector)

        # update parts positions
        n = len(weightvector)
        for p in range(numparts):
                P = model.getPart(p)
                P.x = coords[p, 0]
                if n > 1:
                        P.y = coords[p, 1]
                if n > 2:
                        if not fixedZaxis:
                                P.z = coords[p, 2]

        if fixedZaxis:
                if np.shape(coords)[1] > 2:
                        # restore original z coordinate
                        for p in range(numparts):
                                P = model.getPart(p)
                                coords[p, 2] = P.z._value()