Coverage for src / sdynpy / signal_processing / sdynpy_geometry_fitting.py: 14%

1# -*- coding: utf-8 -*- 

2""" 

3Functions to fit geometry to point data 

4 

5This module defines a Geometry object as well as all of the subcomponents of 

6a geometry object: nodes, elements, tracelines and coordinate system. Geometry 

7plotting is also handled in this module. 

8""" 

9""" 

10Copyright 2022 National Technology & Engineering Solutions of Sandia, 

11LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S. 

12Government retains certain rights in this software. 

13 

14This program is free software: you can redistribute it and/or modify 

15it under the terms of the GNU General Public License as published by 

16the Free Software Foundation, either version 3 of the License, or 

17(at your option) any later version. 

18 

19This program is distributed in the hope that it will be useful, 

20but WITHOUT ANY WARRANTY; without even the implied warranty of 

21MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

22GNU General Public License for more details. 

23 

24You should have received a copy of the GNU General Public License 

25along with this program. If not, see <https://www.gnu.org/licenses/>. 

26""" 

27 

28 

29import numpy as np 

30import scipy.optimize as opt 

31 

32 

33def cone_fit(points, origin, direction, angle): 

34 a = origin 

35 p = points 

36 n = direction 

37 height = np.einsum('...i,i->...', (a - p), n) 

38 distance = np.linalg.norm((a-p) - (height[:, np.newaxis])*n, axis=1) 

39 expected_distance = -np.tan(angle)*height 

40 return np.linalg.norm((distance-expected_distance)) 

41 

42 

43def create_cone(ncircum, naxial, height, angle, origin, direction): 

44 circum_angles = np.linspace(0, 2*np.pi, ncircum, endpoint=False) 

45 heights = np.linspace(0, height, naxial+1)[1:] 

46 points = [] 

47 for i, cangle in enumerate(circum_angles): 

48 for j, h in enumerate(heights): 

49 radius = np.tan(angle)*h 

50 points.append([np.sin(cangle)*radius, np.cos(cangle)*radius, h]) 

51 points = np.array(points) 

52 rotation_matrix_z = direction/np.linalg.norm(direction) 

53 rotation_matrix_y = np.cross(rotation_matrix_z, [1.0, 0.0, 0.0]) 

54 if np.linalg.norm(rotation_matrix_y) < 1e-10: 

55 rotation_matrix_x = np.cross([0.0, 1.0, 0.0], rotation_matrix_z) 

56 rotation_matrix_x /= np.linalg.norm(rotation_matrix_x) 

57 rotation_matrix_y = np.cross(rotation_matrix_z, rotation_matrix_x) 

58 rotation_matrix_y /= np.linalg.norm(rotation_matrix_y) 

59 else: 

60 rotation_matrix_y /= np.linalg.norm(rotation_matrix_y) 

61 rotation_matrix_x = np.cross(rotation_matrix_y, rotation_matrix_z) 

62 rotation_matrix_x /= np.linalg.norm(rotation_matrix_x) 

63 R = np.array((rotation_matrix_x, rotation_matrix_y, rotation_matrix_z)).T 

64 new_points = ((R@points.T).T+origin) 

65 return new_points 

66 

67 

68def cone_error_fn_fixed_angle(points, angle): 

69 return lambda x: cone_fit(points, x[:3], x[3:], angle) 

70 

71 

72def cone_error_fn_free_angle(points): 

73 return lambda x: cone_fit(points, x[:3], x[3:6], x[6]) 

74 

75 

76def fit_cone_fixed_angle(points, angle): 

77 

78 def constraint_eqn(x): 

79 return np.linalg.norm(x[3:]) 

80 

81 opt_fn = cone_error_fn_fixed_angle(points, angle) 

82 origin_guess = np.mean(points, axis=0) 

83 u, s, vh = np.linalg.svd(points-origin_guess) 

84 direction_guess = vh[0] 

85 x0_1 = np.concatenate((origin_guess, direction_guess)) 

86 x0_2 = np.concatenate((origin_guess, -direction_guess)) 

87 # Constrain to a unit vector 

88 constraint = opt.NonlinearConstraint(constraint_eqn, 1.0, 1.0) 

89 result_1 = opt.minimize(opt_fn, x0_1, method='trust-constr', constraints=constraint) 

90 result_2 = opt.minimize(opt_fn, x0_2, method='trust-constr', constraints=constraint) 

91 # Figure out which is the better one 

92 result = result_1 if result_1.fun < result_2.fun else result_2 

93 cone_origin = result.x[:3] 

94 cone_direction = result.x[3:] 

95 return result.fun, cone_origin, cone_direction 

96 

97 

98def distance_point_line(points, origin, direction): 

99 return np.linalg.norm(np.cross((points-origin), direction), axis=-1)/np.linalg.norm(direction) 

100 

101 

102def distance_point_plane(point, plane_point, plane_direction): 

103 return abs(np.einsum('...i,...i->...', point-plane_point, plane_direction))/np.linalg.norm(plane_direction, axis=-1) 

104 

105 

106def cylinder_fit(points, origin, direction, radius): 

107 return np.linalg.norm(distance_point_line(points, origin, direction) - radius) 

108 

109 

110def fit_cylinder(points, origin_guess=None, direction_guess=None, radius_guess=None): 

111 def opt_fn(x): 

112 return cylinder_fit(points, x[:3], x[3:6], x[6:7]) 

113 

114 def constraint_eqn(x): 

115 return np.linalg.norm(x[3:6]) 

116 

117 if origin_guess is None: 

118 origin_guess = np.mean(points, axis=0) 

119 if direction_guess is None or direction_guess == 'largest': 

120 u, s, vh = np.linalg.svd(points-origin_guess) 

121 direction_guess = vh[0] 

122 elif direction_guess == 'smallest': 

123 u, s, vh = np.linalg.svd(points-origin_guess) 

124 direction_guess = vh[-1] 

125 if radius_guess is None: 

126 u, s, vh = np.linalg.svd(points-origin_guess) 

127 radius_guess = s[0]/2 

128 x0 = np.concatenate((origin_guess, direction_guess, [radius_guess])) 

129 # Constrain to a unit vector 

130 constraint = opt.NonlinearConstraint(constraint_eqn, 1.0, 1.0) 

131 result = opt.minimize(opt_fn, x0, method='trust-constr', constraints=constraint) 

132 origin = result.x[:3] 

133 direction = result.x[3:6] 

134 radius = result.x[6] 

135 return result.fun, origin, direction, radius 

136 

137 

138def fit_line_point_cloud(points): 

139 center = np.mean(points, axis=0) 

140 shifted_points = points-center 

141 cov = points.T@shifted_points 

142 evals, evects = np.linalg.eigh(cov) 

143 direction = evects[:, np.argmax(evals)] 

144 direction /= np.linalg.norm(direction) 

145 return lambda t: center+t*direction 

146 

147def intersection_point_multiple_lines(P0,P1): 

148 """Computes the intersection point of multiple lines in a least squares sense 

149  

150 P0 and P1 are points on a line. Each are NxD where N is the number of lines 

151 and D is the dimensionality of the space (D=2 corresponds to lines on a plane, 

152 D=3 is points in 3D space). 

153 """ 

154 n = (P1-P0)/np.linalg.norm(P1-P0,axis=1)[:,np.newaxis] 

155 projectors = np.eye(n.shape[1]) - n[:,:,np.newaxis]*n[:,np.newaxis] # I - n*n.T 

156 R = projectors.sum(axis=0) 

157 q = (projectors @ P0[:,:,np.newaxis]).sum(axis=0) 

158 p = np.linalg.lstsq(R,q,rcond=None)[0] 

159 

160 return p