Wiggle Trace のプロット方法

地震データなどの波形処理をしていると、観測記録をWiggle Traceで表示することが多いので、
そのプロット方法を考えてみた。

##コード

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection, PolyCollection

nt=200 #sample 数
data=np.sin(np.linspace(0,10*np.pi,nt)) #今回はsin波をプロットする
data=np.asarray([data]*10)
amp=0.5 #プロット時の振幅のスケーリング

y=np.arange(nt)
xs_all=[]
xs_pos=[]
for i in range(len(data)):
    pos=data[i].copy()
    pos[pos<0.0]=0 #極性が正のデータ
    xs_pos.append(pos/data[i].max()*amp+i) #トレース毎にスケーリング
    xs_all.append(data[i]/data[i].max()*amp+i)
    
xs_pos=np.asarray(xs_pos)
xs_all=np.asarray(xs_all)

segs_all = np.zeros((data.shape[0],data.shape[1],  2))
segs_all[:, :, 0] = xs_all
segs_all[:, :, 1] = y

segs_pos = np.zeros((data.shape[0],data.shape[1],  2))
segs_pos[:, :, 0] = xs_pos
segs_pos[:, :, 1] = y
fig, ax = plt.subplots(figsize=(10,10))
#極性が正の部分は塗りつぶし
line_segments_pos = PolyCollection(segs_pos, linewidths=0,
                               facecolors='k',edgecolors='k',linestyle='solid')
#極性が負の部分は線だけ引く
line_segments_all = LineCollection(segs_all, color='k',linewidths=0.5,linestyle='solid')
ax.add_collection(line_segments_all)
ax.add_collection(line_segments_pos)
ax.set_xlim(xs_all.min()-0.5, xs_all.max()+0.5)
ax.set_ylim(y.max(), y.min())
ax.set_xlabel('Channel No')
ax.set_ylabel('Sample Number')

##結果

思ったようにできた。
matplotのCollectionを使うことで、for文より早くプロットできるようにしたことが工夫。