def test_get_example_file(self):
"""
get_example_file() should return the path of an example file
"""
import rockfish
# should return a filename
filename = 'mgl0807.p190'
path0 = os.path.join(rockfish.__path__[0], 'navigation',
'tests', 'data', filename)
self.assertEqual(path0, loaders.get_example_file(filename))
# should raise an ioerror if file does not exist
filename = 'does_not.exist'
with self.assertRaises(IOError):
loaders.get_example_file(filename)
def dev_grids(self):
"""
Should use a separate class to manage grids
"""
vm = readVM(get_example_file('1d.vm'))
# should have the grid class
self.assertTrue(hasattr(vm, 'grids'))
self.assertTrue(isinstance(vm.grids, VMGrids))
# should still have the 'sl' alias for the slowness grid
self.assertTrue(hasattr(vm, 'sl'))
sl0 = copy.copy(vm.sl) # save original slowness grid
# modifying vm.sl should update vm.grids.slowness
self.assertTrue(vm.sl is vm.grids.slowness)
vm.sl = 999.99
self.assertTrue(vm.grids.slowness == 999.99)
# and visa versa
vm.grids.slowness = 888.88
self.assertTrue(vm.sl == 888.88)
# grids should also have velocity and twt
self.assertTrue(hasattr(vm.grids, 'velocity'))
self.assertTrue(hasattr(vm.grids, 'twt'))
# these grids should be the same shape as sl
for n0, n1 in zip(vm.sl.shape, vm.grids.velocity.shape):
self.assertEqual(n0, n1)
for n0, n1 in zip(vm.sl.shape, vm.grids.twt.shape):
self.assertEqual(n0, n1)
def setUp(self):
"""
Setup benchmark data.
"""
self.segy = readSEGY(get_example_file("ew0210_o30.segy"), unpack_headers=True)
self.pickdb = PickDatabaseConnection(":memory:")
for pick in uniq_picks:
self.pickdb.update_pick(**pick)
def test_surface2netcdf(self):
"""
Should write a surface to a netcdf file.
"""
vm = readVM(get_example_file('cranis3d.vm'))
surface2netcdf(vm, 0, 'test.grd')
grd = netcdf('test.grd', 'r')
z = grd.variables['z'][:]
self.assertEqual((vm.ny, vm.nx), z.shape)
self.assertAlmostEqual(vm.rf[0][0][0] * 1000., z[0][0], 4)
def test_rayfan_init(self):
"""
Calling Rayfan(file) should read data from a file.
"""
# Example rayfan file
rayfile = get_example_file('goc_l26.14.00.ray')
# Open the file
file = open(rayfile, 'rb')
# Should read 1st rayfan from the file
rfn = Rayfan(file)
def test__get_layers(self):
"""
Should return the layer index of each node on the grid.
"""
for model in TEST_MODELS:
vm = readVM(get_example_file(model))
layers = vm._get_layers()
# should have an entry for each node
self.assertEqual(layers.shape, (vm.nx, vm.ny, vm.nz))
# should have nodes in each layer
# assuming no complete pinchouts
self.assertEqual(len(np.unique(layers)), vm.nr + 1)
def test_read(self):
"""
Should be able to read a P1/90 file
"""
dbfile = "test_p190.sqlite"
for fname, nsrc, nchan in P190_FILES:
_fname = get_example_file(fname)
p1 = P190(_fname, database=dbfile)
self.assertTrue(os.path.isfile(dbfile))
self.assertEqual(len(p1.source_points), nsrc)
self.assertEqual(len(p1.receiver_groups), nchan)
os.remove(dbfile)
def test_define_stretched_layer_velocities(self):
"""
Should fit a 1D velocity function to a layer.
"""
for model in TEST_MODELS:
vm = readVM(get_example_file(model))
# should define constant velocities within layer
for ilyr in range(vm.nr + 1):
vm.sl = np.nan * np.ones((vm.nx, vm.ny, vm.nz))
vm.define_stretched_layer_velocities(ilyr, [10])
self.assertEqual(np.nanmax(vm.sl), 1. / 10)
self.assertEqual(np.nanmin(vm.sl), 1. / 10)
def test_insert_layer_velocities(self):
"""
Should insert velocities into a layer.
"""
for model in TEST_MODELS:
vm = readVM(get_example_file(model))
# should define constant velocities within layer
for ilyr in range(vm.nr + 1):
vm.sl = np.nan * np.ones((vm.nx, vm.ny, vm.nz))
sl = 10 * np.ones((vm.nx, vm.ny, vm.nz))
vm.insert_layer_velocities(ilyr, sl, is_slowness=True)
self.assertEqual(np.nanmax(vm.sl), 10)
self.assertEqual(np.nanmin(vm.sl), 10)
def test_project_model_points(self):
"""
Should map 2d model points to points in the 3d model.
"""
phi = 30.
vm3d = VM(r1=(100, 200, 0), r2=(500, 300, 30), dx=5, dy=5, dz=1)
for model in TEST_MODELS:
vm2d = readVM(get_example_file(model))
# default should return valid coordinates
x = project_model_points(vm2d, vm3d, phi)
self.assertTrue(np.min(x) >= vm2d.r1[0])
self.assertTrue(np.max(x) <= vm2d.r2[0])
# indices = True should return valid indices
ix = project_model_points(vm2d, vm3d, phi)
self.assertTrue(np.min(ix) >= 0)
self.assertTrue(np.max(ix) <= vm2d.nx)
def setUp(self):
"""
Setup benchmark data.
"""
self.segy = readSEGY(get_example_file('ew0210_o30.segy'),
unpack_headers=True)
self.pickdb = PickDatabaseConnection(':memory:')
for pick in uniq_picks:
self.pickdb.update_pick(**pick)
self.default_params = [
'ABSCISSA_KEY', 'GAIN', 'CLIP',
'NORMALIZATION_METHOD', 'OFFSET_GAIN_POWER',
'WIGGLE_PEN_COLOR', 'WIGGLE_PEN_WIDTH',
'NEG_FILL_COLOR', 'POS_FILL_COLOR', 'DISTANCE_UNIT',
'TIME_UNIT', 'SEGY_TIME_UNITS', 'SEGY_DISTANCE_UNITS',
'SEGY_HEADER_ALIASES']
def test_raytrace_branch_id(self):
"""
Raytracing should honor branch ids
"""
#vmfile = get_example_file('jump1d.vm')
vmfile = get_example_file('inactive_layers.vm')
# Create pick database
pickdbfile = 'temp.sqlite'
if os.path.isfile(pickdbfile):
os.remove(pickdbfile)
pickdb = PickDatabaseConnection(pickdbfile)
pickdb.update_pick(event='P1', ensemble=100, trace=1,
branch=1, subbranch=0,
time=5.0, source_x=10, source_y=0.0, source_z=0.006,
receiver_x=40, receiver_y=0.0, receiver_z=4.9)
pickdb.update_pick(event='P2', ensemble=100, trace=1,
branch=2, subbranch=0, time=5.0)
pickdb.update_pick(event='P3', ensemble=100, trace=1,
branch=3, subbranch=0, time=5.0)
pickdb.commit()
# Raytrace
vm = readVM(vmfile)
rayfile = 'temp.ray'
for branch in range(1, 4):
if os.path.isfile(rayfile):
os.remove(rayfile)
pick_keys = {'branch' : branch}
raytrace(vmfile, pickdb, rayfile, **pick_keys)
# Should have created a rayfile
self.assertTrue(os.path.isfile(rayfile))
# Load rayfans
rays = readRayfanGroup(rayfile)
# Should have traced just one ray
self.assertEqual(len(rays.rayfans), 1)
rfn = rays.rayfans[0]
self.assertEqual(len(rfn.paths), 1)
# Rays should turn in specified layer
zmax = max([p[2] for p in rfn.paths[0]])
self.assertGreaterEqual(zmax, vm.rf[branch - 1][0][0])
# cleanup
for filename in [rayfile, pickdbfile]:
if os.path.isfile(filename):
os.remove(filename)
def test_read_write_vm(self):
"""
Reading and writing in the VM format should not change data.
"""
# read data from the disk file
vmfile = get_example_file(BENCHMARK_2D)
vm = readVM(vmfile)
# check values
self.compare_to_benchmark(vm)
# write it to another file
tmp = 'temp.vm'
vm.write(tmp)
# read these data back in
vm = readVM(tmp)
# check values
self.compare_to_benchmark(vm)
# clean up
os.remove(tmp)
def test_boundary_flags(self):
"""
Should convert boundary flags from fortran to python index conventions
"""
for model in TEST_MODELS:
# read anymodel with interfaces
vm = readVM(get_example_file(model))
# set all ir and ij flags to -1 (off)
vm.ir = -1 * np.ones(vm.ir.shape)
vm.ij = -1 * np.ones(vm.ij.shape)
# reading and writing should not change flag values
tempvm = 'temp123.vm'
vm.write(tempvm)
vm = readVM(tempvm)
self.assertEqual(vm.ir.min(), -1)
self.assertEqual(vm.ir.max(), -1)
self.assertEqual(vm.ij.min(), -1)
self.assertEqual(vm.ij.max(), -1)
# cleanup
os.remove(tempvm)
def test__write_csv(self):
"""
Should be able to write data to CSV files
"""
dbfile = "test_p190.sqlite"
for fname, nsrc, nchan in P190_FILES:
# read p190 file
_fname = get_example_file(fname)
p1 = P190(_fname)
# should write three files
fnames = p1.write("test_csv", output_format="csv")
tables = [p1.HEADER_TABLE, p1.COORDINATE_TABLE, p1.RECEIVER_TABLE]
for _table in tables:
self.assertTrue(os.path.isfile(fnames[_table]))
# should have a line for each source, receiver pair
nrec = line_count(fnames["receiver_groups"])
self.assertEqual(nrec - 1, nsrc * nchan)
# cleanup
for _table in tables:
os.remove(fnames[_table])
def test__write_p190(self):
"""
Should be able to write data to the P190 format
"""
dbfile = "test_p190.sqlite"
for fname, nsrc, nchan in P190_FILES:
# read p190 file
_fname = get_example_file(fname)
p1 = P190(_fname)
# should write one p190 file
ofile = "test.p190"
p1.write(ofile, output_format="p190")
self.assertTrue(os.path.isfile(ofile))
# read new p190
p1_new = P190(ofile)
# should have the same data
self.assertEqual(len(p1.receiver_groups), len(p1_new.receiver_groups))
self.assertEqual(len(p1.source_points), len(p1_new.source_points))
# cleanup
os.remove(ofile)
def test_locate_on_surface(self):
"""
Should locate a receiver on a surface.
"""
inst_id = 100
dx = 1
iref = 0
for _vmfile in TEST_MODELS:
vmfile = get_example_file(_vmfile)
vm = readVM(vmfile)
# calculate synthetic times
pickdb = PickDatabaseConnection(':memory:')
x0 = np.mean(vm.x)
y0 = np.mean(vm.y)
picks = []
xsearch = vm.xrange2i(max(vm.r1[0], x0 - dx),
min(vm.r2[0], x0 + dx))
for i, ix in enumerate(xsearch):
x = vm.x[ix]
iy = vm.x2i([y0])[0]
z0 = vm.rf[iref][ix][iy]
pickdb.add_pick(event='Pw', ensemble=inst_id,
trace=i, time=1e30,
source_x=x, source_y=y0, source_z=0.006,
receiver_x=x0, receiver_y=y0, receiver_z=z0,
vm_branch=1, vm_subid=0)
rayfile = 'temp.ray'
raytrace(vmfile, pickdb, rayfile)
raydb = rayfan2db(rayfile, 'temp.syn.sqlite', synthetic=True)
os.remove(rayfile)
# run locate
x, y, z, rms = locate_on_surface(vmfile, raydb, 0, x0=x0,
y0=y0, dx=dx, dy=dx)
# compare result
self.assertAlmostEqual(x, x0, 0)
self.assertAlmostEqual(y, y0, 0)
def test_parallel_raytrace(self):
"""
Should run raytracing in parallel
"""
# Create pick database
pickdbfile = 'temp.sqlite'
if os.path.isfile(pickdbfile):
os.remove(pickdbfile)
pickdb = PickDatabaseConnection(pickdbfile)
for i, event in enumerate( ['P1', 'P2', 'P3']):
branch = i + 1
for ens in range(3):
pickdb.update_pick(event=event, ensemble=ens, trace=1,
branch=branch, subbranch=0,
time=5.0, source_x=10, source_y=0.0,
source_z=0.006, receiver_x=40,
receiver_y=0.0, receiver_z=4.9)
pickdb.commit()
# set velocity model
vmfile = get_example_file('inactive_layers.vm')
# raytrace
for nproc in [1, 2, 8]:
input_dir = 'test.input'
output_dir = 'test.output'
t_start = time.clock()
parallel_raytrace(vmfile, pickdb, branches=[1, 2, 3],
input_dir=input_dir, output_dir=output_dir,
nproc=nproc, ensemble_field='ensemble')
t_elapsed = time.clock() - t_start
shutil.rmtree(input_dir)
shutil.rmtree(output_dir)
os.remove(pickdbfile)
def test_two2three(self, sol=(0, 0), theta=39):
"""
Should create a 3D model from a 2D model.
"""
for test_model in TEST_MODELS:
vm2d = readVM(get_example_file(test_model))
# calculate example endpoints
line_len = np.sqrt((vm2d.r2[0] - vm2d.r1[0]) ** 2 +\
(vm2d.r2[1] - vm2d.r1[1]) ** 2)
eol = (sol[0] + line_len * np.cos(np.deg2rad(theta)),
sol[1] + line_len * np.sin(np.deg2rad(theta)))
# should create 3d model
vm3d = two2three(vm2d, sol, eol, dx=1, dy=1)
# check bounds
self.assertAlmostEqual(vm3d.r1[0], sol[0], 0)
self.assertAlmostEqual(vm3d.r1[1], sol[1], 0)
self.assertAlmostEqual(vm3d.r2[0], eol[0], 0)
self.assertAlmostEqual(vm3d.r2[1], eol[1], 0)
# spot check some values
for iref in range(0, vm3d.nr):
self.assertAlmostEqual(vm3d.rf[iref].min(),
vm2d.rf[iref].min(), 1)
self.assertAlmostEqual(vm3d.rf[iref].max(),
vm2d.rf[iref].max(), 1)
