source: framspy/dissimilarity/alignmodel.py @ 1310

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(numparts): #TODO optimize, only calculate a triangle
                        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]=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()