#%%
import os
import numpy as np
import sys
import time
import multiprocessing as mp
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd
import obspy
import glob
from scipy.signal import iirpeak, tukey, butter, filtfilt, hilbert
from func_disp import *
#Function to calculate dispersion curve through slan-stacking
def CalcDisp_shift_stack(input):
'''
# calculate dispersion using shift and stack (shift in frequency domain)
# Input:
# input - a list of the parameters below:
# f - frequencies
# v - phase velocities
# fs - sample rate
# wavelen - wavelength corresponding to f
# recdist - epicentral distance of the beam center
# beamdist - epicentral distance of the stations in the beam (receiver-source)
# twosided - if the causal and anticausal sides of waveform have not been added, we do it here
# data - waveform, length(station) * npts
# Output:
# disp - dispersion matrix, length(v) * length(f)
'''
nbeamwidth, nepicdist, f, v, fs, wavelen, recdist, beamdist, twosided, data = input
data = data / np.abs(data).max(axis=1)[:, np.newaxis]
nf = len(f)
nv = len(v)
disp = np.empty((nv, nf))
disp[:] = np.nan
for i in range(nf):
newbeam = np.where((np.abs(beamdist) > 30) & (
np.abs(beamdist - recdist) < nbeamwidth * wavelen[i]))[0]
if (recdist>nepicdist[0]*wavelen[i])&(recdist<nepicdist[1]*wavelen[i]):
bb, aa = iirpeak(f[i], 3, fs=fs)
dataf = filtfilt(bb, aa, data[newbeam,:])
if twosided:
halfval=int((dataf.shape[1]-1)/2)
dataf=dataf[:,:(halfval+1)][:,::-1]+dataf[:,halfval:]
for j in range(nv):
nshift = (beamdist[newbeam]-recdist)/v[j]*fs
dataf_shift = np.zeros_like(dataf)
for k in range(len(newbeam)):
dataf_shift[k, :] = np.roll(dataf[k, :], -nshift[k].astype(int))
disp[j, i] = np.linalg.norm(np.sum(dataf_shift, 0)) + 1e-4
disp[:, i] = disp[:, i] / np.max(disp[:, i])
else:
disp[:, i] = np.zeros(nv)
return disp
shots_all=np.load('data/shots_all.npy')
#Consider an example source
i=300
data_shot = shots_all[i,:,500:]+shots_all[i,:,:501][:,::-1]
chan_csv=pd.read_csv('data/Mammoth_south_100km.csv')
chan_use=np.where(chan_csv['status']=='good')[0][5999:]
nch=len(chan_use)//2
lats=np.array(chan_csv[chan_csv['status']=='good'][5999::2]['latitude'])
lons=np.array(chan_csv[chan_csv['status']=='good'][5999::2]['longitude'])
rec_loc=np.array((lons,lats)).T
src_loc = np.array([lons[i],lats[i]])
# epicentral distance
epi_dist = obspy.geodetics.base.degrees2kilometers(obspy.geodetics.base.locations2degrees(
src_loc[1],src_loc[0],rec_loc[:,1],rec_loc[:,0]) ) * 1e3 # in meter
# receiver mutual distance
rec_dist = obspy.geodetics.base.degrees2kilometers(np.array([obspy.geodetics.base.locations2degrees(
rec_loc[x,1],rec_loc[x,0],rec_loc[:,1],rec_loc[:,0]) for x in range(len(rec_loc))] ))* 1e3
#%% visualize to find a suitable frequency range
fs=25
for times in 2**np.arange(5):
fig, ax=plt.subplots(1,figsize=(5,4))
f1, f2 = np.array((0.2,0.5))*times
dataf=filter(data_shot, fs, f1, f2)
vmax = np.max(np.abs(dataf),axis=1)
plt.pcolormesh(np.arange(0,data_shot.shape[0])*2+5999,np.arange(0,data_shot.shape[1])/fs,(dataf/vmax[:,None]).T, vmax=1,vmin=-1,cmap="seismic");plt.colorbar()
plt.ylabel('Time (s)')
plt.xlabel('Channel No.')
plt.ylim((10,0))
plt.xlim((6000,7800))
plt.title('Virtual Source at Receiver %s (BP %s-%s Hz)'%(5999+i*2,f1,f2))
plt.savefig('/home/elibird/200km/figs/virtual_source_ch%s_%s-%s.png'%(5999+i*2,f1,f2))
# plt.pcolormesh(np.arange(800)/fs, epi_dist * np.sign(np.arange(300)-150), dataf,vmax=vmax,vmin=-vmax);plt.colorbar();plt.show()
#%% define parameters
# frequency band for dispersion curves
fs=25
f1, f2 = 0.25, 4.0
# take 100 log-even samples in the desired frequency band
f = np.logspace(np.log10(f1),np.log10(f2),50) # in Hz
# To define the segment length for beamforming, need wavelength for each frequency.
# wavelength = velocity / frequency, make sure wavelen and f correspond.
# wavelength assuming velocity is 2000 m/s
wavelen = 2000 * np.ones_like(f) / f # in meter
# we require subarray size at least 800m (40 channels)
wavelen[wavelen<800] = 800
# Set beam width as some multiple of the wavelength.
n_beam_width = 0.5
# at least 1.5 times wavelength for far-field assumption, less than 3 wavelengths
n_epic_dist = (1.5,3.0)
# phase velocity range for beamforming
v = np.logspace(np.log10(250),np.log10(4000),80) # velocity for searching
# use how many CPUs for parallel
npar = 50
# output path
outputdir = 'data/dcurves'
if True:
chan_csv=pd.read_csv('data/Mammoth_south_100km.csv')
chan_use=np.where(chan_csv['status']=='good')[0][5999:]
nch=len(chan_use)//2
lats=np.array(chan_csv[chan_csv['status']=='good'][5999::2]['latitude'])
lons=np.array(chan_csv[chan_csv['status']=='good'][5999::2]['longitude'])
rec_loc=np.array((lons,lats)).T
fh = open('disp_progress.txt', 'w')
import tqdm
for i in tqdm.tqdm(np.arange(0,nch),file=fh):
# epicentral distance
data_shot = shots_all[i]
src_loc = np.array([lons[i],lats[i]])
epi_dist = obspy.geodetics.base.degrees2kilometers(obspy.geodetics.base.locations2degrees(
src_loc[1],src_loc[0],rec_loc[:,1],rec_loc[:,0]) ) * 1e3 # in meter
# receiver mutual distance
rec_dist = obspy.geodetics.base.degrees2kilometers(np.array([obspy.geodetics.base.locations2degrees(
rec_loc[x,1],rec_loc[x,0],rec_loc[:,1],rec_loc[:,0]) for x in range(len(rec_loc))] ))* 1e3
#%%
# Choose the range of epicentral distances to use for the dispersion curves. I set to 800 so that the whole array will be included
idx_rec = np.where(((np.arange(nch)/2-i/2).astype('int')*2)==np.arange(nch)-i)[0]
rec_epic_dist = epi_dist[idx_rec]
beams = [ np.where(rec_dist[x,:] < n_beam_width * np.max(wavelen))[0] for x in idx_rec]
beam_epic_dist = [epi_dist[x] for x in beams] # epicentral distance (distance to the virtual source) for channels in the beam
# ptb, pte = np.maximum(0, rec_epic_dist / groupvmax * fs), np.minimum(data_shot.shape[1], rec_epic_dist / groupvmin * fs)
# since time 0 doesn't mean event origin time, probably better not define a stacking window
data = [ data_shot[beams[x], :] for x in range(len(beams))]
twosided=True
inputs = [[n_beam_width, n_epic_dist, f, v, fs, wavelen, rec_epic_dist[x], beam_epic_dist[x], twosided, data[x]] for x in range(len(beams))]
# slow, so need parallel here
t1 =time.time()
pool = mp.Pool(npar)
res = pool.map(CalcDisp_shift_stack, inputs)
pool.close()
pool.join()
np.savez(os.path.join(outputdir,'Shot_'+str(i)+'.npz'), rec=idx_rec, res=res, epi_dist=epi_dist, freq=f, vels=v)
fh.close()
idx_rec = np.where(((np.arange(nch)/2-i/2).astype('int')*2)==np.arange(nch)-i)[0]
# %% visualize and check results
#%% frequency-velocity image
rec_check = 410 # which receiver to visualize
irec = np.argmin(np.abs(idx_rec - rec_check))
fvimage = res[irec]
vout = extr_disp([f, v, fvimage])
plt.pcolormesh(f,v,fvimage,cmap='Spectral_r');
# plt.plot(f, vout, 'ko-')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Phase velocity (m/s)')
plt.gca().set_xscale('log')
plt.gca().set_yscale('log')
plt.colorbar()
plt.clim(0,1.0)
plt.xlim(0.4,3.5)
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[10], line 6 4 rec_check = 410 # which receiver to visualize 5 irec = np.argmin(np.abs(idx_rec - rec_check)) ----> 6 fvimage = res[irec] 7 vout = extr_disp([f, v, fvimage]) 8 plt.pcolormesh(f,v,fvimage,cmap='Spectral_r'); NameError: name 'res' is not defined