1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import numpy as num
8from pyrocko import orthodrome as od
10d2r = num.pi/180.
11r2d = 1.0 / d2r
14def arr_vertices(obj):
15 a = num.array(obj, dtype=num.float64)
16 assert len(a.shape) == 2
17 assert a.shape[1] == 3
18 return a
21def arr_faces(obj):
22 a = num.array(obj, dtype=num.int64)
23 assert len(a.shape) == 2
24 return a
27def vsqr(vertices):
28 return num.sum(vertices**2, axis=1)
31def vdot(a, b):
32 return num.sum(a*b, axis=-1)
35def vnorm(vertices):
36 return num.linalg.norm(vertices, axis=1)
39def normalize(vertices):
40 return vertices / vnorm(vertices)[:, num.newaxis]
43def face_centers(vertices, faces):
44 return num.mean(vertices[faces], axis=1)
47def rtp2xyz(rtp):
48 r = rtp[:, 0]
49 theta = rtp[:, 1]
50 phi = rtp[:, 2]
51 vecs = num.empty(rtp.shape, dtype=num.float64)
52 vecs[:, 0] = r*num.sin(theta)*num.cos(phi)
53 vecs[:, 1] = r*num.sin(theta)*num.sin(phi)
54 vecs[:, 2] = -r*num.cos(theta)
55 return vecs
58def xyz2rtp(xyz):
59 x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
60 vecs = num.empty(xyz.shape, dtype=num.float64)
61 vecs[:, 0] = num.sqrt(x**2+y**2+z**2)
62 vecs[:, 1] = num.arctan2(num.sqrt(x**2+y**2), z)
63 vecs[:, 2] = num.arctan2(y, x)
64 return vecs
67def latlon2xyz(latlon, radius=1.0):
68 rtp = num.empty((latlon.shape[0], 3))
69 rtp[:, 0] = radius
70 rtp[:, 1] = (latlon[:, 0] + 90.) * d2r
71 rtp[:, 2] = latlon[:, 1] * d2r
72 return rtp2xyz(rtp)
75def latlondepth2xyz(latlondepth, planetradius):
76 rtp = num.empty((latlondepth.shape[0], 3))
77 rtp[:, 0] = 1.0 - (latlondepth[:, 2] / planetradius)
78 rtp[:, 1] = (latlondepth[:, 0] + 90.) * d2r
79 rtp[:, 2] = latlondepth[:, 1] * d2r
80 return rtp2xyz(rtp)
83def ned2xyz(ned, latlondepth, planetradius):
84 endpoints = num.empty_like(latlondepth)
85 endpoints[:, 0], endpoints[:, 1] = od.ne_to_latlon(
86 latlondepth[:, 0], latlondepth[:, 1], ned[:, 0], ned[:, 1])
87 endpoints[:, 2] = latlondepth[:, 2] + ned[:, 2]
89 start_xyz = latlondepth2xyz(latlondepth, planetradius)
90 end_xyz = latlondepth2xyz(endpoints, planetradius)
92 return end_xyz - start_xyz
95def topo_to_vertices(lat, lon, ele, planetradius):
96 nlat = lat.size
97 nlon = lon.size
98 assert nlat > 1 and nlon > 1 and ele.shape == (nlat, nlon)
99 rtp = num.empty((ele.size, 3))
100 rtp[:, 0] = 1.0 + ele.flatten() / planetradius
101 rtp[:, 1] = (num.repeat(lat, nlon) + 90.) * d2r
102 rtp[:, 2] = num.tile(lon, nlat) * d2r
103 vertices = rtp2xyz(rtp)
104 return vertices
107g_topo_faces = {}
110def topo_to_faces(nlat, nlon):
111 k = (nlat, nlon)
112 if k not in g_topo_faces:
114 nfaces = 2 * (nlat - 1) * (nlon - 1)
115 faces = num.empty((nfaces, 3), dtype=num.int64)
116 ilon = num.arange(nlon - 1)
117 for ilat in range(nlat - 1):
118 ibeg = ilat*(nlon-1) * 2
119 iend = (ilat+1)*(nlon-1) * 2
120 faces[ibeg:iend:2, 0] = ilat*nlon + ilon
121 faces[ibeg:iend:2, 1] = ilat*nlon + ilon + 1
122 faces[ibeg:iend:2, 2] = ilat*nlon + ilon + nlon
123 faces[ibeg+1:iend+1:2, 0] = ilat*nlon + ilon + 1
124 faces[ibeg+1:iend+1:2, 1] = ilat*nlon + ilon + nlon + 1
125 faces[ibeg+1:iend+1:2, 2] = ilat*nlon + ilon + nlon
127 g_topo_faces[k] = faces
129 return g_topo_faces[k]
132g_topo_faces_quad = {}
135def topo_to_faces_quad(nlat, nlon):
136 k = (nlat, nlon)
137 if k not in g_topo_faces_quad:
139 nfaces = (nlat - 1) * (nlon - 1)
140 faces = num.empty((nfaces, 4), dtype=num.int64)
141 ilon = num.arange(nlon - 1)
142 for ilat in range(nlat - 1):
143 ibeg = ilat*(nlon-1)
144 iend = (ilat+1)*(nlon-1)
145 faces[ibeg:iend, 0] = ilat*nlon + ilon
146 faces[ibeg:iend, 1] = ilat*nlon + ilon + 1
147 faces[ibeg:iend, 2] = ilat*nlon + ilon + nlon + 1
148 faces[ibeg:iend, 3] = ilat*nlon + ilon + nlon
150 g_topo_faces_quad[k] = faces
152 return g_topo_faces_quad[k]
155def topo_to_mesh(lat, lon, ele, planetradius):
156 return \
157 topo_to_vertices(lat, lon, ele, planetradius), \
158 topo_to_faces_quad(*ele.shape)
161def refine_triangles(vertices, faces):
163 nv = vertices.shape[0]
164 nf = faces.shape[0]
165 assert faces.shape[1] == 3
167 fedges = num.concatenate((
168 faces[:, (0, 1)],
169 faces[:, (1, 2)],
170 faces[:, (2, 0)]))
172 fedges.sort(axis=1)
173 fedges1 = fedges[:, 0] + fedges[:, 1]*nv
174 sortorder = fedges1.argsort()
176 ne = fedges.shape[0] // 2
177 edges = fedges[sortorder[::2], :]
179 vertices_new = num.zeros((nv+ne, 3), dtype=num.float64)
180 vertices_new[:nv] = vertices
181 vertices_new[nv:] = 0.5 * (vertices[edges[:, 0]] + vertices[edges[:, 1]])
183 faces_new = num.zeros((nf*4, 3), dtype=num.int64)
185 vind = nv + sortorder.argsort()/2
187 faces_new[:nf, 0] = faces[:, 0]
188 faces_new[:nf, 1] = vind[:nf]
189 faces_new[:nf, 2] = vind[2*nf:]
191 faces_new[nf:nf*2, 0] = faces[:, 1]
192 faces_new[nf:nf*2, 1] = vind[nf:2*nf]
193 faces_new[nf:nf*2, 2] = vind[:nf]
195 faces_new[nf*2:nf*3, 0] = faces[:, 2]
196 faces_new[nf*2:nf*3, 1] = vind[2*nf:]
197 faces_new[nf*2:nf*3, 2] = vind[nf:2*nf]
199 faces_new[nf*3:, 0] = vind[:nf]
200 faces_new[nf*3:, 1] = vind[nf:2*nf]
201 faces_new[nf*3:, 2] = vind[2*nf:]
203 return vertices_new, faces_new
206def refine_triangle_parents(nfaces):
207 return num.arange(nfaces) % (nfaces / 4)
210def triangle_barycentric_transforms(vertices, faces):
211 a = vertices[faces[:, 1]] - vertices[faces[:, 0]]
212 b = vertices[faces[:, 2]] - vertices[faces[:, 0]]
214 norm = 1.0 / (vsqr(a) * vsqr(b) - vdot(a, b)**2)
216 transforms = num.zeros((faces.shape[0], 2, 3))
218 transforms[:, 0, :] = norm[:, num.newaxis] \
219 * (vsqr(b)[:, num.newaxis] * a - vdot(a, b)[:, num.newaxis] * b)
220 transforms[:, 1, :] = norm[:, num.newaxis] \
221 * (vsqr(a)[:, num.newaxis] * b - vdot(a, b)[:, num.newaxis] * a)
223 return transforms
226def triangle_barycentric_coordinates(transform, origin, points):
227 c = points - origin[num.newaxis, :]
228 uv = num.dot(transform[:, :], c.T).T
229 return uv
232def triangle_cartesian_system(vertices, faces):
233 a = normalize(vertices[faces[:, 1]] - vertices[faces[:, 0]])
234 b = vertices[faces[:, 2]] - vertices[faces[:, 1]]
235 c = normalize(num.cross(a, b))
236 uvw = num.zeros((faces.shape[0], 3, 3))
237 uvw[:, 0, :] = a
238 uvw[:, 1, :] = num.cross(c, a)
239 uvw[:, 2, :] = c
240 return uvw
243def triangle_planes(vertices, faces):
244 a = vertices[faces[:, 1]] - vertices[faces[:, 0]]
245 b = vertices[faces[:, 2]] - vertices[faces[:, 0]]
246 anormals = normalize(num.cross(a, b))
247 anormals *= vdot(vertices[faces[:, 0]], anormals)[:, num.newaxis]
248 return anormals