def test_get_picks(self):
"""
Should return data rows.
"""
# create a new database in memory
pickdb = PickDatabaseConnection(':memory:')
# add picks, and get our own times
times0 = []
for pick in uniq_picks:
pickdb.add_pick(**pick)
times0.append(pick['time'])
pickdb.commit()
# should return a list of sqlite3.Row objects
for row in pickdb.get_picks(event=pick['event']):
self.assertTrue(isinstance(row, sqlite3.Row))
# should return the same data
times1 = []
for pick in uniq_picks:
row = pickdb.get_picks(**pick)[0]
times1.append(row['time'])
self.assertEqual(sorted(times0), sorted(times1))
# should return all data
self.assertEqual(len(uniq_picks),
len(pickdb.get_picks()))
# should return empty list if no matches
self.assertEqual(len(pickdb.get_picks(event='golly_gee')), 0)
# should also return empty list if no data
pickdb = PickDatabaseConnection(':memory:')
self.assertEqual(len(pickdb.get_picks()), 0)
def test_counts(self):
"""
Count functions should return number of distinct rows in tables.
"""
pickdb = PickDatabaseConnection(':memory:')
# add a single pick
event = 'TestCount'
pick = uniq_picks[0]
pick['event'] = event
pickdb.add_pick(**pick)
pickdb.commit()
# should have a single pick for this event in picks table
n = pickdb.count_picks_by_event(event)
self.assertEqual(n, 1)
# should also have a single entry in the event table
n = pickdb._count_distinct(pickdb.EVENT_TABLE, event=event)
self.assertEqual(n, 1)
# now, another pick for the same event
pick['ensemble'] += 1 # ensure unique
pickdb.add_pick(**pick)
pickdb.commit()
# should have two picks for this event in the picks table
n = pickdb.count_picks_by_event(event)
self.assertEqual(n, 2)
# should still just have one entry in the event table for this event
n = pickdb._count_distinct(pickdb.EVENT_TABLE, event=event)
self.assertEqual(n, 1)
def test_copy(self):
"""
Should create a copy of the database.
"""
db0 = PickDatabaseConnection(':memory:')
for pick in uniq_picks:
db0.add_pick(**pick)
db1 = db0.copy()
self.assertEqual(db0.get_picks()[0], db1.get_picks()[0])
def test_get_events(self):
"""
Should return a list of event names in the database.
"""
pickdb = PickDatabaseConnection(':memory:')
# add picks and get our own list of event names
events = []
for i, pick in enumerate(uniq_picks):
pickdb.add_pick(**pick)
if pick['event'] not in events:
events.append(pick['event'])
# private function should return a list of events
self.assertEqual(sorted(pickdb._get_events()),
sorted(events))
# public 'events' attribute should be a property with
# setfunction=_get_events
self.assertEqual(pickdb.events,
pickdb._get_events())
# should return empty list for empty database
pickdb = PickDatabaseConnection(':memory:')
self.assertEqual(pickdb.events, [])
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)
from rockfish.vmtomo.raytracing import trace
cleanup = True
# Set filenames for model, pick database, etc.
pickdbfile = 'temp.pickdb.sqlite'
vmfile = '../tests/data/goc_l26.15.00.vm'
#vmfile = '../tests/data/cranis_northeast.00.00.vm'
input_dir = 'temp.input'
rayfile = 'temp.ray'
# Build an example pick database
pickdb = PickDatabaseConnection(pickdbfile)
pickdb.add_pick(event='Pg', ensemble=100, trace=1,
vm_branch=3, vm_subid=0,
time=5.0, time_reduced=5.0,
source_x=50, source_y=0.0, source_z=0.006,
receiver_x=100, receiver_y=0.0, receiver_z=2)
pickdb.commit()
# Raytrace the model
trace(vmfile, pickdb, rayfile, input_dir=input_dir, cleanup=False)
# Plot model and rays
fig = plt.figure()
ax = fig.add_subplot(311)
vm = VMFile(vmfile)
vm.plot(ax=ax)
ax.set_aspect(2)
plt.xlim((vm.r1[0], vm.r2[0]))
plt.ylim((vm.r2[2], vm.r1[2]))
def test_add_remove_picks(self):
"""
Should add a pick to the picks table.
"""
pickdb = PickDatabaseConnection(':memory:')
# should add all picks to the database
for pick in uniq_picks:
pickdb.add_pick(**pick)
pickdb.commit()
ndb = len(pickdb.execute('SELECT * FROM picks').fetchall())
self.assertEqual(ndb, len(uniq_picks))
# attempting to add pick without primary fields should raise an error
with self.assertRaises(sqlite3.IntegrityError):
pickdb.add_pick(event='Foobar', time=9.834)
# attempting to add duplicate picks should raise error
for pick in uniq_picks:
with self.assertRaises(sqlite3.IntegrityError):
pickdb.add_pick(**pick)
# directly removing pick and then re-adding should work
for pick in uniq_picks:
pickdb.remove_pick(**pick)
pickdb.add_pick(**pick)
pickdb.commit()
ndb = len(pickdb.execute('SELECT * FROM picks').fetchall())
self.assertEqual(ndb, len(uniq_picks))
# invalid collumn names should raise OperationalError
with self.assertRaises(sqlite3.OperationalError):
pickdb.remove_pick(not_a_field=999)
with self.assertRaises(sqlite3.OperationalError):
pickdb.add_pick(not_a_field=999)
# remove the last pick that we added
pickdb.remove_pick(**pick)
# attempting to remove a non-existant pick should do nothing
pickdb.remove_pick(**pick)
# updating picks should add picks if they don't exist and update picks
# if they do exist
for pick in uniq_picks:
pickdb.update_pick(**pick)
pickdb.commit()
ndb = len(pickdb.execute('SELECT * FROM picks').fetchall())
self.assertEqual(ndb, len(uniq_picks))
# should not be able to add pick without required fields
with self.assertRaises(sqlite3.IntegrityError):
pickdb.add_pick(event='Pg', ensemble=1, trace=1)
def locate_on_surface(vmfile, pickdb, iref, pick_keys={}, x0=None, y0=None,
dx=None, dy=None, plot=False):
"""
Locate a receiver on a surface.
:param vmfile: Filename of the VM Tomography slowness model to
raytrace.
:param pickdb: :class:`rockfish.picking.database.PickDatabaseConnection`
to get picks from for raytracing.
:param iref: index of the surface to move the receiver on
:param pick_keys: Optional. ``dict`` of keys and values to use for
selecting picks from ``pickdb``. Default is to use all picks.
:param x0, y0: Optional. Initial guess at x and y locations. Default is
center of the model.
:param dx, dy: Optional. Distance in x and y to search from ``x0, y0``.
Default is to search the entire model.
:param plot: Optional. Plot results of the location. Default is false.
"""
# Load model
print "Loading VM model..."
vm = readVM(vmfile)
# Setup search domain
print "Setting up search domain..."
if x0 is None:
x0 = np.mean(vm.x)
if y0 is None:
y0 = np.mean(vm.y)
if dx is None:
xsearch = vm.x
else:
ix = vm.xrange2i(max(vm.r1[0], x0 - dx), min(vm.r2[0], x0 + dx))
xsearch = vm.x[ix]
if dy is None:
ysearch = vm.y
else:
iy = vm.yrange2i(max(vm.r1[1], y0 - dy), min(vm.r2[1], y0 + dy))
ysearch = vm.y[iy]
# trace each point in the search region
print "Building traveltime database for {:}x{:} search grid..."\
.format(len(xsearch), len(ysearch))
zz = vm.rf[iref]
ipop = -1
population = []
db = PickDatabaseConnection(':memory:')
for x in xsearch:
for y in ysearch:
z = zz[vm.x2i([x])[0], vm.y2i([y])[0]]
ipop += 1
for _p in pickdb.get_picks(**pick_keys):
p = dict(_p)
p['ensemble'] = ipop
p['receiver_x'] = x
p['receiver_y'] = y
p['receiver_z'] = z
db.add_pick(**p)
population.append((x, y, z))
print "Raytracing..."
rayfile = '.locate.ray'
raytrace(vmfile, db, rayfile, stderr=subprocess.STDOUT,
stdout=subprocess.PIPE)
rays = readRayfanGroup(rayfile)
imin = np.argmin([rfn.rms for rfn in rays.rayfans])
ipop = [rfn.start_point_id for rfn in rays.rayfans][imin]
rms = [rfn.rms for rfn in rays.rayfans][imin]
# # warn if fit point is on edge of search domain
# if (vm.ny > 1) and ((yfit == ysearch[0]) or (yfit == ysearch[-1])):
# on_edge = True
# elif (xfit == xsearch[0]) or (xfit == xsearch[-1]):
# on_edge = True
# else:
# on_edge = False
# if on_edge:
# msg = 'Best-fit location ({:},{:}) is on edge of model domain:'\
# .format(xfit, yfit)
# msg += ' x=[{:},{:}], y=[{:},{:}].'.format(xsearch[0], xsearch[-1],
# ysearch[0], ysearch[-1])
# msg += ' Try using a larger search region.'
# warnings.warn(msg)
# plot
if plot:
fig = plt.figure()
ax = fig.add_subplot(111)
if vm.ny == 1:
ax.plot([p[0] for p in population],
[rfn.rms for rfn in rays.rayfans], '.k')
ax.plot(population[ipop][0], rms, '*r')
plt.xlabel('x (km)')
plt.ylabel('RMS error (s)')
else:
ax.plot([p[0] for p in population], [p[1] for p in population],
'.k')
ax.plot(population[ipop][0], population[ipop][1], '*r')
ax.plot(x0, y0, 'og')
plt.xlabel('x (km)')
plt.xlabel('y (km)')
plt.title('Results of locate_on_surface()')
plt.show()
return population[ipop][0], population[ipop][1], population[ipop][2], rms
