1 | import math |
---|

2 | import numpy as np |
---|

3 | |
---|

4 | def wcentre(matrix, weights): |
---|

5 | sw = weights.sum() |
---|

6 | swx = (matrix*weights).sum(axis=1) |
---|

7 | swx /= sw |
---|

8 | return (matrix.transpose()-swx).transpose()*np.sqrt(weights) |
---|

9 | |
---|

10 | def weightedMDS(distances, weights): |
---|

11 | n = len(weights) |
---|

12 | distances = distances**2 |
---|

13 | for i in range(2): |
---|

14 | distances = wcentre(distances, weights) |
---|

15 | distances = distances.T |
---|

16 | distances *= -0.5 |
---|

17 | _, eigenvalues, vh = np.linalg.svd(distances) |
---|

18 | W = (vh/np.sqrt(weights)).T |
---|

19 | S = np.zeros((n,n)) |
---|

20 | np.fill_diagonal(S, eigenvalues) |
---|

21 | S = S**0.5 |
---|

22 | dcoords = W.dot(S) |
---|

23 | coords = np.zeros((n, 3)) |
---|

24 | coords[:,0]=dcoords[:,0] |
---|

25 | for i in range(1,3): |
---|

26 | if n>i: |
---|

27 | coords[:,i]=dcoords[:,i] |
---|

28 | return coords |
---|

29 | |
---|

30 | |
---|

31 | def align(model, fixedZaxis=False): |
---|

32 | numparts=model.numparts._value() |
---|

33 | distmatrix = np.zeros((numparts, numparts), dtype=float) |
---|

34 | for p1 in range(numparts): |
---|

35 | for p2 in range(numparts): #TODO optimize, only calculate a triangle |
---|

36 | P1=model.getPart(p1) |
---|

37 | P2=model.getPart(p2) |
---|

38 | if fixedZaxis: |
---|

39 | #fixed vertical axis, so pretend all points are on the xy plane |
---|

40 | z_dist = 0 |
---|

41 | else: |
---|

42 | z_dist = (P1.z._value()-P2.z._value())**2 |
---|

43 | distmatrix[p1,p2]=math.sqrt((P1.x._value()-P2.x._value())**2+(P1.y._value()-P2.y._value())**2+z_dist) |
---|

44 | |
---|

45 | if model.numjoints._value() > 0: |
---|

46 | weightvector=np.zeros((numparts), dtype=int) |
---|

47 | else: |
---|

48 | weightvector=np.ones((numparts), dtype=int) |
---|

49 | |
---|

50 | for j in range(model.numjoints._value()): |
---|

51 | J=model.getJoint(j) |
---|

52 | weightvector[J.p1._value()]+=1 |
---|

53 | weightvector[J.p2._value()]+=1 |
---|

54 | weightvector=weightvector.astype(float) # convert to float once, since later it would be promoted to float so many times anyway... |
---|

55 | coords = weightedMDS(distmatrix, weightvector) |
---|

56 | |
---|

57 | # update parts positions |
---|

58 | n = len(weightvector) |
---|

59 | for p in range(numparts): |
---|

60 | P = model.getPart(p) |
---|

61 | P.x = coords[p, 0] |
---|

62 | if n > 1: |
---|

63 | P.y = coords[p, 1] |
---|

64 | if n > 2: |
---|

65 | if not fixedZaxis: |
---|

66 | P.z = coords[p, 2] |
---|

67 | |
---|

68 | |
---|

69 | if fixedZaxis: |
---|

70 | if np.shape(coords)[1] > 2: |
---|

71 | #restore original z coordinate |
---|

72 | for p in range(numparts): |
---|

73 | P=model.getPart(p) |
---|

74 | coords[p,2]=P.z._value() |
---|

75 | |
---|