def test_rotation_error(ned_rotation_error, expected_receiver_fn, tmpdir_factory): # Test simple rotation error ned_duplicate = copy.deepcopy(ned_rotation_error) result = analyze_station_orientations(ned_rotation_error, SYNTH_CURATION_OPTS, DEFAULT_CONFIG_FILTERING, DEFAULT_CONFIG_PROCESSING) print(ned_rotation_error.param, result) assert EXPECTED_SYNTH_KEY in result assert result[EXPECTED_SYNTH_KEY][ 'azimuth_correction'] == -ned_rotation_error.param # Apply correction and check that we can recover proper receiver function ned_duplicate.apply(lambda stream: correct_back_azimuth( None, stream, baz_correction=result[EXPECTED_SYNTH_KEY]['azimuth_correction'])) stacked_rf_R, rfs = compute_ned_stacked_rf(ned_duplicate) actual_plot_filename = 'test_analyze_orientations_actual_{}.png'.format( ned_rotation_error.param) actual_outfile = os.path.join( tmpdir_factory.mktemp('actual_receiver_fn').strpath, actual_plot_filename) plot_rf_stack(rfs, save_file=actual_outfile) print('Actual RF plot saved to tmp file {}'.format(actual_outfile)) expected_data = expected_receiver_fn[0].data actual_data = stacked_rf_R[0].data delta = np.abs(actual_data - expected_data) rmse = np.sqrt(np.mean(np.square(delta))) rms = np.sqrt(np.mean(np.square(expected_data))) print('Rotated RMSE:', rmse, ', RMSE rel:', rmse / rms) assert np.allclose(actual_data, expected_data, rtol=1.0e-3, atol=1.0e-5)
def expected_receiver_fn(_master_event_dataset, tmpdir_factory): ned = copy.deepcopy(_master_event_dataset) baseline_plot_filename = 'test_analyze_orientations_baseline.png' baseline_outfile = os.path.join( tmpdir_factory.mktemp('expected_receiver_fn').strpath, baseline_plot_filename) baseline_rf, details = compute_ned_stacked_rf(ned) plot_rf_stack(details, save_file=baseline_outfile) print('Reference RF plot saved to tmp file {}'.format(baseline_outfile)) plt.close() return baseline_rf
def test_channel_negated(ned_channel_negated, expected_receiver_fn, tmpdir_factory): # Test negation of one or both transverse channels ned_duplicate = copy.deepcopy(ned_channel_negated) plot_folder = tmpdir_factory.mktemp('test_channel_negated_{}'.format( ned_channel_negated.param)).strpath result = analyze_station_orientations(ned_channel_negated, SYNTH_CURATION_OPTS, DEFAULT_CONFIG_FILTERING, DEFAULT_CONFIG_PROCESSING, save_plots_path=plot_folder) print('Negated Analysis plots saved to', plot_folder) print('Negated {}'.format(ned_channel_negated.param), result) assert EXPECTED_SYNTH_KEY in result assert result[EXPECTED_SYNTH_KEY]['azimuth_correction'] != 0.0 # Even though the induced error here is not a rotation, # prospectively apply correction and check how it relates to the # proper, unperturbed receiver function. # This test demonstrates that a channel negation can be repaired by # treating it as if it were a rotation error. ned_duplicate.apply(lambda stream: correct_back_azimuth( None, stream, baz_correction=result[EXPECTED_SYNTH_KEY]['azimuth_correction'])) stacked_rf_R, rfs = compute_ned_stacked_rf(ned_duplicate) actual_plot_filename = 'test_analyze_orientations_actual_{}.png'.format( ned_channel_negated.param) actual_outfile = os.path.join( tmpdir_factory.mktemp('actual_receiver_fn').strpath, actual_plot_filename) plot_rf_stack(rfs, save_file=actual_outfile) print('Actual RF plot saved to tmp file {}'.format(actual_outfile)) expected_data = expected_receiver_fn[0].data actual_data = stacked_rf_R[0].data delta = np.abs(actual_data - expected_data) rmse = np.sqrt(np.mean(np.square(delta))) rms = np.sqrt(np.mean(np.square(expected_data))) rmse_rel = rmse / rms print('Negated RMSE:', rmse, ', RMSE rel:', rmse_rel) assert rmse_rel < 0.05
def main(): output_folder = 'moho_discontinuity' if not os.path.exists(output_folder): os.mkdir(output_folder) # end if # Generate data files corresponding to two different Moho depths V_s = 3.8 # km/s V_p = 6.4 # km/s H_0 = 40 # km k = V_p / V_s log.info("H_0 = {:.3f}".format(H_0)) log.info("k = {:.3f}".format(k)) # Loop over valid range of inclinations and generate synthetic RFs num_traces = 15 inclinations = np.linspace(14.0, 27.0, num_traces) distances = convert_inclination_to_distance(inclinations) final_sample_rate_hz = 10.0 # stream_0 = synthesize_rf_dataset(H_0, V_p, V_s, inclinations, distances, final_sample_rate_hz, log, baz=0) stream_0, _ = synthesize_rf_dataset( H_0, V_p, V_s, inclinations, distances, final_sample_rate_hz, log, # amplitudes=[1, 0.5, 0.2], baz=0) amplitudes=[1, 0.4, 0.3], baz=0) stream_0.write(os.path.join(output_folder, "synth_rf_data_H={}.h5".format(H_0)), format='h5') H_1 = 50 # stream_1 = synthesize_rf_dataset(H_1, V_p, V_s, inclinations, distances, final_sample_rate_hz, log, baz=90) stream_1, _ = synthesize_rf_dataset( H_1, V_p, V_s, inclinations, distances, final_sample_rate_hz, log, # amplitudes=[1, 0.4, 0.3], baz=90) amplitudes=[1, 0.5, 0.2], baz=90) stream_1.write(os.path.join(output_folder, "synth_rf_data_H={}.h5".format(H_1)), format='h5') # Plot each dataset separately, then aggregated together rf_fig_0 = rf_plot_utils.plot_rf_stack(stream_0, time_window=(-5, 30)) plt.title("Exact H = {:.3f}, k = {:.3f}".format(H_0, k)) rf_fig_0.savefig(os.path.join(output_folder, "synth_rf_H={}.png".format(H_0)), dpi=300) rf_fig_1 = rf_plot_utils.plot_rf_stack(stream_1, time_window=(-5, 30)) plt.title("Exact H = {:.3f}, k = {:.3f}".format(H_1, k)) rf_fig_1.savefig(os.path.join(output_folder, "synth_rf_H={}.png".format(H_1)), dpi=300) stream_all = stream_0 + stream_1 rf_fig_all = rf_plot_utils.plot_rf_stack(stream_all, time_window=(-5, 30)) plt.title("Exact H = {:.3f}+{:.3f}, k = {:.3f}".format(H_0, H_1, k)) rf_fig_all.savefig(os.path.join(output_folder, "synth_rf_H={}+{}.png".format(H_0, H_1)), dpi=300) # Run H-k stacking on synthetic data station_db = {'HHR': stream_all} k_grid, h_grid, hk = rf_stacking.compute_hk_stack(station_db['HHR'], include_t3=False, root_order=2) w = (0.5, 0.5) w_str = "w={:1.2f},{:1.2f}".format(*w) stack = rf_stacking.compute_weighted_stack(hk, weighting=w) hk_fig = rf_plot_utils.plot_hk_stack( k_grid, h_grid, stack, num=len(stream_all), title="Synthetic H-k stack ({})".format(w_str)) maxima = rf_stacking.find_local_hk_maxima(k_grid, h_grid, stack, min_rel_value=0.9) log.info("Numerical solutions:{}".format(maxima)) plt.text(k, H_0 + 1.5, "Solution $H_0$ = {:.3f}, k = {:.3f}".format(H_0, k), color="#ffffff", fontsize=16, horizontalalignment='left') plt.text(k, H_1 + 1.5, "Solution $H_1$ = {:.3f}, k = {:.3f}".format(H_1, k), color="#ffffff", fontsize=16, horizontalalignment='left') for h_max, k_max, _, _, _ in maxima: plt.scatter(k_max, h_max, marker='+', c="#000000", s=20) # end for # Save result hk_fig.savefig(os.path.join(output_folder, "synth_stack_{}.png".format(w_str)), dpi=300)
def main(input_file, output_file, event_mask_folder='', apply_amplitude_filter=False, apply_similarity_filter=False, hk_weights=DEFAULT_HK_WEIGHTS): # Read source file data_all = rf_util.read_h5_rf(input_file) # Convert to hierarchical dictionary format data_dict = rf_util.rf_to_dict(data_all) event_mask_dict = None if event_mask_folder and os.path.isdir(event_mask_folder): mask_files = os.listdir(event_mask_folder) mask_files = [ f for f in mask_files if os.path.isfile(os.path.join(event_mask_folder, f)) ] # print(mask_files) pattern = r"([A-Za-z0-9\.]{5,})_event_mask\.txt" pattern = re.compile(pattern) event_mask_dict = dict() for f in mask_files: match_result = pattern.match(f) if not match_result: continue code = match_result[1] # print(code) with open(os.path.join(event_mask_folder, f), 'r') as f: events = f.readlines() events = set([e.strip() for e in events]) event_mask_dict[code] = events # end with # end for # end if if event_mask_dict: print("Loaded {} event masks".format(len(event_mask_dict))) # end if # Plot all data to PDF file fixed_stack_height_inches = 0.8 y_pad_inches = 1.6 total_trace_height_inches = paper_size_A4[ 1] - fixed_stack_height_inches - y_pad_inches max_trace_height = 0.2 with PdfPages(output_file) as pdf: # Would like to use Tex, but lack desktop PC privileges to update packages to what is required plt.rc('text', usetex=False) pbar = tqdm.tqdm(total=len(data_dict)) network = data_dict.network rf_type = data_dict.rotation for st in sorted(data_dict.keys()): station_db = data_dict[st] pbar.update() pbar.set_description("{}.{}".format(network, st)) # Choose RF channel channel = rf_util.choose_rf_source_channel(rf_type, station_db) channel_data = station_db[channel] full_code = '.'.join([network, st, channel]) t_channel = list(channel) t_channel[-1] = 'T' t_channel = ''.join(t_channel) rf_stream = rf.RFStream(channel_data).sort(['back_azimuth']) if event_mask_dict and full_code in event_mask_dict: # Select events from external source event_mask = event_mask_dict[full_code] rf_stream = rf.RFStream([ tr for tr in rf_stream if tr.stats.event_id in event_mask ]).sort(['back_azimuth']) # end if if apply_amplitude_filter: # Label and filter quality rf_util.label_rf_quality_simple_amplitude(rf_type, rf_stream) rf_stream = rf.RFStream([ tr for tr in rf_stream if tr.stats.predicted_quality == 'a' ]).sort(['back_azimuth']) # end if if apply_similarity_filter: rf_stream = rf_util.filter_crosscorr_coeff(rf_stream) # end if if not rf_stream: continue # Find matching T-component data events = [tr.stats.event_id for tr in rf_stream] transverse_data = station_db[t_channel] t_stream = rf.RFStream([ tr for tr in transverse_data if tr.stats.event_id in events ]).sort(['back_azimuth']) if not t_stream: continue # Plot pinwheel of primary and transverse components fig = rf_plot_utils.plot_rf_wheel([rf_stream, t_stream], fontscaling=0.8) fig.set_size_inches(*paper_size_A4) plt.tight_layout() plt.subplots_adjust(hspace=0.15, top=0.95, bottom=0.15) ax = fig.gca() fig.text(-0.32, -0.32, "\n".join(rf_stream[0].stats.processing), fontsize=6, transform=ax.transAxes) pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() num_traces = len(rf_stream) assert len(t_stream) == num_traces # Plot RF stack of primary component trace_ht = min(total_trace_height_inches / num_traces, max_trace_height) fig = rf_plot_utils.plot_rf_stack( rf_stream, trace_height=trace_ht, stack_height=fixed_stack_height_inches, fig_width=paper_size_A4[0]) fig.suptitle("Channel {}".format(rf_stream[0].stats.channel)) # Customize layout to pack to top of page while preserving RF plots aspect ratios _rf_layout_A4(fig) # Save to new page in file pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() # Plot RF stack of transverse component fig = rf_plot_utils.plot_rf_stack( t_stream, trace_height=trace_ht, stack_height=fixed_stack_height_inches, fig_width=paper_size_A4[0]) fig.suptitle("Channel {}".format(t_stream[0].stats.channel)) # Customize layout to pack to top of page while preserving RF plots aspect ratios _rf_layout_A4(fig) # Save to new page in file pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() # Plot H-k stack using primary RF component fig = _produce_hk_stacking(rf_stream, weighting=hk_weights) paper_landscape = (paper_size_A4[1], paper_size_A4[0]) fig.set_size_inches(*paper_landscape) # plt.tight_layout() # plt.subplots_adjust(hspace=0.15, top=0.95, bottom=0.15) pdf.savefig(dpi=300, papertype='a4', orientation='landscape') plt.close() # end for pbar.close()
def main(): save_dist_plot = False output_folder = 'hk_validation' if not os.path.exists(output_folder): os.mkdir(output_folder) # end if n = 0 V_s = 3.8 # km/s, top (crust) layer H_range = np.linspace(30.0, 60.0, 4) # km V_p_range = np.linspace(5.7, 7.22, 5) # km/s include_t3 = True for H in H_range: for V_p in V_p_range: k = V_p / V_s log.info("H = {:.3f}".format(H)) log.info("k = {:.3f}".format(k)) # Generate function mapping ray parameter to teleseismic distance. # The distances are not strictly required for H-k stacking, but rf behaves better when they are there. ts_distance = np.linspace(25, 95, 71) inc = np.zeros_like(ts_distance) model = obspy.taup.TauPyModel(model="iasp91") source_depth_km = 10.0 for i, d in enumerate(ts_distance): ray = model.get_ray_paths(source_depth_km, d, phase_list=['P']) inc[i] = ray[0].incident_angle # pylint: disable=unsupported-assignment-operation # end for if save_dist_plot: plt.plot(inc, ts_distance, '-+') plt.xlabel('P-arrival inclination (deg)') plt.ylabel('Teleseismic distance (deg)') plt.title( "Relationship between incident angle and teleseismic distance\n" "(source depth = {:2.0f} km)".format(source_depth_km)) plt.grid("#80808080", linestyle=":") plt.savefig(os.path.join(output_folder, "inclination_dist.png"), dpi=300) # end if # Create interpolator to convert inclination to teleseismic distance interp_dist = interp1d(inc, ts_distance, bounds_error=False, fill_value=(np.max(ts_distance), np.min(ts_distance))) # Loop over valid range of inclinations and generate synthetic RFs num_traces = 20 inclinations = np.linspace(np.min(inc), np.max(inc), num_traces) distances = interp_dist(inclinations) final_sample_rate_hz = 10.0 stream = synthesize_rf_dataset(H, V_p, V_s, inclinations, distances, final_sample_rate_hz, log, include_t3) stream.write(os.path.join(output_folder, "synth_rf_data.h5"), format='h5') # Plot synthetic RFs rf_fig = rf_plot_utils.plot_rf_stack(stream, time_window=(-5, 30)) plt.title("Exact H = {:.3f}, k = {:.3f}".format(H, k)) rf_fig.savefig(os.path.join(output_folder, "synth_rf_{:03d}.png".format(n)), dpi=300) # Run H-k stacking on synthetic data station_db = {'HHR': stream} k_grid, h_grid, hk = rf_stacking.compute_hk_stack( station_db['HHR'], include_t3=include_t3) fig = rf_plot_utils.plot_hk_stack(k_grid, h_grid, hk[0], title="Synthetic Ps component") fig.savefig(os.path.join(output_folder, "Ps_{:03d}.png".format(n)), dpi=300) plt.close() fig = rf_plot_utils.plot_hk_stack(k_grid, h_grid, hk[1], title="Synthetic PpPs component") fig.savefig(os.path.join(output_folder, "PpPs_{:03d}.png".format(n)), dpi=300) plt.close() if include_t3: fig = rf_plot_utils.plot_hk_stack( k_grid, h_grid, hk[2], title="Synthetic PpSs+PsPs component") fig.savefig(os.path.join(output_folder, "PpSs+PsPs_{:03d}.png".format(n)), dpi=300) plt.close() w = (0.34, 0.33, 0.33) w_str = "w={:1.2f},{:1.2f},{:1.2f}".format(*w) else: w = (0.5, 0.5) w_str = "w={:1.2f},{:1.2f}".format(*w) # end if stack = rf_stacking.compute_weighted_stack(hk, weighting=w) hk_fig = rf_plot_utils.plot_hk_stack( k_grid, h_grid, stack, num=len(stream), title="Synthetic H-k stack ({})".format(w_str)) plt.text(1.45, 65.0, "Expected H = {:.3f}, k = {:.3f}".format(H, k), color="#ffffff", fontsize=16, fontweight='bold') # Determine H-k location of numerical maximum in the solution space. h_max, k_max = rf_stacking.find_global_hk_maximum( k_grid, h_grid, stack) log.info("Numerical solution (H, k) = ({:.3f}, {:.3f})".format( h_max, k_max)) plt.scatter(k_max, h_max, marker='+', c="#000000", s=20) if k_max >= 1.7: plt.text(k_max - 0.01, h_max + 1, "Solution H = {:.3f}, k = {:.3f}".format( h_max, k_max), color="#ffffff", fontsize=16, horizontalalignment='right') else: plt.text(k_max + 0.01, h_max + 1, "Solution H = {:.3f}, k = {:.3f}".format( h_max, k_max), color="#ffffff", fontsize=16) # end if # Save result hk_fig.savefig(os.path.join( output_folder, "synth_stack_{}_{:03d}.png".format(w_str, n)), dpi=300) n += 1
def main(input_file, output_file, event_mask_folder='', apply_amplitude_filter=False, apply_similarity_filter=False, hk_weights=DEFAULT_HK_WEIGHTS, hk_solution_labels=DEFAULT_HK_SOLN_LABEL, hk_hpf_freq=None, hk_vp=DEFAULT_Vp, save_hk_solution=False): # docstring redundant since CLI options are already documented. log.setLevel(logging.INFO) # Read source file log.info("Loading input file {}".format(input_file)) data_all = rf_util.read_h5_rf(input_file) # Convert to hierarchical dictionary format data_dict = rf_util.rf_to_dict(data_all) event_mask_dict = None if event_mask_folder and os.path.isdir(event_mask_folder): log.info( "Applying event mask from folder {}".format(event_mask_folder)) mask_files = os.listdir(event_mask_folder) mask_files = [ f for f in mask_files if os.path.isfile(os.path.join(event_mask_folder, f)) ] pattern = r"([A-Za-z0-9\.]{5,})_event_mask\.txt" pattern = re.compile(pattern) event_mask_dict = dict() for f in mask_files: match_result = pattern.match(f) if not match_result: continue code = match_result[1] with open(os.path.join(event_mask_folder, f), 'r') as _f: events = _f.readlines() events = set([e.strip() for e in events]) event_mask_dict[code] = events # end with # end for # end if if event_mask_dict: log.info("Loaded {} event masks".format(len(event_mask_dict))) # end if # Plot all data to PDF file fixed_stack_height_inches = 0.8 y_pad_inches = 1.6 total_trace_height_inches = paper_size_A4[ 1] - fixed_stack_height_inches - y_pad_inches max_trace_height = 0.2 log.setLevel(logging.WARNING) with PdfPages(output_file) as pdf: # Would like to use Tex, but lack desktop PC privileges to update packages to what is required plt.rc('text', usetex=False) pbar = tqdm.tqdm(total=len(data_dict)) network = data_dict.network rf_type = data_dict.rotation hk_soln = dict() station_coords = dict() for st in sorted(data_dict.keys()): station_db = data_dict[st] pbar.update() pbar.set_description("{}.{}".format(network, st)) # Choose RF channel channel = rf_util.choose_rf_source_channel(rf_type, station_db) channel_data = station_db[channel] if not channel_data: continue # end if full_code = '.'.join([network, st, channel]) t_channel = list(channel) t_channel[-1] = 'T' t_channel = ''.join(t_channel) rf_stream = rf.RFStream(channel_data).sort(['back_azimuth']) if event_mask_dict and full_code in event_mask_dict: # Select events from external source event_mask = event_mask_dict[full_code] rf_stream = rf.RFStream([ tr for tr in rf_stream if tr.stats.event_id in event_mask ]).sort(['back_azimuth']) # end if if apply_amplitude_filter: # Label and filter quality rf_util.label_rf_quality_simple_amplitude(rf_type, rf_stream) rf_stream = rf.RFStream([ tr for tr in rf_stream if tr.stats.predicted_quality == 'a' ]).sort(['back_azimuth']) # end if if not rf_stream: continue if apply_similarity_filter and len(rf_stream) >= 3: rf_stream = rf_util.filter_crosscorr_coeff(rf_stream) # end if if not rf_stream: continue # Find matching T-component data events = [tr.stats.event_id for tr in rf_stream] transverse_data = station_db[t_channel] t_stream = rf.RFStream([ tr for tr in transverse_data if tr.stats.event_id in events ]).sort(['back_azimuth']) # Plot pinwheel of primary and transverse components fig = rf_plot_utils.plot_rf_wheel([rf_stream, t_stream], fontscaling=0.8) fig.set_size_inches(*paper_size_A4) plt.tight_layout() plt.subplots_adjust(hspace=0.15, top=0.95, bottom=0.15) ax = fig.gca() fig.text(-0.32, -0.32, "\n".join(rf_stream[0].stats.processing), fontsize=6, transform=ax.transAxes) pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() num_traces = len(rf_stream) assert len(t_stream) == num_traces or not t_stream # Plot RF stack of primary component trace_ht = min(total_trace_height_inches / num_traces, max_trace_height) fig = rf_plot_utils.plot_rf_stack( rf_stream, trace_height=trace_ht, stack_height=fixed_stack_height_inches, fig_width=paper_size_A4[0]) fig.suptitle("Channel {}".format(rf_stream[0].stats.channel)) # Customize layout to pack to top of page while preserving RF plots aspect ratios _rf_layout_A4(fig) # Save to new page in file pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() # Plot RF stack of transverse component if t_stream: fig = rf_plot_utils.plot_rf_stack( t_stream, trace_height=trace_ht, stack_height=fixed_stack_height_inches, fig_width=paper_size_A4[0]) fig.suptitle("Channel {}".format(t_stream[0].stats.channel)) # Customize layout to pack to top of page while preserving RF plots aspect ratios _rf_layout_A4(fig) # Save to new page in file pdf.savefig(dpi=300, papertype='a4', orientation='portrait') plt.close() # end if # Plot H-k stack using primary RF component fig, maxima = _produce_hk_stacking(rf_stream, weighting=hk_weights, labelling=hk_solution_labels, V_p=hk_vp) if save_hk_solution and hk_hpf_freq is None: hk_soln[st] = maxima station_coords[st] = (channel_data[0].stats.station_latitude, channel_data[0].stats.station_longitude) # end if paper_landscape = (paper_size_A4[1], paper_size_A4[0]) fig.set_size_inches(*paper_landscape) # plt.tight_layout() # plt.subplots_adjust(hspace=0.15, top=0.95, bottom=0.15) pdf.savefig(dpi=300, papertype='a4', orientation='landscape') plt.close() if hk_hpf_freq is not None: # Repeat H-k stack with high pass filtering fig, maxima = _produce_hk_stacking( rf_stream, weighting=hk_weights, labelling=hk_solution_labels, V_p=hk_vp, filter_options={ 'type': 'highpass', 'freq': hk_hpf_freq, 'corners': 1, 'zerophase': True }) if save_hk_solution: hk_soln[st] = maxima station_coords[st] = ( channel_data[0].stats.station_latitude, channel_data[0].stats.station_longitude) # end if fig.set_size_inches(*paper_landscape) pdf.savefig(dpi=300, papertype='a4', orientation='landscape') plt.close() # end if # end for pbar.close() # end with # Save H-k solutions to CSV file if hk_soln: assert len(hk_soln) == len(station_coords) # Sort H-k solutions by depth from low to high update_dict = {} for st, hks in hk_soln.items(): sorted_hks = sorted([tuple(hk) for hk in hks]) update_dict[st] = np.array( list(station_coords[st]) + [i for hk in sorted_hks for i in hk]) # end for hk_soln.update(update_dict) df = pd.DataFrame.from_dict(hk_soln, orient='index') colnames = [('H{}'.format(i), 'k{}'.format(i)) for i in range((len(df.columns) - 2) // 2)] colnames = ['Latitude', 'Longitude'] + list( itertools.chain.from_iterable(colnames)) df.columns = colnames csv_fname, _ = os.path.splitext(output_file) csv_fname += '.csv' df.index.name = 'Station' df.to_csv(csv_fname)
def main(): """Main entry function for RF picking tool. """ infile = filedialog.askopenfilename(initialdir=".", title="Select RF file", filetypes=(("h5 files", "*.h5"), )) output_folder = filedialog.askdirectory( initialdir=os.path.split(infile)[0], title='Select output folder') if not os.path.isdir(output_folder): log.info("Creating output folder {}".format(output_folder)) os.makedirs(output_folder, exist_ok=True) # end if log.info("Output files will be emitted to {}".format(output_folder)) log.info("Loading %s", infile) data_all = rf_util.read_h5_rf(infile) data_dict = rf_util.rf_to_dict(data_all) stations = sorted(list(data_dict.keys())) # Assuming same rotation type for all RFs. This is consistent with the existing workflow. rf_type = data_all[0].stats.rotation for st in stations: station_db = data_dict[st] # Choose RF channel channel = rf_util.choose_rf_source_channel(rf_type, station_db) channel_data = station_db[channel] # Check assumption for tr in channel_data: assert tr.stats.rotation == rf_type, 'Mismatching RF rotation type' # Label and filter quality rf_util.label_rf_quality_simple_amplitude(rf_type, channel_data) rf_stream = rf.RFStream([ tr for tr in channel_data if tr.stats.predicted_quality == 'a' ]).sort(['back_azimuth']) if not rf_stream: log.info("No data survived filtering for %s, skipping", st) continue # Plot RF stack of primary component fig = rf_plot_utils.plot_rf_stack(rf_stream) fig.set_size_inches(8, 9) fig.suptitle("Channel {}".format(rf_stream[0].stats.channel)) ax0 = fig.axes[0] # Make sure we draw once first before capturing blit background fig.canvas.draw() # Disallow resizing to avoid having to handle blitting with resized window. win = fig.canvas.window() win.setFixedSize(win.size()) blit_background = fig.canvas.copy_from_bbox(ax0.bbox) mask = np.array([False] * len(rf_stream)) rect_select = RectangleSelector(ax0, lambda e0, e1: on_select(e0, e1, mask), useblit=True, rectprops=dict(fill=False, edgecolor='red')) cid = fig.canvas.mpl_connect( 'button_release_event', lambda e: on_release(e, ax0, mask, blit_background, rect_select)) plt.show() fig.canvas.mpl_disconnect(cid) rect_select = None selected_event_ids = [ tr.stats.event_id for i, tr in enumerate(rf_stream) if mask[i] ] log.info("{} streams selected".format(len(selected_event_ids))) log.info("Selected event ids:") log.info(selected_event_ids) network = rf_stream[0].stats.network outfile = os.path.join( output_folder, '.'.join([network, st, channel]) + '_event_mask.txt') log.info("Writing mask to file {}".format(outfile)) if os.path.exists(outfile): log.warning("Overwriting existing file {} !".format(outfile)) with open(outfile, 'w') as f: f.write('\n'.join(selected_event_ids))