Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/geometry.py: 70%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Geometry helper functions for coordinate transforms and mesh manipulations. 

8''' 

9 

10import numpy as num 

11 

12from pyrocko import orthodrome as od 

13 

14d2r = num.pi/180. 

15r2d = 1.0 / d2r 

16 

17 

18def arr_vertices(obj): 

19 a = num.array(obj, dtype=num.float64) 

20 assert len(a.shape) == 2 

21 assert a.shape[1] == 3 

22 return a 

23 

24 

25def arr_faces(obj): 

26 a = num.array(obj, dtype=num.int64) 

27 assert len(a.shape) == 2 

28 return a 

29 

30 

31def vsqr(vertices): 

32 return num.sum(vertices**2, axis=1) 

33 

34 

35def vdot(a, b): 

36 return num.sum(a*b, axis=-1) 

37 

38 

39def vnorm(vertices): 

40 return num.linalg.norm(vertices, axis=1) 

41 

42 

43def normalize(vertices): 

44 return vertices / vnorm(vertices)[:, num.newaxis] 

45 

46 

47def face_centers(vertices, faces): 

48 return num.mean(vertices[faces], axis=1) 

49 

50 

51def rtp2xyz(rtp): 

52 r = rtp[:, 0] 

53 theta = rtp[:, 1] 

54 phi = rtp[:, 2] 

55 vecs = num.empty(rtp.shape, dtype=num.float64) 

56 vecs[:, 0] = r*num.sin(theta)*num.cos(phi) 

57 vecs[:, 1] = r*num.sin(theta)*num.sin(phi) 

58 vecs[:, 2] = -r*num.cos(theta) 

59 return vecs 

60 

61 

62def xyz2rtp(xyz): 

63 x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2] 

64 vecs = num.empty(xyz.shape, dtype=num.float64) 

65 vecs[:, 0] = num.sqrt(x**2+y**2+z**2) 

66 vecs[:, 1] = num.arctan2(num.sqrt(x**2+y**2), z) 

67 vecs[:, 2] = num.arctan2(y, x) 

68 return vecs 

69 

70 

71def latlon2xyz(latlon, radius=1.0): 

72 rtp = num.empty((latlon.shape[0], 3)) 

73 rtp[:, 0] = radius 

74 rtp[:, 1] = (latlon[:, 0] + 90.) * d2r 

75 rtp[:, 2] = latlon[:, 1] * d2r 

76 return rtp2xyz(rtp) 

77 

78 

79def latlondepth2xyz(latlondepth, planetradius): 

80 rtp = num.empty((latlondepth.shape[0], 3)) 

81 rtp[:, 0] = 1.0 - (latlondepth[:, 2] / planetradius) 

82 rtp[:, 1] = (latlondepth[:, 0] + 90.) * d2r 

83 rtp[:, 2] = latlondepth[:, 1] * d2r 

84 return rtp2xyz(rtp) 

85 

86 

87def ned2xyz(ned, latlondepth, planetradius): 

88 endpoints = num.empty_like(latlondepth) 

89 endpoints[:, 0], endpoints[:, 1] = od.ne_to_latlon( 

90 latlondepth[:, 0], latlondepth[:, 1], ned[:, 0], ned[:, 1]) 

91 endpoints[:, 2] = latlondepth[:, 2] + ned[:, 2] 

92 

93 start_xyz = latlondepth2xyz(latlondepth, planetradius) 

94 end_xyz = latlondepth2xyz(endpoints, planetradius) 

95 

96 return end_xyz - start_xyz 

97 

98 

99def topo_to_vertices(lat, lon, ele, planetradius): 

100 nlat = lat.size 

101 nlon = lon.size 

102 assert nlat > 1 and nlon > 1 and ele.shape == (nlat, nlon) 

103 rtp = num.empty((ele.size, 3)) 

104 rtp[:, 0] = 1.0 + ele.flatten() / planetradius 

105 rtp[:, 1] = (num.repeat(lat, nlon) + 90.) * d2r 

106 rtp[:, 2] = num.tile(lon, nlat) * d2r 

107 vertices = rtp2xyz(rtp) 

108 return vertices 

109 

110 

111g_topo_faces = {} 

112 

113 

114def topo_to_faces(nlat, nlon): 

115 k = (nlat, nlon) 

116 if k not in g_topo_faces: 

117 

118 nfaces = 2 * (nlat - 1) * (nlon - 1) 

119 faces = num.empty((nfaces, 3), dtype=num.int64) 

120 ilon = num.arange(nlon - 1) 

121 for ilat in range(nlat - 1): 

122 ibeg = ilat*(nlon-1) * 2 

123 iend = (ilat+1)*(nlon-1) * 2 

124 faces[ibeg:iend:2, 0] = ilat*nlon + ilon 

125 faces[ibeg:iend:2, 1] = ilat*nlon + ilon + 1 

126 faces[ibeg:iend:2, 2] = ilat*nlon + ilon + nlon 

127 faces[ibeg+1:iend+1:2, 0] = ilat*nlon + ilon + 1 

128 faces[ibeg+1:iend+1:2, 1] = ilat*nlon + ilon + nlon + 1 

129 faces[ibeg+1:iend+1:2, 2] = ilat*nlon + ilon + nlon 

130 

131 g_topo_faces[k] = faces 

132 

133 return g_topo_faces[k] 

134 

135 

136g_topo_faces_quad = {} 

137 

138 

139def topo_to_faces_quad(nlat, nlon): 

140 k = (nlat, nlon) 

141 if k not in g_topo_faces_quad: 

142 

143 nfaces = (nlat - 1) * (nlon - 1) 

144 faces = num.empty((nfaces, 4), dtype=num.int64) 

145 ilon = num.arange(nlon - 1) 

146 for ilat in range(nlat - 1): 

147 ibeg = ilat*(nlon-1) 

148 iend = (ilat+1)*(nlon-1) 

149 faces[ibeg:iend, 0] = ilat*nlon + ilon 

150 faces[ibeg:iend, 1] = ilat*nlon + ilon + 1 

151 faces[ibeg:iend, 2] = ilat*nlon + ilon + nlon + 1 

152 faces[ibeg:iend, 3] = ilat*nlon + ilon + nlon 

153 

154 g_topo_faces_quad[k] = faces 

155 

156 return g_topo_faces_quad[k] 

157 

158 

159def topo_to_mesh(lat, lon, ele, planetradius): 

160 return \ 

161 topo_to_vertices(lat, lon, ele, planetradius), \ 

162 topo_to_faces_quad(*ele.shape) 

163 

164 

165def refine_triangles(vertices, faces): 

166 

167 nv = vertices.shape[0] 

168 nf = faces.shape[0] 

169 assert faces.shape[1] == 3 

170 

171 fedges = num.concatenate(( 

172 faces[:, (0, 1)], 

173 faces[:, (1, 2)], 

174 faces[:, (2, 0)])) 

175 

176 fedges.sort(axis=1) 

177 fedges1 = fedges[:, 0] + fedges[:, 1]*nv 

178 sortorder = fedges1.argsort() 

179 

180 ne = fedges.shape[0] // 2 

181 edges = fedges[sortorder[::2], :] 

182 

183 vertices_new = num.zeros((nv+ne, 3), dtype=num.float64) 

184 vertices_new[:nv] = vertices 

185 vertices_new[nv:] = 0.5 * (vertices[edges[:, 0]] + vertices[edges[:, 1]]) 

186 

187 faces_new = num.zeros((nf*4, 3), dtype=num.int64) 

188 

189 vind = nv + sortorder.argsort()/2 

190 

191 faces_new[:nf, 0] = faces[:, 0] 

192 faces_new[:nf, 1] = vind[:nf] 

193 faces_new[:nf, 2] = vind[2*nf:] 

194 

195 faces_new[nf:nf*2, 0] = faces[:, 1] 

196 faces_new[nf:nf*2, 1] = vind[nf:2*nf] 

197 faces_new[nf:nf*2, 2] = vind[:nf] 

198 

199 faces_new[nf*2:nf*3, 0] = faces[:, 2] 

200 faces_new[nf*2:nf*3, 1] = vind[2*nf:] 

201 faces_new[nf*2:nf*3, 2] = vind[nf:2*nf] 

202 

203 faces_new[nf*3:, 0] = vind[:nf] 

204 faces_new[nf*3:, 1] = vind[nf:2*nf] 

205 faces_new[nf*3:, 2] = vind[2*nf:] 

206 

207 return vertices_new, faces_new 

208 

209 

210def refine_triangle_parents(nfaces): 

211 return num.arange(nfaces) % (nfaces / 4) 

212 

213 

214def triangle_barycentric_transforms(vertices, faces): 

215 a = vertices[faces[:, 1]] - vertices[faces[:, 0]] 

216 b = vertices[faces[:, 2]] - vertices[faces[:, 0]] 

217 

218 norm = 1.0 / (vsqr(a) * vsqr(b) - vdot(a, b)**2) 

219 

220 transforms = num.zeros((faces.shape[0], 2, 3)) 

221 

222 transforms[:, 0, :] = norm[:, num.newaxis] \ 

223 * (vsqr(b)[:, num.newaxis] * a - vdot(a, b)[:, num.newaxis] * b) 

224 transforms[:, 1, :] = norm[:, num.newaxis] \ 

225 * (vsqr(a)[:, num.newaxis] * b - vdot(a, b)[:, num.newaxis] * a) 

226 

227 return transforms 

228 

229 

230def triangle_barycentric_coordinates(transform, origin, points): 

231 c = points - origin[num.newaxis, :] 

232 uv = num.dot(transform[:, :], c.T).T 

233 return uv 

234 

235 

236def triangle_cartesian_system(vertices, faces): 

237 a = normalize(vertices[faces[:, 1]] - vertices[faces[:, 0]]) 

238 b = vertices[faces[:, 2]] - vertices[faces[:, 1]] 

239 c = normalize(num.cross(a, b)) 

240 uvw = num.zeros((faces.shape[0], 3, 3)) 

241 uvw[:, 0, :] = a 

242 uvw[:, 1, :] = num.cross(c, a) 

243 uvw[:, 2, :] = c 

244 return uvw 

245 

246 

247def triangle_planes(vertices, faces): 

248 a = vertices[faces[:, 1]] - vertices[faces[:, 0]] 

249 b = vertices[faces[:, 2]] - vertices[faces[:, 0]] 

250 anormals = normalize(num.cross(a, b)) 

251 anormals *= vdot(vertices[faces[:, 0]], anormals)[:, num.newaxis] 

252 return anormals