%matplotlib inline
from __future__ import division, print_function
import math
import numpy as np
import numpy.testing
import multiprocessing
import threading
import numba
from time import time
from fatiando.mesher import TesseroidMesh, Tesseroid
from fatiando import constants, utils, gridder
from fatiando.gravmag.tesseroid import gz as gz_cython
from fatiando.vis import mpl
def make_data():
model = TesseroidMesh((0, 5, 0, 10, 1000, -5000), (1, 5, 5))
model.addprop('density', 1000*np.ones(model.size))
x, y, z = gridder.regular((0, 5, 0, 10), (50, 50), z=50000)
return x, y, z, model
def test_gz(func, **kwargs):
print("Testing {}:".format(func.__name__))
x, y, z, model = make_data()
# Test the func against the Cython version in Fatiando
f1 = func(x, y, z, model, **kwargs)
f2 = gz_cython(x, y, z, model)
np.testing.assert_allclose(f1, f2, err_msg="Results don't match current implementation")
print(" - Results match current implementation")
# Time the execution
times = []
for i in xrange(5):
t1 = time()
func(x, y, z, model, **kwargs)
t2 = time()
times.append(t2 - t1)
print(" - Execution time: {:f} s".format(np.median(times)))
nodes = np.array([-0.577350269189625731058868041146,
0.577350269189625731058868041146])
stack_size = 100
def _check_input(lon, lat, height, ratio, out):
lon = np.radians(lon)
radius = height + constants.MEAN_EARTH_RADIUS
coslat = np.cos(np.radians(lat))
sinlat = np.sin(np.radians(lat))
if out is None:
result = np.zeros_like(lon)
else:
result = out
return lon, coslat, sinlat, radius, ratio, result
def _get_engine(engine):
if engine == 'numpy':
func = _numpy_engine
elif engine == 'numba':
func = _numba_engine
else:
raise ValueError("engine trouble")
return func
def _get_strides(size, njobs):
n = size//njobs
strides = [(i*n, (i + 1)*n) for i in xrange(njobs - 1)]
strides.append((strides[-1][-1], size))
return strides
def _split_arrays(arrays, args, njobs):
strides = _get_strides(arrays[0].size, njobs)
return [[x[low:high] for x in arrays] + args
for low, high in strides]
def _dispatcher(args):
lon, coslat, sinlat, radius, result, model, engine, ratio = args
func = _get_engine(engine)
for c in model:
density = c.props['density']
func(lon, coslat, sinlat, radius, c, density, ratio, result)
return result
def gz(lon, lat, height, model, ratio=1.6, engine=None, njobs=None, out=None):
lon, coslat, sinlat, radius, ratio, result = _check_input(lon, lat, height, ratio, out)
if njobs is None or njobs == 1:
_dispatcher([lon, coslat, sinlat, radius, result, model, engine, ratio])
else:
chunks = _split_arrays([lon, coslat, sinlat, radius, result], [model, engine, ratio], njobs)
pool = multiprocessing.Pool(njobs)
result = np.hstack(pool.map(_dispatcher, chunks))
pool.close()
result *= constants.G*constants.SI2MGAL
return result
def _numpy_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):
for l in xrange(lon.size):
stack = [cell]
while stack:
t = stack.pop()
distance, Llon, Llat, Lr = distance_size_numpy(lon[l], coslat[l], sinlat[l], radius[l], t)
nlon = 1 if distance/Llon > ratio else 2
nlat = 1 if distance/Llat > ratio else 2
nr = 1 if distance/Lr > ratio else 2
new_cells = nlon*nlat*nr
if new_cells > 1:
if new_cells + len(stack) > stack_size:
print(distance, Llon, Llat, Lr)
raise OverflowError('Tesseroid stack overflowed')
stack.extend(t.split(nlon, nlat, nr))
else:
lonc, sinlatc, coslatc, rc, scale = scale_nodes_numpy(t.get_bounds(), nodes)
result[l] += scale*density*kernelz_numpy(lon[l], coslat[l], sinlat[l], radius[l],
lonc, sinlatc, coslatc, rc)
def scale_nodes_numpy(bounds, nodes):
w, e, s, n, top, bottom = bounds
d2r = np.pi/180
dlon = d2r*(e - w)
dlat = d2r*(n - s)
dr = top - bottom
# Scale the GLQ nodes to the integration limits
lonc = 0.5*dlon*nodes + d2r*0.5*(e + w)
latc = 0.5*dlat*nodes + d2r*0.5*(n + s)
sinlatc = np.sin(latc)
coslatc = np.cos(latc)
rc = (0.5*dr*nodes +
0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))
scale = dlon*dlat*dr*0.125
return lonc, sinlatc, coslatc, rc, scale
def distance_size_numpy(lon, coslat, sinlat, radius, cell):
w, e, s, n, top, bottom = cell.get_bounds()
d2r = np.pi/180
rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS
lont = d2r*0.5*(w + e)
latt = d2r*0.5*(s + n)
sinlatt = np.sin(latt)
coslatt = np.cos(latt)
cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)
distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)
# Calculate the dimensions of the tesseroid in meters
rtop = top + constants.MEAN_EARTH_RADIUS
Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))
Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))
Lr = top - bottom
return distance, Llon, Llat, Lr
def kernelz_numpy(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):
r_sqr = radius**2
result = 0
for i in range(2):
coslon = np.cos(lon - lonc[i])
for j in range(2):
cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon
for k in range(2):
l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi
kappa = (rc[k]**2)*coslatc[j]
result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)
result *= -1
return result
@numba.jit(looplift=True, nogil=True)
def _numba_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):
bounds = np.array(cell.get_bounds())
stack = np.empty((stack_size, 6))
lonc = np.empty_like(nodes)
sinlatc = np.empty_like(nodes)
coslatc = np.empty_like(nodes)
rc = np.empty_like(nodes)
for l in range(result.size):
for i in range(6):
stack[0, i] = bounds[i]
stktop = 0
while stktop >= 0:
w, e, s, n, top, bottom = stack[stktop, :]
stktop -= 1
distance, Llon, Llat, Lr = distance_size_numba(lon[l], coslat[l], sinlat[l], radius[l],
w, e, s, n, top, bottom)
nlon, nlat, nr, new_cells = divisions_per_dimension(distance, Llon, Llat, Lr, ratio)
if new_cells > 1:
if new_cells + (stktop + 1) > stack_size:
raise OverflowError
stktop = split_numba(w, e, s, n, top, bottom, nlon, nlat, nr,
stack, stktop)
else:
scale = scale_nodes_numba(w, e, s, n, top, bottom, nodes,
lonc, sinlatc, coslatc, rc)
result[l] += density*scale*kernelz_numba(lon[l], coslat[l], sinlat[l], radius[l],
lonc, sinlatc, coslatc, rc)
@numba.jit(nopython=True, nogil=True)
def scale_nodes_numba(w, e, s, n, top, bottom, nodes, lonc, sinlatc, coslatc, rc):
d2r = np.pi/180
dlon = d2r*(e - w)
dlat = d2r*(n - s)
dr = top - bottom
# Scale the GLQ nodes to the integration limits
for i in range(len(nodes)):
lonc[i] = 0.5*dlon*nodes[i] + d2r*0.5*(e + w)
latc = 0.5*dlat*nodes[i] + d2r*0.5*(n + s)
sinlatc[i] = np.sin(latc)
coslatc[i] = np.cos(latc)
rc[i] = (0.5*dr*nodes[i] +
0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))
scale = dlon*dlat*dr*0.125
return scale
@numba.jit(nopython=True, nogil=True)
def distance_size_numba(lon, coslat, sinlat, radius, w, e, s, n, top, bottom):
d2r = np.pi/180
rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS
lont = d2r*0.5*(w + e)
latt = d2r*0.5*(s + n)
sinlatt = np.sin(latt)
coslatt = np.cos(latt)
cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)
distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)
# Calculate the dimensions of the tesseroid in meters
rtop = top + constants.MEAN_EARTH_RADIUS
Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))
Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))
Lr = top - bottom
return distance, Llon, Llat, Lr
@numba.jit(nopython=True, nogil=True)
def kernelz_numba(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):
r_sqr = radius**2
result = 0
for i in range(2):
coslon = np.cos(lon - lonc[i])
for j in range(2):
cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon
for k in range(2):
l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi
kappa = (rc[k]**2)*coslatc[j]
result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)
result *= -1
return result
@numba.jit(nopython=True, nogil=True)
def split_numba(w, e, s, n, top, bottom, nlon, nlat, nr, stack, stktop):
dlon = (e - w)/nlon
dlat = (n - s)/nlat
dr = (top - bottom)/nr
for i in xrange(nlon):
for j in xrange(nlat):
for k in xrange(nr):
stktop += 1
stack[stktop, 0] = w + i*dlon
stack[stktop, 1] = w + (i + 1)*dlon
stack[stktop, 2] = s + j*dlat
stack[stktop, 3] = s + (j + 1)*dlat
stack[stktop, 4] = bottom + (k + 1)*dr
stack[stktop, 5] = bottom + k*dr
return stktop
@numba.jit(nopython=True, nogil=True)
def divisions_per_dimension(distance, Llon, Llat, Lr, ratio):
nlon = 1 if distance/Llon > ratio else 2
nlat = 1 if distance/Llat > ratio else 2
nr = 1 if distance/Lr > ratio else 2
return nlon, nlat, nr, nlon*nlat*nr
test_gz(gz_cython)
Testing gz: - Results match current implementation - Execution time: 0.259152 s
test_gz(gz, engine='numba')
Testing gz: - Results match current implementation - Execution time: 0.212037 s
test_gz(gz, engine='numpy')
Testing gz: - Results match current implementation - Execution time: 13.163021 s