decon.py 文件源码

def damped_lstsq(a,b,damping=1.0,plot=False,return_G=False):
    '''
    Gm = d
       G : discrete convolution matrix
       m : signal we are trying to recover (receiver function)
       d : the convolved data (signal a)

       m = (G^T G)^(-1)G^T * d
    '''

    #build G
    padding = np.zeros(a.shape[0] - 1, a.dtype)
    first_col = np.r_[a, padding]
    first_row = np.r_[a[0], padding]
    G = toeplitz(first_col, first_row)

    #reshape b
    shape = G.shape
    shape = shape[0]
    len_b = len(b)
    zeros = np.zeros((shape-len_b))
    b = np.append(b,zeros)

    #solve system (use either lsmr or scipy solve)
    #sol = lsmr(G,b,damp=damping)
    sol = np.linalg.solve(G+damping*np.eye(G.shape[0]),b,overwrite_a=False,overwrite_b=False)
    m_est = sol[0]
    rf = m_est

    if plot==True:
       fig,axes = plt.subplots(3,sharex=True)
       axes[0].plot(a)
       axes[1].plot(b)
       axes[2].plot(rf)
       plt.show()

    if return_G==True:
       return G
    else:
       return rf