diff --git a/seispy/para.py b/seispy/para.py index c1c6994c..48786235 100644 --- a/seispy/para.py +++ b/seispy/para.py @@ -134,6 +134,7 @@ def __init__(self): self.use_remote_data=False self.data_server_user = None self.data_server_password = None + self.n_proc = 1 self.stainfo = StaInfo() def get_para(self): @@ -276,6 +277,8 @@ def read_para(cls, cfg_file): pa.data_server_password = value elif key == 'data_server': pa.data_server = value + elif key == 'n_proc': + pa.n_proc = cf.getint('fetch', 'n_proc') else: exec('pa.stainfo.{} = value'.format(key)) sections.remove('fetch') diff --git a/seispy/plotR.py b/seispy/plotR.py index e1c7c76e..07e789a0 100644 --- a/seispy/plotR.py +++ b/seispy/plotR.py @@ -109,16 +109,15 @@ def plot(self, out_path=None, outformat='g', show=False): """ self.plot_waves() self.set_fig() - if outformat is None and not show: - return - elif outformat == 'g': - self.fig.savefig(join(out_path, self.stadata.staname+'_{}_{}order_{:.1f}.png'.format( - self.stadata.comp, self.key, self.stadata.f0[0])), - dpi=400, bbox_inches='tight') - elif outformat == 'f': - self.fig.savefig(join(out_path, self.stadata.staname+'_{}_{}order_{:.1f}.pdf'.format( - self.stadata.comp, self.key, self.stadata.f0[0])), - format='pdf', bbox_inches='tight') + if out_path is not None: + if outformat == 'g': + self.fig.savefig(join(out_path, self.stadata.staname+'_{}_{}order_{:.1f}.png'.format( + self.stadata.comp, self.key, self.stadata.f0[0])), + dpi=400, bbox_inches='tight') + elif outformat == 'f': + self.fig.savefig(join(out_path, self.stadata.staname+'_{}_{}order_{:.1f}.pdf'.format( + self.stadata.comp, self.key, self.stadata.f0[0])), + format='pdf', bbox_inches='tight') if show: plt.show() diff --git a/seispy/plotRT.py b/seispy/plotRT.py index 027926ec..232d1659 100644 --- a/seispy/plotRT.py +++ b/seispy/plotRT.py @@ -148,16 +148,15 @@ def plot(self, out_path=None, outformat='g', show=False): """ self.plot_waves() self.set_fig() - if outformat is None and not show: - return - elif outformat == 'g': - self.fig.savefig(join(out_path, self.stadata.staname+'_{}T_{}order_{:.1f}.png'.format( - self.stadata.comp, self.key, self.stadata.f0[0])), - dpi=400, bbox_inches='tight') - elif outformat == 'f': - self.fig.savefig(join(out_path, self.stadata.staname+'_{}T_{}order_{:.1f}.pdf'.format( - self.stadata.comp, self.key, self.stadata.f0[0])), - format='pdf', bbox_inches='tight') + if out_path is not None: + if outformat == 'g': + self.fig.savefig(join(out_path, self.stadata.staname+'_{}T_{}order_{:.1f}.png'.format( + self.stadata.comp, self.key, self.stadata.f0[0])), + dpi=400, bbox_inches='tight') + elif outformat == 'f': + self.fig.savefig(join(out_path, self.stadata.staname+'_{}T_{}order_{:.1f}.pdf'.format( + self.stadata.comp, self.key, self.stadata.f0[0])), + format='pdf', bbox_inches='tight') if show: plt.show() diff --git a/seispy/psrayp.py b/seispy/psrayp.py index 1d7e6177..db9d243d 100644 --- a/seispy/psrayp.py +++ b/seispy/psrayp.py @@ -50,7 +50,6 @@ def taup_rayp(self, this_dis=50, this_dep=10, taup='taup_time'): def get_rayp(self): for i in range(self.dis.shape[0]): - print('{}'.format(self.dis[i])) for j in range(self.dep.shape[0]): self.rayp[i, j, :] = self.taup_rayp(this_dis=self.dis[i], this_dep=self.dep[j])[1] diff --git a/seispy/rf.py b/seispy/rf.py index a8ff9a9e..f0d75a8b 100644 --- a/seispy/rf.py +++ b/seispy/rf.py @@ -21,6 +21,7 @@ import argparse import sys import pickle +import concurrent.futures def pickphase(eqs, para, logger): @@ -109,6 +110,44 @@ def read_catalog(logpath:str, b_time, e_time, stla:float, stlo:float, return eq_lst +def process_row(i, size, row, para, model, query, tb, te, logger): + new_col = ['dis', 'bazi', 'data', 'datestr'] + datestr = row['date'].strftime('%Y.%j.%H.%M.%S') + daz = distaz(para.stainfo.stla, para.stainfo.stlo, row['evla'], row['evlo']) + arrivals = model.get_travel_times(row['evdp'], daz.delta, phase_list=[para.phase]) + + if not arrivals: + logger.RFlog.error('The phase of {} with source depth {} and distance {} is not exists'.format( + para.phase, row['evdp'], daz.delta)) + return None + + if len(arrivals) > 1: + logger.RFlog.error('More than one phase were calculated with source depth of {} and distance of {}'.format( + row['evdp'], daz.delta)) + return None + + arr_time = arrivals[0].time + t1 = row['date'] + arr_time - tb + t2 = row['date'] + arr_time + te + try: + logger.RFlog.info('Fetch waveforms of event {} ({}/{}) from {}'.format(datestr, i+1, size, para.data_server)) + st = query.client.get_waveforms(para.stainfo.network, para.stainfo.station, + para.stainfo.location, para.stainfo.channel, t1, t2) + _add_header(st, row['date'], para.stainfo) + except Exception as e: + logger.RFlog.error('Error in fetching waveforms of event {}: {}'.format(datestr, str(e).strip())) + return None + + try: + this_eq = EQ.from_stream(st) + except Exception as e: + logger.RFlog.error('{}'.format(e)) + return None + + this_eq.get_time_offset(row['date']) + this_df = pd.DataFrame([[daz.delta, daz.baz, this_eq, datestr]], columns=new_col, index=[i]) + return this_df + def fetch_waveform(eq_lst, para, model, logger): """Fetch waveforms from remote data server @@ -129,43 +168,22 @@ def fetch_waveform(eq_lst, para, model, logger): query = para.stainfo.query except: logger.RFlog.error('Please load station information and search earthquake before fetch waveform') - new_col = ['dis', 'bazi', 'data', 'datestr'] eqall = [] - for i, row in eq_lst.iterrows(): - datestr = row['date'].strftime('%Y.%j.%H.%M.%S') - daz = distaz(para.stainfo.stla, para.stainfo.stlo, row['evla'], row['evlo']) - arrivals = model.get_travel_times(row['evdp'], daz.delta, phase_list=[para.phase]) - if not arrivals: - logger.RFlog.error('The phase of {} with source depth {} and distance {} is not exists'.format( - para.phase, row['evdp'], daz.delta)) - continue - if len(arrivals) > 1: - logger.RFlog.error('More than one phase were calculated with source depth of {} and distance of {}'.format( - row['evdp'], daz.delta)) - else: - arr_time = arrivals[0].time - t1 = row['date']+arr_time-tb - t2 = row['date']+arr_time+te - try: - logger.RFlog.info('Fetch waveforms of ({}/{}) event {} from {}'.format( - i+1, eq_lst.shape[0], datestr, para.data_server)) - st = query.client.get_waveforms(para.stainfo.network, para.stainfo.station, - para.stainfo.location, para.stainfo.channel, t1, t2) - _add_header(st, row['date'], para.stainfo) - except Exception as e: - logger.RFlog.error('Error in fetching waveforms of event {}: {}'.format(datestr, str(e).strip())) - continue - try: - this_eq = EQ.from_stream(st) - except Exception as e: - logger.RFlog.error('{}'.format(e)) - continue - this_eq.get_time_offset(row['date']) - this_df = pd.DataFrame([[daz.delta, daz.baz, this_eq, datestr]], columns=new_col, index=[i]) - eqall.append(this_df) + + # parallel downloading waveforms + with concurrent.futures.ProcessPoolExecutor(max_workers=para.n_proc) as executor: + futures = {executor.submit(process_row, i, eq_lst.shape[0], row, para, model, query, tb, te, logger): i for i, row in eq_lst.iterrows()} + + for future in concurrent.futures.as_completed(futures): + result = future.result() + if result is not None: + eqall.append(result) + if not eqall: logger.RFlog.error('No waveforms fetched') sys.exit(1) + + # list to DataFrame eq_match = pd.concat(eqall) ind = eq_match.index.drop_duplicates(keep=False) eq_match = eq_match.loc[ind] diff --git a/seispy/rfcorrect.py b/seispy/rfcorrect.py index afe89389..67bd5622 100644 --- a/seispy/rfcorrect.py +++ b/seispy/rfcorrect.py @@ -470,23 +470,19 @@ def harmonic(self, tb=-5, te=10, is_stack=True): self.harmo.harmo_trans() return self.harmo.harmonic_trans, self.harmo.unmodel_trans - def plotrt(self, outformat=None, **kwargs): + def plotrt(self, **kwargs): """Plot radial and transverse RFs - - :param outformat: Output format for plot, defaults to None - :type outformat: str, optional + :param kwargs: Parameters for plot, see :meth:`seispy.plotRT.plotrt` for detail """ if self.only_r: raise ValueError('Transverse RFs are nessary or use RFStation.plotr instead.') - return _plotrt(self, outformat=outformat, **kwargs) + return _plotrt(self, **kwargs) - def plotr(self, outformat=None, **kwargs): + def plotr(self,**kwargs): """Plot radial RFs - - :param outformat: Output format for plot, defaults to None - :type outformat: str, optional + :param kwargs: Parameters for plot, see :meth:`seispy.plotR.plotr` for detail """ - return _plotr(self, outformat=outformat, **kwargs) + return _plotr(self, **kwargs) class SACStation(RFStation): @@ -716,7 +712,6 @@ def psrf_1D_raytracing(stadatar, YAxisRange, velmod='iasp91', srayp=None, sphere for i in range(stadatar.ev_num): tps[i], x_s, x_p, raylength_s[i], raylength_p[i] = xps_tps_map( dep_mod, stadatar.rayp[i], stadatar.rayp[i], is_raylen=True, sphere=sphere, phase=phase) - print(type(stadatar.bazi[i]), type(rad2deg(x_s))) pplat_s[i], pplon_s[i] = latlon_from(stadatar.stla, stadatar.stlo, stadatar.bazi[i], rad2deg(x_s)) pplat_p[i], pplon_p[i] = latlon_from(stadatar.stla, stadatar.stlo, stadatar.bazi[i], rad2deg(x_p)) elif isinstance(srayp, str) or isinstance(srayp, np.lib.npyio.NpzFile): diff --git a/seispy/scripts.py b/seispy/scripts.py index 6eb38f56..274f286e 100644 --- a/seispy/scripts.py +++ b/seispy/scripts.py @@ -247,7 +247,7 @@ def plot_rt(): if arg.format not in ('f', 'g'): raise ValueError('Error: The format must be in \'f\' and \'g\'') rfsta = RFStation(arg.rfpath, prime_comp=arg.c) - rfsta.plotrt(rfsta, enf=arg.enf, out_path=arg.output, key=arg.k, outformat=arg.format, xlim=[-2, arg.x]) + rfsta.plotrt(enf=arg.enf, out_path=arg.output, key=arg.k, outformat=arg.format, xlim=[-2, arg.x]) def plot_r(): @@ -268,7 +268,7 @@ def plot_r(): if arg.format not in ('f', 'g'): parser.error('Error: The format must be in \'f\' and \'g\'') rfsta = RFStation(arg.rfpath, prime_comp=arg.c) - rfsta.plotr(rfsta, arg.output, enf=arg.enf, key=arg.k, xlim=[-2, arg.x], outformat=arg.format) + rfsta.plotr(arg.output, enf=arg.enf, key=arg.k, xlim=[-2, arg.x], outformat=arg.format) def get_events(): diff --git a/setup.py b/setup.py index 1aaa05b1..0e91e5b6 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ long_description = fh.read() -VERSION = "1.3.7" +VERSION = "1.3.8" setup(name='python-seispy', version=VERSION, author='Mijian Xu', @@ -18,7 +18,7 @@ package_dir={'seispy': 'seispy'}, package_data={'': ['data/*']}, install_requires=[ - 'numpy>=1.19.0', + 'numpy>=1.19.0,<2.0', 'scipy>=1.9.1', 'matplotlib>=3.5.0', 'pandas>=1.0.0',