# -*- coding: utf-8 -*-
# Copyright 2019- Weijia Sun, MIT license
"""
stacking
"""
import os
import glob
import numpy as np
from scipy.signal import hilbert
import commentjson
from obspy import read, Stream
from acc.io import _load_json
from tqdm import tqdm
from functools import partial
from multiprocessing import Pool
import multiprocessing
import logging
logging.basicConfig(level=logging.NOTSET, format='%(asctime)s %(filename)s[line:%(lineno)d] %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
filename="acc.processing.log", filemode="a"
)
[docs]def linear_stack(stream, normalize=True):
if normalize:
stream.normalize(global_max=False)
stk = np.zeros_like(stream[0].data)
for tr in stream:
stk += tr.data
stk /= len(stream)
tr = stream[0]
tr.data = np.copy(stk)
tr.stats.starttime = 0
tr.stats.update({"number_of_ray": len(stream), "stack_method":{"linear"}})
return tr
[docs]def pws_stack(stream, power=2, normalize=True):
if normalize:
stream.normalize(global_max=False)
# safe guard to be a positive integer number
power = int(abs(power))
tr_lin = linear_stack(stream, normalize)
if power == 0:
return tr_lin
phase = np.zeros_like(stream[0].data)
for tr in stream:
# hilbert transform
hilb = hilbert(tr.data)
phase += np.abs(hilb / np.abs(hilb))
phase /= len(stream)
power = int(power)
stk = tr_lin.data * np.power(phase, power)
tr = stream[0]
tr.data = np.copy(stk)
tr.stats.starttime = 0
tr.stats.update({"number_of_ray": len(stream), "stack_method":{"PWS", power}})
return tr
[docs]def stack(jsonfile):
kwargs = _load_json(jsonfile)
path = kwargs["io"]["outpath"]
stations = glob.glob(path + "/1_results/*")
njobs = kwargs["njobs"]
do_work = partial(_stack, **kwargs)
numbers = []
if njobs == 1:
logging.info('do work sequential (%d cores)', njobs)
for sta in tqdm(stations, total=len(stations)):
num = do_work(sta)
numbers.append(num)
else:
logging.info('do work parallel (%d cores)', njobs)
pool = multiprocessing.Pool(njobs)
for num in tqdm(pool.imap_unordered(do_work, stations), total=len(stations)):
numbers.append(num)
pool.close()
pool.join()
[docs]def _stack(path, **kwargs):
files = glob.glob(path + "/*")
# stations = glob.glob(path + "/*")
# print(stations)
# for sta in stations:
# files = glob.glob(sta + "/*")
st = Stream()
for file in files:
try:
tr = read(file)[0]
except:
continue
st.append(tr)
# get channels
channels = []
for tr in st:
chn = tr.stats.channel
if chn not in channels:
channels.append(chn)
method = kwargs["stack"]["method"]
power = kwargs["stack"]["power"]
outpath = kwargs["io"]["outpath"] + "/1a_stack"
try:
os.makedirs(outpath)
except:
pass
stream = st.copy()
for chn in channels:
st = stream.select(channel=chn)
if method == "linear":
tr = linear_stack(st, normalize=True)
elif method == "PWS":
tr = pws_stack(st, power, normalize=True)
elif method == "bootstrap_linear" or method == "bootstrap_PWS":
# note tr here is a stream containing two traces, mean and std
tr = _bootstrap(st, normalize=True, **kwargs)
filen = outpath + "/" + tr.id + "_%s.pkl" % method
tr.write(filename=filen, format="PICKLE")
return 0
[docs]def _bootstrap(st, normalize=True, **kwargs):
if normalize:
stream.normalize(global_max=False)
# generate randoma number
ntrace = len(st)
options = kwargs["stack"]
n_iter = options["n_iter"]
percentage = options["percentage"]
seed = kwargs["seed"]
rand_list = _gen_random_number(low=0, high=ntrace,
n_iter=n_iter, percentage=percentage, seed=seed)
# stacking
method = kwargs["stack"]["method"]
power = kwargs["stack"]["power"]
# method is bootstrap_linear or bootstrap_PWS
method = method.split("_")[1]
npts = st[0].stats.npts
boot = np.zeros([n_iter, npts])
m = 0
for r in rand_list:
# random stream
st_rand = Stream()
for i in r:
tr = st[i]
st_rand.append(tr)
# stack
if method == "linear":
tr_stk = linear_stack(st_rand, normalize=True)
elif method == "PWS":
tr_stk = pws_stack(st_rand, power, normalize=True)
boot[i] = np.copy(tr_stk.data)
# mean value
tr_mean = st[0].copy()
tr_std = st[0].copy()
tr_mean.data = np.mean(boot, axis=0)
tr_mean.stats.update({"statistics":"mean"})
tr_std.data = np.std(boot, axis=0)
tr_st.stats.update({"statistics":"std"})
st2 = Stream()
st2.append(tr_mean)
st2.append(tr_std)
return st2
[docs]def _gen_random_number(low=0, high=100, n_iter=100, percentage=0.9, seed=59):
np.random.seed(seed)
n_trace = int(percentage * high)
rand_list = []
for i in range(n_iter):
r = np.random.randint(low=high, high=high)
rand_list.append(r)
# reset random seed. use the multiplication of the first nonzero value
# generated random number from the previous random seed as the new seed
s = 1
counter = 0
for k in r:
if k != 0:
s *= k
if counter > 5:
break
counter += 1
np.random.seed(seed=s)
return rand_list