Skip to content

Commit

Permalink
update plotrt
Browse files Browse the repository at this point in the history
  • Loading branch information
@xumi1993
xumi1993 committed Nov 23, 2023
1 parent 83e1c5e commit de2f9de
Show file tree
Hide file tree
Showing 6 changed files with 204 additions and 153 deletions.
20 changes: 14 additions & 6 deletions seispy/pickrf/pickfigure.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from matplotlib.figure import Figure
from os.path import join, basename
from seispy.setuplog import setuplog
from seispy.plotRT import init_figure, set_fig, plot_waves
from seispy.rfcorrect import RFStation


def indexpags(evt_num, maxidx=20):
Expand Down Expand Up @@ -158,6 +158,7 @@ def read_sac(self, dt=0.1, order='baz'):
self.log.RFlog.info('A total of {} PRFs loaded'.format(self.evt_num))
self.baz = np.array([tr.stats.sac.baz for tr in self.rrf])
self.gcarc = np.array([tr.stats.sac.gcarc for tr in self.rrf])
self.rayp = np.array([tr.stats.sac.user0 for tr in self.rrf])
self._sort(order)
self.axpages, self.rfidx = indexpags(self.evt_num, self.maxidx)
self.staname = (self.rrf[0].stats.network+'.'+self.rrf[0].stats.station).strip('.')
Expand Down Expand Up @@ -286,8 +287,15 @@ def finish(self):
def plot(self):
plt.ion()
plt.rcParams['toolbar'] = 'None'
stadata = StaData(self.filenames, self.rrf, self.trf, self.baz, self.goodrf)
stadata.time_axis = self.time_axis
self.plotfig, axr, axt, axb, axr_sum, axt_sum = init_figure()
plot_waves(axr, axt, axb, axr_sum, axt_sum, stadata, enf=self.enf)
set_fig(axr, axt, axb, axr_sum, axt_sum, stadata, self.staname)
goodidx = np.where(self.goodrf == 1)[0]
newrfs = obspy.Stream([self.rrf[idx] for idx in goodidx])
newtrfs = obspy.Stream([self.trf[idx] for idx in goodidx])
stadata = RFStation.read_stream(newrfs, self.rayp[goodidx],
self.baz[goodidx], prime_comp=self.comp,
stream_t=newtrfs)
stadata.event = np.array([self.filenames[i] for i in goodidx])
prt = stadata.plotrt(enf=self.enf, xlim=self.xlim, out_path=None)
self.plotfig = prt.fig
# self.plotfig, axr, axt, axb, axr_sum, axt_sum = init_figure()
# plot_waves(axr, axt, axb, axr_sum, axt_sum, stadata, enf=self.enf)
# set_fig(axr, axt, axb, axr_sum, axt_sum, stadata, self.staname)
2 changes: 1 addition & 1 deletion seispy/pickrf/rpickfigure.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,5 +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])
pr = stadata.plotr(enf=self.enf, xlim=self.xlim, output=None)
pr = stadata.plotr(enf=self.enf, xlim=self.xlim, out_path=None)
self.plotfig = pr.fig
10 changes: 6 additions & 4 deletions seispy/plotR.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,11 +106,13 @@ def plot(self, out_path=None, outformat='g', show=False):
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')
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+'_R_{}order_{:.1f}.pdf'.format(self.key, self.stadata.f0[0])),
format='pdf', bbox_inches='tight')
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()

Expand Down
303 changes: 165 additions & 138 deletions seispy/plotRT.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,133 +5,168 @@
from os.path import join


def init_figure():
h = plt.figure(figsize=(11.7, 8.3))
gs = GridSpec(17, 3)
gs.update(wspace=0.25)
axr_sum = plt.subplot(gs[0, 0])
axr_sum.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
axr = plt.subplot(gs[1:, 0])
axr.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
axt_sum = plt.subplot(gs[0, 1])
axt_sum.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
axt = plt.subplot(gs[1:, 1])
axt.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
axb = plt.subplot(gs[1:, 2])
axb.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
return h, axr, axt, axb, axr_sum, axt_sum


def plot_waves(axr, axt, axb, axr_sum, axt_sum, stadata, enf=3):
"""Plot PRFs with R and T components
:param axr: The axis of R component
:type axr: matplotlib.axes
:param axt: The axis of T component
:type axt: matplotlib.axes
:param axb: The axis of back-azimuth
:type axb: matplotlib.axes
:param axr_sum: The axis of R component sum
:type axr_sum: matplotlib.axes
:param axt_sum: The axis of T component sum
:type axt_sum: matplotlib.axes
:param stadata: The PRFs data
:type stadata: seispy.core.StaData
:param enf: The enlarge factor, defaults to 3
:type enf: int, optional
"""

bound = np.zeros(stadata.rflength)
for i in range(stadata.ev_num):
datar = stadata.datar[i] * enf + (i + 1)
datat = stadata.datat[i] * enf + (i + 1)
# axr.plot(time_axis, stadata.datar[i], linewidth=0.2, color='black')
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)
# axt.plot(time_axis, stadata.datat[i], linewidth=0.2, color='black')
axt.fill_between(stadata.time_axis, datat, bound + i + 1, where=datat > i+1, facecolor='red',
alpha=0.7)
axt.fill_between(stadata.time_axis, datat, bound + i + 1, where=datat < i+1, facecolor='blue',
alpha=0.7)
datar = np.mean(stadata.datar, axis=0)
datar /= np.max(datar)
datat = np.mean(stadata.datat, axis=0)
datat /= np.max(datar)
axr_sum.fill_between(stadata.time_axis, datar, bound, where=datar > 0, facecolor='red', alpha=0.7)
axr_sum.fill_between(stadata.time_axis, datar, bound, where=datar < 0, facecolor='blue', alpha=0.7)
axt_sum.fill_between(stadata.time_axis, datat, bound, where=datat > 0, facecolor='red', alpha=0.7)
axt_sum.fill_between(stadata.time_axis, datat, bound, where=datat < 0, facecolor='blue', alpha=0.7)
axb.scatter(stadata.bazi, np.arange(stadata.ev_num) + 1, s=7)


def set_fig(axr, axt, axb, axr_sum, axt_sum, stadata, station, xmin=-2, xmax=30, comp='R'):
y_range = np.arange(stadata.ev_num) + 1
x_range = np.arange(0, xmax+2, 2)
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'))

# set axr_sum
axr_sum.set_title('{} components ({})'.format(comp, station), fontsize=16)
axr_sum.set_xlim(xmin, xmax)
axr_sum.set_xticks(x_range)
axr_sum.set_xticklabels([])
axr_sum.set_ylim(-0.5, 1.25)
axr_sum.set_yticks([0.375])
axr_sum.set_yticklabels(['Sum'], fontsize=8)
axr_sum.tick_params(axis='y', left=False)
axr_sum.add_line(Line2D([0, 0], axr_sum.get_ylim(), color='black'))

# set axt
axt.set_xlim(xmin, xmax)
axt.set_xticks(x_range)
axt.set_xticklabels(x_range, fontsize=8)
axt.set_ylim(0, stadata.ev_num + space)
axt.set_yticks(y_range)
bazi = ['{:.1f}'.format(ba) for ba in stadata.bazi]
axt.set_yticklabels(bazi, fontsize=5)
axt.set_xlabel('Time after P (s)', fontsize=13)
axt.set_ylabel(r'Back-azimuth ($\circ$)', fontsize=13)
axt.add_line(Line2D([0, 0], axt.get_ylim(), color='black'))

# set axt_sum
axt_sum.set_title('T components ({})'.format(station), fontsize=16)
axt_sum.set_xlim(xmin, xmax)
axt_sum.set_xticks(x_range)
axt_sum.set_xticklabels([])
axt_sum.set_ylim(-0.5, 1.25)
axt_sum.set_yticks([0.375])
axt_sum.set_yticklabels(['Sum'], fontsize=8)
axt_sum.tick_params(axis='y', left=False)
axt_sum.add_line(Line2D([0, 0], axt_sum.get_ylim(), color='black'))

# 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)


def plotrt(rfsta, enf=3, out_path='./', key='bazi', outformat='g', xlim=[-2, 30], show=False):
class PlotRT:
def __init__(self, stadata, enf=3, xlim=[2, 80], key='bazi') -> None:
"""Plot receiver function in R and T components
:param stadata: receiver function data
:type stadata: seispy.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.init_figure()

def init_figure(self):
self.fig = plt.figure(figsize=(11.7, 8.3))
gs = GridSpec(17, 3)
gs.update(wspace=0.25)
self.axr_sum = plt.subplot(gs[0, 0])
self.axr_sum.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
self.axr = plt.subplot(gs[1:, 0])
self.axr.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
self.axt_sum = plt.subplot(gs[0, 1])
self.axt_sum.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
self.axt = plt.subplot(gs[1:, 1])
self.axt.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')
self.axb = plt.subplot(gs[1:, 2])
self.axb.grid(color='gray', linestyle='--', linewidth=0.4, axis='x')

def plot_waves(self):
"""Plot PRFs with R and T components
"""

bound = np.zeros(self.stadata.rflength)
for i in range(self.stadata.ev_num):
datar = self.stadata.data_prime[i] * self.enf + (i + 1)
datat = self.stadata.datat[i] * self.enf + (i + 1)
# axr.plot(time_axis, stadata.datar[i], linewidth=0.2, color='black')
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)
# axt.plot(time_axis, stadata.datat[i], linewidth=0.2, color='black')
self.axt.fill_between(self.stadata.time_axis, datat, bound + i + 1, where=datat > i+1, facecolor='red',
alpha=0.7)
self.axt.fill_between(self.stadata.time_axis, datat, bound + i + 1, where=datat < i+1, facecolor='blue',
alpha=0.7)
datar = np.mean(self.stadata.data_prime, axis=0)
datar /= np.max(datar)
datat = np.mean(self.stadata.datat, axis=0)
datat /= np.max(datar)
self.axr_sum.fill_between(self.stadata.time_axis, datar, bound, where=datar > 0, facecolor='red', alpha=0.7)
self.axr_sum.fill_between(self.stadata.time_axis, datar, bound, where=datar < 0, facecolor='blue', alpha=0.7)
self.axt_sum.fill_between(self.stadata.time_axis, datat, bound, where=datat > 0, facecolor='red', alpha=0.7)
self.axt_sum.fill_between(self.stadata.time_axis, datat, bound, where=datat < 0, facecolor='blue', alpha=0.7)
self.axb.scatter(self.stadata.bazi, np.arange(self.stadata.ev_num) + 1, s=7)


def set_fig(self):
"""Set figure
"""
y_range = np.arange(self.stadata.ev_num) + 1
x_range = np.arange(0, self.xlim[1]+2, 2)
space = 2

# 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'))

# set axr_sum
self.axr_sum.set_title('{} components ({})'.format(self.stadata.comp, self.stadata.staname), fontsize=16)
self.axr_sum.set_xlim(self.xlim)
self.axr_sum.set_xticks(x_range)
self.axr_sum.set_xticklabels([])
self.axr_sum.set_ylim(-0.5, 1.25)
self.axr_sum.set_yticks([0.375])
self.axr_sum.set_yticklabels(['Sum'], fontsize=8)
self.axr_sum.tick_params(axis='y', left=False)
self.axr_sum.add_line(Line2D([0, 0], self.axr_sum.get_ylim(), color='black'))

# set axt
self.axt.set_xlim(self.xlim)
self.axt.set_xticks(x_range)
self.axt.set_xticklabels(x_range, fontsize=8)
self.axt.set_ylim(0, self.stadata.ev_num + space)
self.axt.set_yticks(y_range)
bazi = ['{:.1f}'.format(ba) for ba in self.stadata.bazi]
self.axt.set_yticklabels(bazi, fontsize=5)
self.axt.set_xlabel('Time after P (s)', fontsize=13)
self.axt.set_ylabel(r'Back-azimuth ($\circ$)', fontsize=13)
self.axt.add_line(Line2D([0, 0], self.axt.get_ylim(), color='black'))

# set axt_sum
self.axt_sum.set_title('T components ({})'.format(self.stadata.staname), fontsize=16)
self.axt_sum.set_xlim(self.xlim)
self.axt_sum.set_xticks(x_range)
self.axt_sum.set_xticklabels([])
self.axt_sum.set_ylim(-0.5, 1.25)
self.axt_sum.set_yticks([0.375])
self.axt_sum.set_yticklabels(['Sum'], fontsize=8)
self.axt_sum.tick_params(axis='y', left=False)
self.axt_sum.add_line(Line2D([0, 0], self.axt_sum.get_ylim(), color='black'))

# 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)


def plot(self, out_path=None, outformat='g', show=False):
"""Plot PRFs with R and T components
:param rfsta: Path to PRFs
:type rfsta: seispy.rfcorrect.RFStation
:param enf: The enlarge factor, defaults to 3
: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, avialible 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
"""
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 show:
plt.show()

def plotrt(rfsta, out_path='./', xlim=[-2, 30], key='bazi', enf=6, outformat='g', show=False):
"""Plot PRFs with R and T components
:param rfsta: Path to PRFs
:type rfsta: seispy.rfcorrect.RFStation
:param enf: The enlarge factor, defaults to 3
: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
Expand All @@ -140,22 +175,14 @@ def plotrt(rfsta, enf=3, out_path='./', key='bazi', outformat='g', xlim=[-2, 30]
: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, 30]
:type xlim: list, optional
:param show: show figure
:type show: bool
"""
h, axr, axt, axb, axr_sum, axt_sum = init_figure()
rfsta.sort(key)
plot_waves(axr, axt, axb, axr_sum, axt_sum, rfsta, enf=enf)
set_fig(axr, axt, axb, axr_sum, axt_sum, rfsta, rfsta.staname, xmin=xlim[0], xmax=xlim[1], comp=rfsta.comp)
if outformat is None and not show:
return h
elif outformat == 'g':
h.savefig(join(out_path, rfsta.staname+'_RT_{}order_{:.1f}.png'.format(key, rfsta.f0[0])),
dpi=400, bbox_inches='tight')
elif outformat == 'f':
h.savefig(join(out_path, rfsta.staname+'_RT_{}order_{:.1f}.pdf'.format(key, rfsta.f0[0])),
format='pdf', bbox_inches='tight')
if show:
plt.show()
return h
prt = PlotRT(rfsta, enf=enf, xlim=xlim, key=key)
prt.plot(out_path, outformat, show)
return prt


if __name__ == '__main__':
Expand Down
Loading

0 comments on commit de2f9de

Please sign in to comment.