<jupyter_start><jupyter_text>Demo2 of CubeNetA simple example of how to use CubeNet to directly pick P/S arrivals from a given cube<jupyter_code>import sys; sys.path.append('../../'); # To import CubeNet
import torch
import CubeNet

import numpy as np
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>Step 1: Load Seismograms and Stations<jupyter_code># Load Prepared waveform file
data_dict = np.load('Cube.npz');
data = data_dict['cube'];
# Dimension of cube data: batch size, channel (X/Y/Z), dx, dy, sampling points
print('Shape of cube data:',data.shape);
# Look at seismograms (Z component)
plt.figure(figsize=(14,16))
counter = 0; ytick_lst = [];
for i in range(data.shape[2]):
    for ii in range(data.shape[3]):
        ytick_lst.append('(' + str(i) + ',' + str(ii) + ')');
        plt.plot(counter + data[0,2,i,ii,:]/(max(data[0,2,i,ii,:]) + 1e-12) , 'k' ,lw=0.8);
        counter += 1;
plt.xlim([0,data.shape[-1]]);plt.xticks(fontsize=14);plt.xlabel('Sampling Pts',fontsize=16);
plt.ylim([-1,data.shape[2]**2+1]);plt.yticks(fontsize=14);plt.ylabel('Trace ID',fontsize=16);
plt.yticks(np.arange(0,64,3,dtype=int),ytick_lst[slice(0,64,3)]);<jupyter_output><empty_output><jupyter_text>Step 2: Put them in given API<jupyter_code># convert data to torch Tensor
torch_data = torch.Tensor(data);
# Using CubeNet to predict seismic phases
net = CubeNet.network; net.eval();

# convert data to cpu/gpu
# net = net.cuda(); torch_data = torch_data.cuda();
net = net.cpu(); torch_data = torch_data.cpu();

# prediction
pick_result = net(torch_data);
pick_result = pick_result.detach().cpu().numpy()

print('Shape of pick_result: ',pick_result.shape)
# Create a funtion to plot the picking results

def wavenorm(data,scale=1):
    _range = np.max(data) - np.min(data);
    if _range != 0:
        return scale*(2*((data - np.min(data)) / _range)-1);
    if _range == 0:
        return np.zeros(data.shape);
    
def pickplot1D(Data,P_result,S_result,chose_TraceID,plot_comp='Z',threshold=0.3):
    if plot_comp == 'X':
        plot_data = wavenorm(Data[0,0,chose_TraceID[0],chose_TraceID[1],:],scale=0.8);
    elif plot_comp == 'Y':
        plot_data = wavenorm(Data[0,1,chose_TraceID[0],chose_TraceID[1],:],scale=0.8);
    else:
        plot_data = wavenorm(Data[0,2,chose_TraceID[0],chose_TraceID[1],:],scale=0.8);
        
    if np.max(P_result[chose_TraceID[0],chose_TraceID[1],:]) > threshold:
        P_picks = np.argmax(P_result[chose_TraceID[0],chose_TraceID[1],:]);
    else:
        P_picks = None;
        
    if np.max(S_result[chose_TraceID[0],chose_TraceID[1],:]) > threshold:
        S_picks = np.argmax(S_result[chose_TraceID[0],chose_TraceID[1],:]);
    else:
        S_picks = None;
    
    plt.figure(figsize=(18,6))
    plt.subplot(3,1,1);
    plt.plot(plot_data,'k',linewidth=0.8);
    if P_picks:
        plt.plot([P_picks,P_picks],[-1,1],'r--');
    if S_picks:
        plt.plot([S_picks,S_picks],[-1,1],'b--');
    plt.xticks([]);plt.yticks(fontsize=14);
    plt.subplot(3,1,2);
    plt.plot(P_result[chose_TraceID[0],chose_TraceID[1],:],'r--',linewidth=0.8);
    plt.xticks([]);plt.yticks(fontsize=14);
    plt.ylim([-.1,1]);
    plt.subplot(3,1,3);
    plt.plot(S_result[chose_TraceID[0],chose_TraceID[1],:],'b--',linewidth=0.8);
    plt.ylim([-.1,1]);
    plt.xticks(fontsize=14);plt.yticks(fontsize=14);
    plt.xlabel('Sampling Points',fontsize=18);
    
def pickplot2D(Data,P_result,S_result,plot_comp='Z',threshold=0.3):
    plt.figure(figsize=(14,16));
    tick_len = 0.8; tick_w = 1.4;
    counter = 0; ytick_lst = [];
    for i in range(Data.shape[2]):
        for ii in range(Data.shape[3]):
            ytick_lst.append('(' + str(i) + ',' + str(ii) + ')');
            if plot_comp == 'X':
                plot_data = wavenorm(Data[0,0,i,ii,:],scale=0.8);
            elif plot_comp == 'Y':
                plot_data = wavenorm(Data[0,1,i,ii,:],scale=0.8);
            else:
                plot_data = wavenorm(Data[0,2,i,ii,:],scale=0.8);
            plot_data = plot_data - np.average(plot_data);
            plt.plot(plot_data+counter,'k',linewidth=0.7);
            P_picks = np.argmax(P_result[i,ii,:]); S_picks = np.argmax(S_result[i,ii,:]);
            if P_result[i,ii,P_picks] > threshold:
                plt.plot([P_picks,P_picks],[-tick_len+counter,tick_len+counter],'r',linewidth=tick_w);
            if S_result[i,ii,S_picks] > threshold:
                plt.plot([S_picks,S_picks],[-tick_len+counter,tick_len+counter],'b',linewidth=tick_w);
            counter += 1;
    plt.xticks(fontsize=14);plt.yticks(np.arange(0,64,3,dtype=int),ytick_lst[slice(0,64,3)],fontsize=14);
    plt.ylim([-0.5,Data.shape[2]**2+0.5]); plt.xlim([0,Data.shape[-1]]);
    plt.xlabel('Sampling Points',fontsize=18);
    plt.ylabel('Trace ID',fontsize=18);
# plot picking result of different traces
chose_TraceID = (3,4);
P_result = pick_result[0,0,:,:,:];
S_result = pick_result[0,1,:,:,:];
pickplot1D(data,P_result,S_result,chose_TraceID,plot_comp='X');
# Pick Result in total
pickplot2D(data,P_result,S_result,plot_comp='Z');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1889301779777039, 'fraction_numerical': 0.03461764130647369, 'mean_word_length': 4.139698492462312, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '387133', 'license': 'MIT', 'n_tokens_mistral': 2125, 'n_tokens_neox': 2061, 'n_words': 293, 'path': 'demo/demo2/run.ipynb', 'repo_name': 'Billy-Chen0327/CubeNet', 'size': '5113'}