1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function, division
8import numpy as num
10from pyrocko import orthodrome as od
12d2r = num.pi/180.
13r2d = 1.0 / d2r
16def arr_vertices(obj):
17 a = num.array(obj, dtype=num.float64)
18 assert len(a.shape) == 2
19 assert a.shape[1] == 3
20 return a
23def arr_faces(obj):
24 a = num.array(obj, dtype=num.int64)
25 assert len(a.shape) == 2
26 return a
29def vsqr(vertices):
30 return num.sum(vertices**2, axis=1)
33def vdot(a, b):
34 return num.sum(a*b, axis=-1)
37def vnorm(vertices):
38 return num.linalg.norm(vertices, axis=1)
41def normalize(vertices):
42 return vertices / vnorm(vertices)[:, num.newaxis]
45def face_centers(vertices, faces):
46 return num.mean(vertices[faces], axis=1)
49def rtp2xyz(rtp):
50 r = rtp[:, 0]
51 theta = rtp[:, 1]
52 phi = rtp[:, 2]
53 vecs = num.empty(rtp.shape, dtype=num.float64)
54 vecs[:, 0] = r*num.sin(theta)*num.cos(phi)
55 vecs[:, 1] = r*num.sin(theta)*num.sin(phi)
56 vecs[:, 2] = -r*num.cos(theta)
57 return vecs
60def xyz2rtp(xyz):
61 x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
62 vecs = num.empty(xyz.shape, dtype=num.float64)
63 vecs[:, 0] = num.sqrt(x**2+y**2+z**2)
64 vecs[:, 1] = num.arctan2(num.sqrt(x**2+y**2), z)
65 vecs[:, 2] = num.arctan2(y, x)
66 return vecs
69def latlon2xyz(latlon, radius=1.0):
70 rtp = num.empty((latlon.shape[0], 3))
71 rtp[:, 0] = radius
72 rtp[:, 1] = (latlon[:, 0] + 90.) * d2r
73 rtp[:, 2] = latlon[:, 1] * d2r
74 return rtp2xyz(rtp)
77def latlondepth2xyz(latlondepth, planetradius):
78 rtp = num.empty((latlondepth.shape[0], 3))
79 rtp[:, 0] = 1.0 - (latlondepth[:, 2] / planetradius)
80 rtp[:, 1] = (latlondepth[:, 0] + 90.) * d2r
81 rtp[:, 2] = latlondepth[:, 1] * d2r
82 return rtp2xyz(rtp)
85def ned2xyz(ned, latlondepth, planetradius):
86 endpoints = num.empty_like(latlondepth)
87 endpoints[:, 0], endpoints[:, 1] = od.ne_to_latlon(
88 latlondepth[:, 0], latlondepth[:, 1], ned[:, 0], ned[:, 1])
89 endpoints[:, 2] = latlondepth[:, 2] + ned[:, 2]
91 start_xyz = latlondepth2xyz(latlondepth, planetradius)
92 end_xyz = latlondepth2xyz(endpoints, planetradius)
94 return end_xyz - start_xyz
97def topo_to_vertices(lat, lon, ele, planetradius):
98 nlat = lat.size
99 nlon = lon.size
100 assert nlat > 1 and nlon > 1 and ele.shape == (nlat, nlon)
101 rtp = num.empty((ele.size, 3))
102 rtp[:, 0] = 1.0 + ele.flatten() / planetradius
103 rtp[:, 1] = (num.repeat(lat, nlon) + 90.) * d2r
104 rtp[:, 2] = num.tile(lon, nlat) * d2r
105 vertices = rtp2xyz(rtp)
106 return vertices
109g_topo_faces = {}
112def topo_to_faces(nlat, nlon):
113 k = (nlat, nlon)
114 if k not in g_topo_faces:
116 nfaces = 2 * (nlat - 1) * (nlon - 1)
117 faces = num.empty((nfaces, 3), dtype=num.int64)
118 ilon = num.arange(nlon - 1)
119 for ilat in range(nlat - 1):
120 ibeg = ilat*(nlon-1) * 2
121 iend = (ilat+1)*(nlon-1) * 2
122 faces[ibeg:iend:2, 0] = ilat*nlon + ilon
123 faces[ibeg:iend:2, 1] = ilat*nlon + ilon + 1
124 faces[ibeg:iend:2, 2] = ilat*nlon + ilon + nlon
125 faces[ibeg+1:iend+1:2, 0] = ilat*nlon + ilon + 1
126 faces[ibeg+1:iend+1:2, 1] = ilat*nlon + ilon + nlon + 1
127 faces[ibeg+1:iend+1:2, 2] = ilat*nlon + ilon + nlon
129 g_topo_faces[k] = faces
131 return g_topo_faces[k]
134g_topo_faces_quad = {}
137def topo_to_faces_quad(nlat, nlon):
138 k = (nlat, nlon)
139 if k not in g_topo_faces_quad:
141 nfaces = (nlat - 1) * (nlon - 1)
142 faces = num.empty((nfaces, 4), dtype=num.int64)
143 ilon = num.arange(nlon - 1)
144 for ilat in range(nlat - 1):
145 ibeg = ilat*(nlon-1)
146 iend = (ilat+1)*(nlon-1)
147 faces[ibeg:iend, 0] = ilat*nlon + ilon
148 faces[ibeg:iend, 1] = ilat*nlon + ilon + 1
149 faces[ibeg:iend, 2] = ilat*nlon + ilon + nlon + 1
150 faces[ibeg:iend, 3] = ilat*nlon + ilon + nlon
152 g_topo_faces_quad[k] = faces
154 return g_topo_faces_quad[k]
157def topo_to_mesh(lat, lon, ele, planetradius):
158 return \
159 topo_to_vertices(lat, lon, ele, planetradius), \
160 topo_to_faces_quad(*ele.shape)
163def refine_triangles(vertices, faces):
165 nv = vertices.shape[0]
166 nf = faces.shape[0]
167 assert faces.shape[1] == 3
169 fedges = num.concatenate((
170 faces[:, (0, 1)],
171 faces[:, (1, 2)],
172 faces[:, (2, 0)]))
174 fedges.sort(axis=1)
175 fedges1 = fedges[:, 0] + fedges[:, 1]*nv
176 sortorder = fedges1.argsort()
178 ne = fedges.shape[0] // 2
179 edges = fedges[sortorder[::2], :]
181 vertices_new = num.zeros((nv+ne, 3), dtype=num.float64)
182 vertices_new[:nv] = vertices
183 vertices_new[nv:] = 0.5 * (vertices[edges[:, 0]] + vertices[edges[:, 1]])
185 faces_new = num.zeros((nf*4, 3), dtype=num.int64)
187 vind = nv + sortorder.argsort()/2
189 faces_new[:nf, 0] = faces[:, 0]
190 faces_new[:nf, 1] = vind[:nf]
191 faces_new[:nf, 2] = vind[2*nf:]
193 faces_new[nf:nf*2, 0] = faces[:, 1]
194 faces_new[nf:nf*2, 1] = vind[nf:2*nf]
195 faces_new[nf:nf*2, 2] = vind[:nf]
197 faces_new[nf*2:nf*3, 0] = faces[:, 2]
198 faces_new[nf*2:nf*3, 1] = vind[2*nf:]
199 faces_new[nf*2:nf*3, 2] = vind[nf:2*nf]
201 faces_new[nf*3:, 0] = vind[:nf]
202 faces_new[nf*3:, 1] = vind[nf:2*nf]
203 faces_new[nf*3:, 2] = vind[2*nf:]
205 return vertices_new, faces_new
208def refine_triangle_parents(nfaces):
209 return num.arange(nfaces) % (nfaces / 4)
212def triangle_barycentric_transforms(vertices, faces):
213 a = vertices[faces[:, 1]] - vertices[faces[:, 0]]
214 b = vertices[faces[:, 2]] - vertices[faces[:, 0]]
216 norm = 1.0 / (vsqr(a) * vsqr(b) - vdot(a, b)**2)
218 transforms = num.zeros((faces.shape[0], 2, 3))
220 transforms[:, 0, :] = norm[:, num.newaxis] \
221 * (vsqr(b)[:, num.newaxis] * a - vdot(a, b)[:, num.newaxis] * b)
222 transforms[:, 1, :] = norm[:, num.newaxis] \
223 * (vsqr(a)[:, num.newaxis] * b - vdot(a, b)[:, num.newaxis] * a)
225 return transforms
228def triangle_barycentric_coordinates(transform, origin, points):
229 c = points - origin[num.newaxis, :]
230 uv = num.dot(transform[:, :], c.T).T
231 return uv
234def triangle_cartesian_system(vertices, faces):
235 a = normalize(vertices[faces[:, 1]] - vertices[faces[:, 0]])
236 b = vertices[faces[:, 2]] - vertices[faces[:, 1]]
237 c = normalize(num.cross(a, b))
238 uvw = num.zeros((faces.shape[0], 3, 3))
239 uvw[:, 0, :] = a
240 uvw[:, 1, :] = num.cross(c, a)
241 uvw[:, 2, :] = c
242 return uvw
245def triangle_planes(vertices, faces):
246 a = vertices[faces[:, 1]] - vertices[faces[:, 0]]
247 b = vertices[faces[:, 2]] - vertices[faces[:, 0]]
248 anormals = normalize(num.cross(a, b))
249 anormals *= vdot(vertices[faces[:, 0]], anormals)[:, num.newaxis]
250 return anormals