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 rayfan2db(rayfan_file, raydb_file=':memory:', synthetic=False, noise=None,
pickdb=None, raypaths=False):
"""
Read a rayfan file and store its data in a database.
Data are stored in a modified version of a
:class:`rockfish.picking.database.PickDatabaseConnection`.
Parameters
----------
rayfan_file: {str, file}
An open file-like object or a string which is assumed to be a
filename of a rayfan binary file.
raydb_file: str, optional
The filename of the new database. Default is to create
a new database in memory.
synthetic: bool, optional
Determines whether or not to record traced traveltimes as picked
traveltimes.
noise: {float, None}
Maximum amplitude of random noise to add the travel times. If
``None``, no noise is added.
pickdb: :class:`rockfish.database.PickDatabaseConnection`, optional
An active connection to the pick database that was used
trace rays in ``rayfan_file``. Values for extra fields (e.g.,
'trace_in_file') are copied from this database to the new
database along with rayfan data. Default is ignore these extra fields.
raypaths: bool, optional
If ``True``, raypath coordinates are stored as text in a new table
'raypaths'.
"""
raydb = PickDatabaseConnection(raydb_file)
rays = readRayfanGroup(rayfan_file)
print "Adding {:} traveltimes to {:} ..."\
.format(rays.nrays, raydb_file)
# add fields for raypaths
if raypaths:
raydb._create_table_if_not_exists(RAYPATH_TABLE, RAYPATH_FIELDS)
ndb0 = raydb.execute('SELECT COUNT(rowid) FROM picks').fetchone()[0]
for rfn in rays.rayfans:
for i, _t in enumerate(rfn.travel_times):
if noise is not None:
_noise = noise * 2 * (np.random.random() - 0.5)
else:
_noise = 0.0
sx, sy, sz = rfn.paths[i][0]
rx, ry, rz = rfn.paths[i][-1]
if pickdb is not None:
event = pickdb.vmbranch2event[rfn.event_ids[i]]
else:
event = rfn.event_ids[i]
if synthetic:
time = rfn.travel_times[i] + _noise
time_reduced = time
predicted = None
residual = 0.
else:
time = rfn.pick_times[i]
time_reduced = time
predicted = rfn.travel_times[i]
residual = rfn.residuals[i]
d = {'event': event,
'ensemble': rfn.start_point_id,
'trace': rfn.end_point_ids[i],
'vm_branch': rfn.event_ids[i],
'vm_subid': rfn.event_subids[i],
'time' : time,
'time_reduced' : time_reduced,
'predicted' : predicted,
'residual' : residual,
'error': rfn.pick_errors[i],
'source_x': sx, 'source_y': sy, 'source_z': sz,
'receiver_x': rx, 'receiver_y': ry, 'receiver_z': rz,
'offset': rfn.offsets[i],
'faz': rfn.azimuths[i],
'method': 'rayfan2db({:})'.format(rayfan_file),
'data_file': rays.file.name}
# Copy data from pickdb
if pickdb is not None:
pick = pickdb.get_picks(event=[event],
ensemble=[d['ensemble']],
trace=[d['trace']])
if len(pick) > 0:
for f in ['trace_in_file', 'line', 'site', 'data_file']:
try:
d[f] = pick[0][f]
except KeyError:
pass
# add data to standard tables
raydb.update_pick(**d)
# add raypath data to new table
if raypaths:
d = {'event': event,
'ensemble': rfn.start_point_id,
'trace': rfn.end_point_ids[i],
'ray_btm_x': rays.bottom_points[i][0],
'ray_btm_y': rays.bottom_points[i][1],
'ray_btm_z': rays.bottom_points[i][2],
'ray_x': str([p[0] for p in rfn.paths[i]]),
'ray_y': str([p[1] for p in rfn.paths[i]]),
'ray_z': str([p[2] for p in rfn.paths[i]])}
raydb._insert(RAYPATH_TABLE, **d)
raydb.commit()
ndb = raydb.execute('SELECT COUNT(rowid) FROM picks').fetchone()[0]
if (ndb - ndb0) != rays.nrays:
msg = 'Only added {:} of {:} travel times to the database.'\
.format(ndb - ndb0, rays.nrays)
warnings.warn(msg)
return raydb
