Skip to content

Commit

Permalink
move plotr to class
Browse files Browse the repository at this point in the history
  • Loading branch information
@xumi1993
xumi1993 committed Nov 22, 2023
1 parent befb1f0 commit e51a046
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 92 deletions.
3 changes: 2 additions & 1 deletion seispy/pickrf/rpickfigure.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
pr = stadata.plotr(enf=self.enf, xlim=self.xlim, output=None)
self.plotfig = pr.fig
198 changes: 108 additions & 90 deletions seispy/plotR.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion seispy/scripts.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit e51a046

Please sign in to comment.