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