1 | import math |
---|
2 | import numpy as np |
---|
3 | |
---|
4 | |
---|
5 | def wcentre(matrix, weights): |
---|
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 | |
---|
12 | def weightedMDS(distances, weights): |
---|
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 |
---|
31 | |
---|
32 | |
---|
33 | def align(model, fixedZaxis=False): |
---|
34 | numparts = model.numparts._value() |
---|
35 | distmatrix = np.zeros((numparts, numparts), dtype=float) |
---|
36 | for p1 in range(numparts): |
---|
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) |
---|
40 | if fixedZaxis: |
---|
41 | # fixed vertical axis, so pretend all points are on the xy plane |
---|
42 | z_dist = 0 |
---|
43 | else: |
---|
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 | |
---|
47 | if model.numjoints._value() > 0: |
---|
48 | weightvector = np.zeros((numparts), dtype=int) |
---|
49 | else: |
---|
50 | weightvector = np.ones((numparts), dtype=int) |
---|
51 | |
---|
52 | for j in range(model.numjoints._value()): |
---|
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... |
---|
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: |
---|
72 | # restore original z coordinate |
---|
73 | for p in range(numparts): |
---|
74 | P = model.getPart(p) |
---|
75 | coords[p, 2] = P.z._value() |
---|