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