Skip to content

Commit e320bdb

Browse files
author
Malcolm White
committed
UPGRADE:: Implement ray-tracing in spherical coordinates
1 parent a443acb commit e320bdb

File tree

1 file changed

+221
-46
lines changed

1 file changed

+221
-46
lines changed

‎pykonal/pykonal.pyx‎

Lines changed: 221 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,8 @@ class EikonalSolver(object):
189189
self._update()
190190

191191

192-
def trace_ray(self, *args, method='euler', tolerance=1e-2):
193-
if method.upper() == 'EULER':
194-
return (self._trace_ray_euler(*args, tolerance=tolerance))
195-
else:
196-
raise (NotImplementedError('Only Euler integration is implemented yet'))
192+
#def trace_ray(self, *args, step_size=None):
193+
#return (self._trace_ray_euler_spherical(*args))
197194

198195

199196
def transfer_travel_times_from(self, old, origin, rotate=False, set_alive=False):
@@ -238,10 +235,137 @@ class EikonalSolver(object):
238235
vvi = return_nan_on_error(LinearInterpolator3D(old.vgrid, old.vv))
239236
self.vv = np.apply_along_axis(vvi, -1, vgrid_new)
240237

238+
def _get_gradient(self):
239+
if self.coord_sys == 'cartesian':
240+
gg = np.moveaxis(
241+
np.stack(
242+
np.gradient(
243+
self.uu,
244+
*[
245+
self.pgrid.node_intervals[iax]
246+
for iax in range(self.ndim) if iax not in self.iax_null
247+
],
248+
axis=[
249+
iax
250+
for iax in range(self.ndim)
251+
if iax not in self.iax_null
252+
]
253+
)
254+
),
255+
0, -1
256+
)
257+
for iax in self.iax_null:
258+
gg = np.insert(gg, iax, np.zeros(self.pgrid.npts), axis=-1)
259+
else:
260+
grid = self.pgrid[...]
261+
d0, d1, d2 = self.pgrid.node_intervals
262+
n0, n1, n2 = self.pgrid.npts
263+
264+
if 0 not in self.iax_null:
265+
g0 = np.concatenate(
266+
[
267+
# Second-order forward difference evaluated along the lower edge
268+
(
269+
(
270+
self.uu[2]
271+
- 4*self.uu[1]
272+
+ 3*self.uu[0]
273+
) / (2*d0)
274+
).reshape(1, n1, n2),
275+
# Second order central difference evaluated in the interior
276+
(self.uu[2:] - self.uu[:-2]) / (2*d0),
277+
# Second order backward difference evaluated along the upper edge
278+
(
279+
(
280+
self.uu[-3]
281+
- 4*self.uu[-2]
282+
+ 3*self.uu[-1]
283+
) / (2*d0)
284+
).reshape(1, n1, n2)
285+
],
286+
axis=0
287+
)
288+
else:
289+
g0 = np.zeros((n0, n1, n2))
241290

242-
def _trace_ray_euler(self, start):
291+
if 1 not in self.iax_null:
292+
g1 = np.concatenate(
293+
[
294+
# Second-order forward difference evaluated along the lower edge
295+
(
296+
(
297+
self.uu[:,2]
298+
- 4*self.uu[:,1]
299+
+ 3*self.uu[:,0]
300+
) / (2*grid[:,0,:,0]*d1)
301+
).reshape(n0, 1, n2),
302+
# Second order central difference evaluated in the interior
303+
(
304+
self.uu[:,2:]
305+
- self.uu[:,:-2]
306+
) / (2*grid[:,1:-1,:,0]*d1),
307+
# Second order backward difference evaluated along the upper edge
308+
(
309+
(
310+
self.uu[:,-3]
311+
- 4*self.uu[:,-2]
312+
+ 3*self.uu[:,-1]
313+
) / (2*grid[:,-1,:,0]*d1)
314+
).reshape(n0, 1, n2)
315+
],
316+
axis=1
317+
)
318+
else:
319+
g1 = np.zeros((n0, n1, n2))
320+
321+
if 2 not in self.iax_null:
322+
g2 = np.concatenate(
323+
[
324+
# Second-order forward difference evaluated along the lower edge
325+
(
326+
(
327+
self.uu[:,:,2]
328+
- 4*self.uu[:,:,1]
329+
+ 3*self.uu[:,:,0]
330+
) / (2*grid[:,:,0,0]*np.sin(grid[:,:,0,1])*d2)
331+
).reshape(n0, n1, 1),
332+
333+
# Second order central difference evaluated in the interior
334+
(
335+
self.uu[:,:,2:]
336+
- self.uu[:,:,:-2]
337+
) / (2*grid[:,:,1:-1,0]*np.sin(grid[:,:,1:-1,1])*d2),
338+
# Second order backward difference evaluated along the upper edge
339+
(
340+
(
341+
self.uu[:,:,-3]
342+
- 4*self.uu[:,:,-2]
343+
+ 3*self.uu[:,:,-1]
344+
) / (2*grid[:,:,-1,0]*np.sin(grid[:,:,-1,1])*d2)
345+
).reshape(n0, n1, 1)
346+
],
347+
axis=2
348+
)
349+
else:
350+
g2 = np.zeros((n0, n1, n2))
351+
return (np.moveaxis(np.stack([g0, g1, g2]), 0, -1))
352+
353+
def _get_step_size(self):
354+
# This is will not work in spherical coordinates when rho0 = 0.
355+
return(
356+
self.norm[
357+
tuple(
358+
slice(self.norm.shape[iax])
359+
if iax not in self.iax_null
360+
else 0
361+
for iax in range(self.ndim)
362+
)
363+
].min() / 2
364+
)
365+
366+
def trace_ray(self, start, step_size=None):
243367
cdef cpp_vector[_REAL_t *] ray
244-
cdef _REAL_t step_size, gx, gy, gz, norm
368+
cdef _REAL_t g0, g1, g2, norm
245369
cdef _REAL_t *point_new
246370
cdef _REAL_t[3] point_last, point_2last
247371
cdef Py_ssize_t i
@@ -251,50 +375,30 @@ class EikonalSolver(object):
251375
point_new[0], point_new[1], point_new[2] = start
252376
ray.push_back(point_new)
253377
# step_size <-- half the smallest node_interval
254-
step_size = np.min(
255-
[
256-
self.pgrid.node_intervals[iax]
257-
for iax in range(self.ndim) if iax not in self.iax_null
258-
]
259-
) / 2
260-
# Create an interpolator for the gradient field
261-
gg = np.moveaxis(
262-
np.stack(
263-
np.gradient(
264-
self.uu,
265-
*[
266-
self.pgrid.node_intervals[iax]
267-
for iax in range(self.ndim) if iax not in self.iax_null
268-
],
269-
axis=[
270-
iax
271-
for iax in range(self.ndim)
272-
if iax not in self.iax_null
273-
]
274-
)
275-
),
276-
0, -1
277-
)
278-
for iax in self.iax_null:
279-
gg = np.insert(gg, iax, np.zeros(self.pgrid.npts), axis=-1)
280-
grad_x = LinearInterpolator3D(self.pgrid, gg[...,0].astype(DTYPE_REAL))
281-
grad_y = LinearInterpolator3D(self.pgrid, gg[...,1].astype(DTYPE_REAL))
282-
grad_z = LinearInterpolator3D(self.pgrid, gg[...,2].astype(DTYPE_REAL))
378+
if step_size is None:
379+
step_size = self._get_step_size()
380+
gg = self._get_gradient()
381+
grad_0 = LinearInterpolator3D(self.pgrid, gg[...,0].astype(DTYPE_REAL))
382+
grad_1 = LinearInterpolator3D(self.pgrid, gg[...,1].astype(DTYPE_REAL))
383+
grad_2 = LinearInterpolator3D(self.pgrid, gg[...,2].astype(DTYPE_REAL))
283384
# Create an interpolator for the travel-time field
284385
uu = LinearInterpolator3D(self.pgrid, self.uu)
285386
point_last = ray.back()
286387
while True:
287-
gx = grad_x.interpolate(point_last)
288-
gy = grad_y.interpolate(point_last)
289-
gz = grad_z.interpolate(point_last)
290-
norm = libc.math.sqrt(gx**2 + gy**2 + gz**2)
291-
gx /= norm
292-
gy /= norm
293-
gz /= norm
388+
g0 = grad_0.interpolate(point_last)
389+
g1 = grad_1.interpolate(point_last)
390+
g2 = grad_2.interpolate(point_last)
391+
norm = libc.math.sqrt(g0**2 + g1**2 + g2**2)
392+
g0 /= norm
393+
g1 /= norm
394+
g2 /= norm
395+
if self.coord_sys == 'spherical':
396+
g1 /= point_last[0]
397+
g2 /= point_last[0] * np.sin(point_last[1])
294398
point_new = <_REAL_t *> malloc(3 * sizeof(_REAL_t))
295-
point_new[0] = point_last[0] - step_size * gx
296-
point_new[1] = point_last[1] - step_size * gy
297-
point_new[2] = point_last[2] - step_size * gz
399+
point_new[0] = point_last[0] - step_size * g0
400+
point_new[1] = point_last[1] - step_size * g1
401+
point_new[2] = point_last[2] - step_size * g2
298402
point_2last = ray.back()
299403
ray.push_back(point_new)
300404
point_last = ray.back()
@@ -309,6 +413,76 @@ class EikonalSolver(object):
309413
return (ray_np)
310414

311415

416+
#def _trace_ray_euler_cartesian(self, start):
417+
# cdef cpp_vector[_REAL_t *] ray
418+
# cdef _REAL_t step_size, gx, gy, gz, norm
419+
# cdef _REAL_t *point_new
420+
# cdef _REAL_t[3] point_last, point_2last
421+
# cdef Py_ssize_t i
422+
# cdef np.ndarray[_REAL_t, ndim=2] ray_np
423+
424+
# point_new = <_REAL_t *> malloc(3 * sizeof(_REAL_t))
425+
# point_new[0], point_new[1], point_new[2] = start
426+
# ray.push_back(point_new)
427+
# # step_size <-- half the smallest node_interval
428+
# step_size = np.min(
429+
# [
430+
# self.pgrid.node_intervals[iax]
431+
# for iax in range(self.ndim) if iax not in self.iax_null
432+
# ]
433+
# ) / 2
434+
# # Create an interpolator for the gradient field
435+
# gg = np.moveaxis(
436+
# np.stack(
437+
# np.gradient(
438+
# self.uu,
439+
# *[
440+
# self.pgrid.node_intervals[iax]
441+
# for iax in range(self.ndim) if iax not in self.iax_null
442+
# ],
443+
# axis=[
444+
# iax
445+
# for iax in range(self.ndim)
446+
# if iax not in self.iax_null
447+
# ]
448+
# )
449+
# ),
450+
# 0, -1
451+
# )
452+
# for iax in self.iax_null:
453+
# gg = np.insert(gg, iax, np.zeros(self.pgrid.npts), axis=-1)
454+
# grad_x = LinearInterpolator3D(self.pgrid, gg[...,0].astype(DTYPE_REAL))
455+
# grad_y = LinearInterpolator3D(self.pgrid, gg[...,1].astype(DTYPE_REAL))
456+
# grad_z = LinearInterpolator3D(self.pgrid, gg[...,2].astype(DTYPE_REAL))
457+
# # Create an interpolator for the travel-time field
458+
# uu = LinearInterpolator3D(self.pgrid, self.uu)
459+
# point_last = ray.back()
460+
# while True:
461+
# gx = grad_x.interpolate(point_last)
462+
# gy = grad_y.interpolate(point_last)
463+
# gz = grad_z.interpolate(point_last)
464+
# norm = libc.math.sqrt(gx**2 + gy**2 + gz**2)
465+
# gx /= norm
466+
# gy /= norm
467+
# gz /= norm
468+
# point_new = <_REAL_t *> malloc(3 * sizeof(_REAL_t))
469+
# point_new[0] = point_last[0] - step_size * gx
470+
# point_new[1] = point_last[1] - step_size * gy
471+
# point_new[2] = point_last[2] - step_size * gz
472+
# point_2last = ray.back()
473+
# ray.push_back(point_new)
474+
# point_last = ray.back()
475+
# if uu.interpolate(point_2last) <= uu.interpolate(point_last):
476+
# break
477+
# ray_np = np.zeros((ray.size()-1, 3), dtype=DTYPE_REAL)
478+
# for i in range(ray.size()-1):
479+
# ray_np[i, 0] = ray[i][0]
480+
# ray_np[i, 1] = ray[i][1]
481+
# ray_np[i, 2] = ray[i][2]
482+
# free(ray[i])
483+
# return (ray_np)
484+
485+
312486
def _update(self):
313487
'''
314488
Update travel-time grid.
@@ -708,6 +882,7 @@ cdef class LinearInterpolator3D(object):
708882
)
709883
)
710884
idx[iax] = (point[iax] - self._min_coords[iax]) / self._node_intervals[iax]
885+
# TODO:: Accounting for node_interval is likely uneccessary here.
711886
delta[iax] = (idx[iax] % 1.) * self._node_intervals[iax]
712887
i1 = <Py_ssize_t> idx[0]
713888
i2 = <Py_ssize_t> idx[1]

0 commit comments

Comments
 (0)