[1208] | 1 | import math |
---|
| 2 | import numpy as np |
---|
| 3 | |
---|
[1325] | 4 | |
---|
[1208] | 5 | def wcentre(matrix, weights): |
---|
[1325] | 6 | sw = weights.sum() |
---|
| 7 | swx = (matrix * weights).sum(axis=1) |
---|
| 8 | swx /= sw |
---|
| 9 | return (matrix.transpose() - swx).transpose() * np.sqrt(weights) |
---|
| 10 | |
---|
| 11 | |
---|
[1208] | 12 | def weightedMDS(distances, weights): |
---|
[1325] | 13 | n = len(weights) |
---|
| 14 | distances = distances ** 2 |
---|
| 15 | for i in range(2): |
---|
| 16 | distances = wcentre(distances, weights) |
---|
| 17 | distances = distances.T |
---|
| 18 | distances *= -0.5 |
---|
| 19 | _, eigenvalues, vh = np.linalg.svd(distances) |
---|
| 20 | W = (vh / np.sqrt(weights)).T |
---|
| 21 | S = np.zeros((n, n)) |
---|
| 22 | np.fill_diagonal(S, eigenvalues) |
---|
| 23 | S = S ** 0.5 |
---|
| 24 | dcoords = W.dot(S) |
---|
| 25 | coords = np.zeros((n, 3)) |
---|
| 26 | coords[:, 0] = dcoords[:, 0] |
---|
| 27 | for i in range(1, 3): |
---|
| 28 | if n > i: |
---|
| 29 | coords[:, i] = dcoords[:, i] |
---|
| 30 | return coords |
---|
[1208] | 31 | |
---|
[1325] | 32 | |
---|
[1208] | 33 | def align(model, fixedZaxis=False): |
---|
[1325] | 34 | numparts = model.numparts._value() |
---|
[1208] | 35 | distmatrix = np.zeros((numparts, numparts), dtype=float) |
---|
| 36 | for p1 in range(numparts): |
---|
[1325] | 37 | for p2 in range(p1 + 1, numparts): # only calculate a triangle since Euclidean distance is symmetrical |
---|
| 38 | P1 = model.getPart(p1) |
---|
| 39 | P2 = model.getPart(p2) |
---|
[1208] | 40 | if fixedZaxis: |
---|
[1325] | 41 | # fixed vertical axis, so pretend all points are on the xy plane |
---|
[1208] | 42 | z_dist = 0 |
---|
| 43 | else: |
---|
[1325] | 44 | z_dist = (P1.z._value() - P2.z._value()) ** 2 |
---|
| 45 | distmatrix[p1, p2] = distmatrix[p2, p1] = math.sqrt((P1.x._value() - P2.x._value()) ** 2 + (P1.y._value() - P2.y._value()) ** 2 + z_dist) |
---|
| 46 | |
---|
[1208] | 47 | if model.numjoints._value() > 0: |
---|
[1325] | 48 | weightvector = np.zeros((numparts), dtype=int) |
---|
[1208] | 49 | else: |
---|
[1325] | 50 | weightvector = np.ones((numparts), dtype=int) |
---|
| 51 | |
---|
[1208] | 52 | for j in range(model.numjoints._value()): |
---|
[1325] | 53 | J = model.getJoint(j) |
---|
| 54 | weightvector[J.p1._value()] += 1 |
---|
| 55 | weightvector[J.p2._value()] += 1 |
---|
| 56 | weightvector = weightvector.astype(float) # convert to float once, since later it would be promoted to float so many times anyway... |
---|
[1208] | 57 | coords = weightedMDS(distmatrix, weightvector) |
---|
| 58 | |
---|
| 59 | # update parts positions |
---|
| 60 | n = len(weightvector) |
---|
| 61 | for p in range(numparts): |
---|
| 62 | P = model.getPart(p) |
---|
| 63 | P.x = coords[p, 0] |
---|
| 64 | if n > 1: |
---|
| 65 | P.y = coords[p, 1] |
---|
| 66 | if n > 2: |
---|
| 67 | if not fixedZaxis: |
---|
| 68 | P.z = coords[p, 2] |
---|
| 69 | |
---|
| 70 | if fixedZaxis: |
---|
| 71 | if np.shape(coords)[1] > 2: |
---|
[1325] | 72 | # restore original z coordinate |
---|
[1208] | 73 | for p in range(numparts): |
---|
[1325] | 74 | P = model.getPart(p) |
---|
| 75 | coords[p, 2] = P.z._value() |
---|