# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: seismology
#     language: python
#     name: seismology
# ---

# +
# %matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt

from os.path import join
from obspy.taup import TauPyModel
import obspy
# -

model = TauPyModel(model="ak135")


# +
def plot_single(network="AH",event="201111080259120",station="ANQ",minfreq=1/50,maxfreq=1/10,t1=-200,t2=600,datadir_url="/mnt/research/seismolab2/japan_slab/data/upload_temp_ziyi/extract_cea/extract/processed",syncdir_url="/mnt/research/seismolab2/japan_slab/data/upload_temp_ziyi/extract_cea/extract/sync/nosmooth"):
    # urls
    z_url=join(datadir_url,event,network+"."+station+"*Z")
    r_url=join(datadir_url,event,network+"."+station+"*R")
    t_url=join(datadir_url,event,network+"."+station+"*T")
    syncz_url=join(syncdir_url,event,f"{network}.{station}.MXZ.sem.sac")
    syncr_url=join(syncdir_url,event,f"{network}.{station}.MXR.sem.sac")
    synct_url=join(syncdir_url,event,f"{network}.{station}.MXT.sem.sac")
    # read 
    try:    
        z=obspy.read(z_url)
        r=obspy.read(r_url)
        t=obspy.read(t_url)
        syncz=obspy.read(syncz_url)
        syncr=obspy.read(syncr_url)
        synct=obspy.read(synct_url)
    except:
        return False,1
    # filter
    syncz[0].data=syncz[0].data*1e9
    syncr[0].data=syncr[0].data*1e9
    synct[0].data=synct[0].data*1e9
    z.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    r.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    t.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    syncz.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    syncr.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    synct.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    # fig
    fig = plt.figure(figsize=(9,16))
    axz = fig.add_subplot(3, 1, 1)
    axr = fig.add_subplot(3, 1, 2)
    axt = fig.add_subplot(3, 1, 3)
    
    for index,(wavez,waver,wavet,syncwavez,syncwaver,syncwavet) in enumerate(zip(z,r,t,syncz,syncr,synct)):
        try:
            assert ((wavez.stats.station==waver.stats.station) and (wavez.stats.station==wavet.stats.station))
        except:
            return False,2
            
        wavez.trim(wavez.stats.starttime+wavez.stats.sac.o,wavez.stats.endtime)
        waver.trim(waver.stats.starttime+waver.stats.sac.o,waver.stats.endtime)
        wavet.trim(wavet.stats.starttime+wavet.stats.sac.o,wavet.stats.endtime)
        syncwavez.trim(syncwavez.stats.starttime+syncwavez.stats.sac.o,syncwavez.stats.endtime)
        syncwaver.trim(syncwaver.stats.starttime+syncwaver.stats.sac.o,syncwaver.stats.endtime)
        syncwavet.trim(syncwavet.stats.starttime+syncwavet.stats.sac.o,syncwavet.stats.endtime)
        
        plot_arrivals={
            "first_p":None,
            "first_s":None,
            "pP":None,
            "PcP":None
        }

        arrivals = model.get_travel_times(source_depth_in_km=wavez.stats.sac.evdp,distance_in_degree=wavez.stats.sac.gcarc)
        for arrival in arrivals:
            if((arrival.name=='p') or (arrival.name=="P")):
                plot_arrivals["first_p"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='s') or (arrival.name=="S")):
                plot_arrivals["first_s"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='pP')):
                plot_arrivals["pP"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='PcP')):
                plot_arrivals["PcP"]=arrival.time
                break
        
        t1_utc=wavez.stats.starttime+plot_arrivals["first_p"]+t1
        end_utc=wavez.stats.starttime+plot_arrivals["first_p"]+t2
        t0_utc=wavez.stats.starttime
        print("starttime",wavez.stats.starttime)
        print("p arrival",wavez.stats.starttime+plot_arrivals["first_p"])
        print("s arrival",wavez.stats.starttime+plot_arrivals["first_s"])
        if(t1_utc<t0_utc):
            start_utc=t0_utc
        else:
            start_utc=t1_utc
            
        wavez.trim(start_utc,end_utc)
        waver.trim(start_utc,end_utc)
        wavet.trim(start_utc,end_utc)
        syncwavez.trim(start_utc,end_utc)
        syncwaver.trim(start_utc,end_utc)
        syncwavet.trim(start_utc,end_utc)
        
        if(t1_utc<t0_utc):
            plotx=np.linspace(-plot_arrivals["first_p"],t2,wavez.stats.npts)
            syncplotx=np.linspace(-plot_arrivals["first_p"],t2,syncwavez.stats.npts)
        else:
            plotx=np.linspace(t1,t2,wavez.stats.npts)   
            syncplotx=np.linspace(t1,t2,syncwavez.stats.npts)
        
        ploty_z=wavez.data
        ploty_r=waver.data
        ploty_t=wavet.data
        syncploty_z=syncwavez.data
        syncploty_r=syncwaver.data
        syncploty_t=syncwavet.data
        
        axz.plot(plotx,ploty_z,"k")
        axr.plot(plotx,ploty_r,"k")
        axt.plot(plotx,ploty_t,"k")
        axz.plot(syncplotx,syncploty_z,"r")
        axr.plot(syncplotx,syncploty_r,"r")
        axt.plot(syncplotx,syncploty_t,"r")
        
        axz.scatter(0,wavez.stats.sac.gcarc,color="r",label="first p")
        axr.scatter(0,waver.stats.sac.gcarc,color="r",label="first p")
        axt.scatter(0,wavet.stats.sac.gcarc,color="r",label="first p")
        
        if(plot_arrivals["first_s"]-plot_arrivals["first_p"]<t2):
            axz.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="b",label="first s")
            axr.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="b",label="first s")
            axt.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="b",label="first s")

        try:
            if(plot_arrivals["pP"]-plot_arrivals["first_p"]<t2):
                axz.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="g",label="pP")
                axr.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="g",label="pP")
                axt.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="g",label="pP")
            if(plot_arrivals["PcP"]-plot_arrivals["first_p"]<t2):
                axz.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="purple",label="PnP")
                axr.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="purple",label="PnP")
                axt.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="purple",label="PnP")
        except:
            pass
    
    axz.set_xlabel("t(s)")
    axz.set_ylabel("amplitude")
    axz.set_title(f"{event} {network} {station} z")
    axz.legend()
    axr.set_xlabel("t(s)")
    axr.set_ylabel("amplitude")
    axr.set_title(f"{event} {network} {station} r")
    axr.legend()
    axt.set_xlabel("t(s)")
    axt.set_ylabel("amplitude")
    axt.set_title(f"{event} {network} {station} t")
    axt.legend()
    plt.show()
    
#     print(wavez.stats.sac.gcarc)


# -

bbq=plot_single(station="TCH",network="SD",event="201206092100193")


def plot_multiple(network="AH",event="201111080259120",scale=0.5,minfreq=1/50,maxfreq=1/10,t1=-20,t2=600,textinfo={"x":700,"fontsize":1},datadir_url="/mnt/research/seismolab2/japan_slab/data/upload_temp_ziyi/extract_cea/extract/processed",syncdir_url="/mnt/research/seismolab2/japan_slab/data/upload_temp_ziyi/extract_cea/extract/sync/nosmooth"):
    # urls
    z_url=join(datadir_url,event,network+"*Z")
    r_url=join(datadir_url,event,network+"*R")
    t_url=join(datadir_url,event,network+"*T")
    # read 
    try:    
        z=obspy.read(z_url)
        r=obspy.read(r_url)
        t=obspy.read(t_url)
    except:
        return False,1
    # filter
    z.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    r.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    t.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
    
    # fig
    fig = plt.figure(figsize=(9,15))
    axz = fig.add_subplot(3, 1, 1)
    axr = fig.add_subplot(3, 1, 2)
    axt = fig.add_subplot(3, 1, 3)
    outputinfo=[]
    
    for index,(wavez,waver,wavet) in enumerate(zip(z,r,t)):
        try:
            assert ((wavez.stats.station==waver.stats.station) and (wavez.stats.station==wavet.stats.station))
        except:
            return False,2
        
        # readin sync data
        syncwavez=obspy.read(join(syncdir_url,event,f"{wavez.stats.network}.{wavez.stats.station}.MXZ.sem.sac"))[0]
        syncwaver=obspy.read(join(syncdir_url,event,f"{wavez.stats.network}.{wavez.stats.station}.MXR.sem.sac"))[0]
        syncwavet=obspy.read(join(syncdir_url,event,f"{wavez.stats.network}.{wavez.stats.station}.MXT.sem.sac"))[0]
        syncwavez.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
        syncwaver.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
        syncwavet.filter("bandpass",freqmin=minfreq,freqmax=maxfreq)
            
        wavez.trim(wavez.stats.starttime+wavez.stats.sac.o,wavez.stats.endtime)
        waver.trim(waver.stats.starttime+waver.stats.sac.o,waver.stats.endtime)
        wavet.trim(wavet.stats.starttime+wavet.stats.sac.o,wavet.stats.endtime)
        syncwavez.trim(syncwavez.stats.starttime+syncwavez.stats.sac.o,syncwavez.stats.endtime)
        syncwaver.trim(syncwaver.stats.starttime+syncwaver.stats.sac.o,syncwaver.stats.endtime)
        syncwavet.trim(syncwavet.stats.starttime+syncwavet.stats.sac.o,syncwavet.stats.endtime)
        
        plot_arrivals={
            "first_p":None,
            "first_s":None,
            "pP":None,
            "PcP":None
        }

        arrivals = model.get_travel_times(source_depth_in_km=wavez.stats.sac.evdp,distance_in_degree=wavez.stats.sac.gcarc)
        for arrival in arrivals:
            if((arrival.name=='p') or (arrival.name=="P")):
                plot_arrivals["first_p"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='s') or (arrival.name=="S")):
                plot_arrivals["first_s"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='pP')):
                plot_arrivals["pP"]=arrival.time
                break

        for arrival in arrivals:
            if((arrival.name=='PcP')):
                plot_arrivals["PcP"]=arrival.time
                break
        
        t1_utc=wavez.stats.starttime+plot_arrivals["first_p"]+t1
        end_utc=wavez.stats.starttime+plot_arrivals["first_p"]+t2
        t0_utc=wavez.stats.starttime
        if(t1_utc<t0_utc):
            start_utc=t0_utc
        else:
            start_utc=t1_utc
            
        wavez.trim(start_utc,end_utc)
        waver.trim(start_utc,end_utc)
        wavet.trim(start_utc,end_utc)
        syncwavez.trim(start_utc,end_utc)
        syncwaver.trim(start_utc,end_utc)
        syncwavet.trim(start_utc,end_utc)
        
        syncwavez.data=syncwavez.data*1e9
        syncwaver.data=syncwaver.data*1e9
        syncwavet.data=syncwavet.data*1e9

        wave_all=obspy.Stream()
        wave_all=wave_all+wavez
        wave_all=wave_all+waver
        wave_all=wave_all+wavet
        wave_all=wave_all+syncwavez
        wave_all=wave_all+syncwaver
        wave_all=wave_all+syncwavet
        wave_all.normalize(global_max=True)
        
        if(t1_utc<t0_utc):
            plotx=np.linspace(-plot_arrivals["first_p"],t2,wavez.stats.npts)
            syncplotx=np.linspace(-plot_arrivals["first_p"],t2,syncwavez.stats.npts)
        else:
            plotx=np.linspace(t1,t2,wavez.stats.npts)      
            syncplotx=np.linspace(t1,t2,syncwavez.stats.npts)
        
        ploty_z=wavez.data*scale+wavez.stats.sac.gcarc
        ploty_r=waver.data*scale+wavez.stats.sac.gcarc
        ploty_t=wavet.data*scale+wavez.stats.sac.gcarc
        syncploty_z=syncwavez.data*scale+wavez.stats.sac.gcarc
        syncploty_r=syncwaver.data*scale+wavez.stats.sac.gcarc
        syncploty_t=syncwavet.data*scale+wavez.stats.sac.gcarc
        
        axz.plot(plotx,ploty_z,"k")
        axr.plot(plotx,ploty_r,"k")
        axt.plot(plotx,ploty_t,"k")
        axz.plot(syncplotx,syncploty_z,"r")
        axr.plot(syncplotx,syncploty_r,"r")
        axt.plot(syncplotx,syncploty_t,"r")
        axz.text(textinfo["x"],wavez.stats.sac.gcarc,wavez.stats.station,fontsize=textinfo["fontsize"])
        axr.text(textinfo["x"],wavez.stats.sac.gcarc,wavez.stats.station,fontsize=textinfo["fontsize"])
        axt.text(textinfo["x"],wavez.stats.sac.gcarc,wavez.stats.station,fontsize=textinfo["fontsize"])
        outputinfo.append((wavez.stats.sac.gcarc,wavez.stats.station))
        
        axz.scatter(0,wavez.stats.sac.gcarc,color="r")
        axr.scatter(0,waver.stats.sac.gcarc,color="r")
        axt.scatter(0,wavet.stats.sac.gcarc,color="r")
        
        if(plot_arrivals["first_s"]-plot_arrivals["first_p"]<t2):
            axz.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="b")
            axr.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="b")
            axt.scatter(plot_arrivals["first_s"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="b")

        try:
            if(plot_arrivals["pP"]-plot_arrivals["first_p"]<t2):
                axz.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="g")
                axr.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="g")
                axt.scatter(plot_arrivals["pP"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="g")
            if(plot_arrivals["PcP"]-plot_arrivals["first_p"]<t2):
                axz.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],wavez.stats.sac.gcarc,color="purple")
                axr.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],waver.stats.sac.gcarc,color="purple")
                axt.scatter(plot_arrivals["PcP"]-plot_arrivals["first_p"],wavet.stats.sac.gcarc,color="purple")
        except:
            pass
    
    axz.set_xlabel("t(s)")
    axz.set_ylabel("distance(degree)")
    axz.set_title(f"{event} {network} z")
    axr.set_xlabel("t(s)")
    axr.set_ylabel("distance(degree)")
    axr.set_title(f"{event} {network} r")
    axt.set_xlabel("t(s)")
    axt.set_ylabel("distance(degree)")
    axt.set_title(f"{event} {network} t")
    plt.show()
    print(sorted(outputinfo,key=lambda x:x[0]))


plot_multiple(network="SD",event="201505021650497",scale=0.2,t1=-100,t2=600,textinfo={"x":700,"fontsize":0})

StarcoderdataJupyterScriptsDedupFiltered