diff --git a/seispy/pickrf/rpickfigure.py b/seispy/pickrf/rpickfigure.py index 8c3bcf95..b31a7669 100644 --- a/seispy/pickrf/rpickfigure.py +++ b/seispy/pickrf/rpickfigure.py @@ -169,4 +169,5 @@ def plot(self): stadata = RFStation.read_stream(newrfs, self.rayp[goodidx], self.baz[goodidx], prime_comp=self.comp) stadata.event = np.array([self.filenames[i] for i in goodidx]) - self.plotfig = stadata.plotr(enf=self.enf, xlim=self.xlim) \ No newline at end of file + pr = stadata.plotr(enf=self.enf, xlim=self.xlim, output=None) + self.plotfig = pr.fig \ No newline at end of file diff --git a/seispy/plotR.py b/seispy/plotR.py index b61f9475..40d5332b 100644 --- a/seispy/plotR.py +++ b/seispy/plotR.py @@ -5,109 +5,127 @@ from os.path import join -def init_figure(): - h = plt.figure(figsize=(8, 10)) - h.tight_layout() - gs = GridSpec(15, 3) - gs.update(wspace=0.25) - axr = plt.subplot(gs[1:, 0:-1]) - axr.grid(color='gray', linestyle='--', linewidth=0.4, axis='x') - axb = plt.subplot(gs[1:, -1]) - axb.grid(color='gray', linestyle='--', linewidth=0.4, axis='x') - axs = plt.subplot(gs[0, 0:-1]) - axs.grid(color='gray', linestyle='--', linewidth=0.4, axis='x') - return h, axr, axb, axs +class PlotR: + def __init__(self, stadata, enf=12, xlim=[2, 80], key='bazi'): + """Plot receiver function in R component + :param stadata: receiver function data + :type stadata: sespy.rfcorrect.RFStation + :param enf: enlarge factor, defaults to 12 + :type enf: int, optional + :param xlim: xlim, defaults to [2, 80] + :type xlim: list, optional + :param key: sort key, defaults to 'bazi' + :type key: str, optional + """ + self.stadata = stadata + self.stadata.sort(key) + self.enf = enf + self.xlim = xlim + self.key = key + self.fig, self.axr, self.axb, self.axs = self.init_figure() -def plot_waves(axr, axb, axs, stadata, enf=12): - bound = np.zeros(stadata.rflength) - for i in range(stadata.ev_num): - datar = stadata.data_prime[i] * enf + (i + 1) - axr.fill_between(stadata.time_axis, datar, bound + i+1, where=datar > i+1, facecolor='red', - alpha=0.7) - axr.fill_between(stadata.time_axis, datar, bound + i+1, where=datar < i+1, facecolor='blue', - alpha=0.7) - # rfcorr, _ = stadata.moveoutcorrect( 0.1, np.arange(300), sphere=False) - # rfsum = np.mean(rfcorr, axis=0) - rfsum = np.mean(stadata.data_prime, axis=0) - axs.fill_between(stadata.time_axis, rfsum, 0, where=rfsum > 0, facecolor='red', alpha=0.7) - axs.fill_between(stadata.time_axis, rfsum, 0, where=rfsum < 0, facecolor='blue', alpha=0.7) - # axs.plot(stadata.time_axis, rfsum, color='gray', lw=0.5) - axb.scatter(stadata.bazi, np.arange(stadata.ev_num) + 1, s=7) - # axp = axb.twiny() - # axp.scatter(stadata.rayp, np.arange(stadata.ev_num) + 1, s=7) - # return axp + def plot_waves(self): + bound = np.zeros(self.stadata.rflength) + for i in range(self.stadata.ev_num): + datar = self.stadata.data_prime[i] * self.enf + (i + 1) + self.axr.fill_between(self.stadata.time_axis, datar, bound + i+1, where=datar > i+1, facecolor='red', + alpha=0.7) + self.axr.fill_between(self.stadata.time_axis, datar, bound + i+1, where=datar < i+1, facecolor='blue', + alpha=0.7) + # rfcorr, _ = stadata.moveoutcorrect( 0.1, np.arange(300), sphere=False) + # rfsum = np.mean(rfcorr, axis=0) + rfsum = np.mean(self.stadata.data_prime, axis=0) + self.axs.fill_between(self.stadata.time_axis, rfsum, 0, where=rfsum > 0, facecolor='red', alpha=0.7) + self.axs.fill_between(self.stadata.time_axis, rfsum, 0, where=rfsum < 0, facecolor='blue', alpha=0.7) + # axs.plot(stadata.time_axis, rfsum, color='gray', lw=0.5) + self.axb.scatter(self.stadata.bazi, np.arange(self.stadata.ev_num) + 1, s=7) + # axp = axb.twiny() + # axp.scatter(stadata.rayp, np.arange(stadata.ev_num) + 1, s=7) + # return axp -def set_fig(axr, axb, axs, stadata, xmin=-2, xmax=80): - y_range = np.arange(stadata.ev_num) + 1 - x_range = np.arange(0, xmax+2, 5) - space = 2 + def set_fig(self): + y_range = np.arange(self.stadata.ev_num) + 1 + x_range = np.arange(0, self.xlim[1]+2, 5) + space = 2 - # set axr - axr.set_xlim(xmin, xmax) - axr.set_xticks(x_range) - axr.set_xticklabels(x_range, fontsize=8) - axr.set_ylim(0, stadata.ev_num + space) - axr.set_yticks(y_range) - axr.set_yticklabels(stadata.event, fontsize=5) - axr.set_xlabel('Time after P (s)', fontsize=13) - axr.set_ylabel('Event', fontsize=13) - axr.add_line(Line2D([0, 0], axr.get_ylim(), color='black')) - # axr.set_title('R components ({})'.format(stadata.staname), fontsize=16) + # set axr + self.axr.set_xlim(self.xlim) + self.axr.set_xticks(x_range) + self.axr.set_xticklabels(x_range, fontsize=8) + self.axr.set_ylim(0, self.stadata.ev_num + space) + self.axr.set_yticks(y_range) + self.axr.set_yticklabels(self.stadata.event, fontsize=5) + self.axr.set_xlabel('Time after P (s)', fontsize=13) + self.axr.set_ylabel('Event', fontsize=13) + self.axr.add_line(Line2D([0, 0], self.axr.get_ylim(), color='black')) + # axr.set_title('R components ({})'.format(stadata.staname), fontsize=16) - # set axb - axb.set_xlim(0, 360) - axb.set_xticks(np.linspace(0, 360, 7)) - axb.set_xticklabels(np.linspace(0, 360, 7, dtype='i'), fontsize=8) - axb.set_ylim(0, stadata.ev_num + space) - axb.set_yticks(y_range) - axb.set_yticklabels(y_range, fontsize=5) - axb.set_xlabel(r'Back-azimuth ($^\circ$)', fontsize=13) - # axp.set_xlabel('Ray-parameter (s/km)', fontsize=13) + # set axb + self.axb.set_xlim(0, 360) + self.axb.set_xticks(np.linspace(0, 360, 7)) + self.axb.set_xticklabels(np.linspace(0, 360, 7, dtype='i'), fontsize=8) + self.axb.set_ylim(0, self.stadata.ev_num + space) + self.axb.set_yticks(y_range) + self.axb.set_yticklabels(y_range, fontsize=5) + self.axb.set_xlabel(r'Back-azimuth ($^\circ$)', fontsize=13) + # axp.set_xlabel('Ray-parameter (s/km)', fontsize=13) - axs.set_title('{} components ({})'.format(stadata.comp.upper(), stadata.staname), fontsize=16) - axs.set_xlim(xmin, xmax) - axs.set_xticks(x_range) - axs.set_xticklabels([]) - axs.set_yticks([np.sum(axs.get_ylim())/3]) - axs.tick_params(axis='y', left=False) - axs.set_yticklabels(['Sum'], fontsize=8) - axs.add_line(Line2D([0, 0], axs.get_ylim(), color='black')) + self.axs.set_title('{} components ({})'.format(self.stadata.comp.upper(), self.stadata.staname), fontsize=16) + self.axs.set_xlim(self.xlim) + self.axs.set_xticks(x_range) + self.axs.set_xticklabels([]) + self.axs.set_yticks([np.sum(self.axs.get_ylim())/3]) + self.axs.tick_params(axis='y', left=False) + self.axs.set_yticklabels(['Sum'], fontsize=8) + self.axs.add_line(Line2D([0, 0], self.axs.get_ylim(), color='black')) + + def plot(self, out_path=None, outformat='g', show=False): + """ Plot receiver function in R component + :param out_path: output path + :type out_path: str + :param outformat: output format + :type outformat: str + :param show: show figure + :type show: bool + """ + 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+'_R_{}order_{:.1f}.png'.format(self.key, self.stadata.f0[0])), + dpi=400, bbox_inches='tight') + elif outformat == 'f': + self.fig.savefig(join(out_path, self.stadata.staname+'_R_{}order_{:.1f}.pdf'.format(self.key, self.stadata.f0[0])), + format='pdf', bbox_inches='tight') + if show: + plt.show() def plotr(rfsta, out_path='./', xlim=[-2, 80], key='bazi', enf=6, outformat='g', show=False): """ Plot receiver function in R component - :param rfsta: receiver function data - :type rfsta: sespy.rfcorrect.RFStation - :param out_path: output path - :type out_path: str - :param xlim: xlim - :type xlim: list - :param key: sort key - :type key: str - :param enf: enhancement factor - :type enf: float - :param outformat: output format - :type outformat: str + :param rfsta: Path to PRFs + :type rfsta: seispy.rfcorrect.RFStation + :param enf: The enlarge factor, defaults to 6 + :type enf: int, optional + :param out_path: The output path, defaults to current directory + :type out_path: str, optional + :param key: The key to sort PRFs, available for ``event``, ``evla``, ``evlo``, ``evdp``, + ``dis``, ``bazi``, ``rayp``, ``mag``, ``f0``, defaults to ``bazi`` + :type key: str, optional + :param outformat: File format of the image file, g as \'png\', f as \'pdf\', defaults to 'g' + :type outformat: str, optional + :param xlim: xlim, defaults to [-2, 80] + :type xlim: list, optional :param show: show figure :type show: bool + :return: PlotR object + :rtype: PlotR """ - h, axr, axb, axs = init_figure() - rfsta.sort(key) - plot_waves(axr, axb, axs, rfsta, enf=enf) - set_fig(axr, axb, axs, rfsta, xlim[0], xlim[1]) - if outformat is None and not show: - return h - elif outformat == 'g': - h.savefig(join(out_path, rfsta.staname+'_R_{}order_{:.1f}.png'.format(key, rfsta.f0[0])), - dpi=400, bbox_inches='tight') - elif outformat == 'f': - h.savefig(join(out_path, rfsta.staname+'_R_{}order_{:.1f}.pdf'.format(key, rfsta.f0[0])), - format='pdf', bbox_inches='tight') - if show: - plt.show() - return h - + pr = PlotR(rfsta, enf=enf, xlim=xlim, key=key) + pr.plot(out_path=out_path, outformat=outformat, show=show) + return pr if __name__ == '__main__': pass \ No newline at end of file diff --git a/seispy/scripts.py b/seispy/scripts.py index 0cdddfc1..392d1e46 100644 --- a/seispy/scripts.py +++ b/seispy/scripts.py @@ -261,5 +261,5 @@ def plot_r(): elif arg.format == 'f': fmt = 'pdf' rfsta = RFStation(arg.rfpath) - plotr(rfsta, arg.output, enf=arg.enf, key=arg.k, xlim=[-2, arg.x], format=fmt) + rfsta.plotr(rfsta, arg.output, enf=arg.enf, key=arg.k, xlim=[-2, arg.x], format=fmt)