DASPy-toolbox 1.0.0__py3-none-any.whl

daspy/advanced_tools/strain2vel.py

@@ -0,0 +1,245 @@
+# Purpose: Convert strain rate data to velocity
+# Author: Minzhe Hu
+# Date: 2024.6.8
+# Email: hmz2018@mail.ustc.edu.cn
+import numpy as np
+from numpy.fft import irfft2, ifftshift
+from scipy.signal import hilbert
+from daspy.basic_tools.freqattributes import next_pow_2, fk_transform
+from daspy.basic_tools.preprocessing import padding, cosine_taper
+from daspy.basic_tools.filter import bandpass
+from daspy.advanced_tools.fdct import fdct_wrapping, ifdct_wrapping
+from daspy.advanced_tools.denoising import _velocity_bin
+from daspy.advanced_tools.decomposition import fk_fan_mask
+
+
+def fk_rescaling(data, dx, fs, taper=(0.02, 0.05), pad='default', fmax=None,
+                 kmin=(1 / 2000, 1 / 3000), vmax=(15000, 30000), edge=0.2,
+                 turning=None, verbose=False):
+    """
+    Convert strain/strain rate to velocity/acceleration by fk rescaling.
+
+    :param data: numpy.ndarray. Data to do fk rescaling.
+    :param dx: Channel interval in m.
+    :param fs: Sampling rate in Hz.
+    :param taper: float or sequence of floats. Each float means decimal
+        percentage of Tukey taper for corresponding dimension (ranging from 0 to
+        1). Default is 0.1 which tapers 5% from the beginning and 5% from the
+        end.
+    :param pad: Pad the data or not. It can be float or sequence of floats. Each
+        float means padding percentage before FFT for corresponding dimension.
+        If set to 0.1 will pad 5% before the beginning and after the end.
+        'default' means pad both dimensions to next power of 2. None or False
+        means don't pad data before or during Fast Fourier Transform.
+    :param fmax, kmin, vmax: float or or sequence of 2 floats. Sequence of 2
+        floats represents the start and end of taper. Setting these parameters
+        can reduce artifacts.
+    :param edge: float. The width of fan mask taper edge.
+    :param turning: Sequence of int. Channel number of turning points.
+    :param verbose: If True, return converted data, f-k spectrum, frequency
+        sequence, wavenumber sequence and f-k mask.
+    :return: Converted data and some variables in the process if verbose==True.
+    """
+    if turning is not None:
+        data_vel = np.zeros_like(data)
+        start_ch = [0, *turning]
+        end_ch = [*turning, len(data)]
+        for (s, e) in zip(start_ch, end_ch):
+            data_vel[s:e] = fk_rescaling(data[s:e], dx, fs, taper=taper,
+                                         pad=pad, fmax=fmax, kmin=kmin,
+                                         vmax=vmax, edge=edge, verbose=False)
+    else:
+        data_tp = cosine_taper(data, taper)
+
+        if pad == 'default':
+            nch, nt = data.shape
+            dn = (next_pow_2(nch) - nch, next_pow_2(nt) - nt)
+            nfft = None
+        elif pad is None or pad is False:
+            dn = 0
+            nfft = None
+        else:
+            dn = np.round(np.array(pad) * data.shape).astype(int)
+            nfft = 'default'
+
+        data_pd = padding(data_tp, dn)
+        nch, nt = data_pd.shape
+
+        fk, f, k = fk_transform(data_pd, dx, fs, taper=taper, nfft=nfft)
+
+        ff = np.tile(f, (len(k), 1))
+        kk = np.tile(k, (len(f), 1)).T
+        vv = - np.divide(ff, kk, out=np.ones_like(ff) * 1e10, where=kk != 0)
+
+        mask = fk_fan_mask(f, k, fmax=fmax, kmin=kmin, vmax=vmax, edge=edge) * vv
+        mask[kk == 0] = 0
+
+        data_vel = irfft2(ifftshift(fk * mask, axes=0)).real[:nch, :nt]
+        data_vel = padding(data_vel, dn, reverse=True)
+
+        if verbose:
+            return data_vel, fk, f, k, mask
+    return data_vel
+
+
+def curvelet_conversion(data, dx, fs, pad=0.3, scale_begin=2, nbscales=None,
+                        nbangles=16, turning=None):
+    """
+    Use curevelet transform to convert strain/strain rate to
+    velocity/acceleration. {Yang et al. , 2023, Geophys. Res. Lett.}
+
+    :param data: numpy.ndarray. Data to convert.
+    :param dx: Channel interval in m.
+    :param fs: Sampling rate in Hz.
+    :param pad: float or sequence of floats. Each float means padding percentage
+        before FFT for corresponding dimension. If set to 0.1 will pad 5% before
+        the beginning and after the end.
+    :param scale_begin: int. The beginning scale to do conversion.
+    :param nbscales: int. Number of scales including the coarsest wavelet level.
+        Default set to ceil(log2(min(M,N)) - 3).
+    :param nbangles: int. Number of angles at the 2nd coarsest level,
+        minimum 8, must be a multiple of 4.
+    :param turning: Sequence of int. Channel number of turning points.
+    :return: numpy.ndarray. Converted data.
+    """
+    if turning is not None:
+        print(1)
+        data_vel = np.zeros_like(data)
+        start_ch = [0, *turning]
+        end_ch = [*turning, len(data)]
+        for (s, e) in zip(start_ch, end_ch):
+            data_vel[s:e] = curvelet_conversion(data[s:e], dx, fs, pad=pad,
+                                                scale_begin=scale_begin,
+                                                nbscales=nbscales,
+                                                nbangles=nbangles, turning=None)
+    else:
+        if pad is None or pad is False:
+            pad = 0
+        dn = np.round(np.array(pad) * data.shape).astype(int)
+        data_pd = padding(data, dn)
+
+        C = fdct_wrapping(data_pd, is_real=True, finest=1, nbscales=nbscales,
+                          nbangles_coarse=nbangles)
+
+        # rescale with velocity
+        np.seterr(divide='ignore')
+        for s in range(0, scale_begin - 1):
+            for w in range(len(C[s])):
+                C[s][w] *= 0
+
+        for s in range(scale_begin - 1, len(C)):
+            nbangles = len(C[s])
+            velocity = _velocity_bin(nbangles, fs, dx)
+            factors = np.mean(velocity, axis=1)
+            for w in range(nbangles):
+                if abs(factors[w]) == np.inf:
+                    factors[w] = abs(velocity[w]).min() * \
+                        np.sign(velocity[w, 0]) * 2
+                C[s][w] *= factors[w]
+
+        data_vel = ifdct_wrapping(C, is_real=True, size=data_pd.shape)
+        data_vel = padding(data_vel, dn, reverse=True)
+
+    return data_vel
+
+
+def slowness(g, dx, fs, slm, sls, swin=2):
+    """
+    Estimate the slowness time series by calculate semblance.
+    {Lior et al., 2021, Solid Earth}
+
+    :param g: 2-dimensional array. time series of adjacent channels used for
+        estimating slowness
+    :param dx: float. Spatical sampling rate (in m)
+    :param fs: float. Sampling rate of records
+    :param slm: float. Slowness x max
+    :param sls: float. Slowness step
+    :param swin: int. Slowness smooth window
+    :return: Sequences of slowness and sembalence.
+    """
+    L = (len(g) - 1) // 2
+    nt = len(g[0])
+    h = np.imag(hilbert(g))
+    grdpnt = round(slm / sls)
+    sem = np.zeros((2 * grdpnt + 1, nt))
+    gap = round(slm * dx * L * fs)
+
+    h_ex = np.zeros((len(g), nt + 2 * gap))
+    h_ex[:, gap:-gap] = h
+    g_ex = np.zeros((len(g), nt + 2 * gap))
+    g_ex[:, gap:-gap] = g
+
+    for i in range(2 * grdpnt + 1):
+        px = (i - grdpnt) * sls
+        if abs(px) < 1e-5:
+            continue
+        gt = np.zeros(g.shape)
+        ht = np.zeros(h.shape)
+        for j in range(-L, L):
+            shift = round(px * j * dx * fs)
+            gt[j + L] = g_ex[j + L, gap + shift:gap + shift + nt]
+            ht[j + L] = h_ex[j + L, gap + shift:gap + shift + nt]
+        sem[i] = (np.sum(gt, axis=0)**2 + np.sum(ht, axis=0)**2) / \
+            np.sum(gt**2 + ht**2, axis=0) / (2 * L + 1)
+    p = (np.argmax(sem, axis=0) - grdpnt) * sls
+    # smooth P
+    for i in range(swin, nt - swin):
+        win = p[i - swin:i + swin + 1]
+        sign = np.sign(sum(np.sign(win)))
+        win = [px for px in win if np.sign(px) == sign]
+        p[i] = np.mean(win)
+
+    return p, sem
+
+
+def slant_stacking(data, dx, fs, L=None, slm=0.01,
+                   sls=0.000125, frqlow=0.1, frqhigh=15, turning=None,
+                   channel='all'):
+    """
+    Convert strain to velocity based on slant-stack.
+
+    :param data: 2-dimensional array. Axis 0 is channel number and axis 1 is
+        time series
+    :param dx: float. Spatical sampling rate (in m)
+    :param L: int. the number of adjacent channels over which slowness is
+        estimated
+    :param slm: float. Slowness x max
+    :param sls: float. slowness step
+    :param freqmin: Pass band low corner frequency.
+    :param freqmax: Pass band high corner frequency.
+    :param turning: Sequence of int. Channel number of turning points.
+    :param channel: int or list or 'all'. convert a certain channel number /
+        certain channel range / all channels.
+    :return: Converted velocity data
+    """
+    if L is None:
+        L = round(50 / dx)
+
+    nch, nt = data.shape
+    if isinstance(channel, str) and channel == 'all':
+        channel = list(range(nch))
+    elif isinstance(channel, int):
+        channel = [channel]
+
+    if turning is not None:
+        data_vel = np.zeros((0, len(data[0])))
+        start_ch = [0, *turning]
+        end_ch = [*turning, len(data)]
+        for (s, e) in zip(start_ch, end_ch):
+            channel_seg = [ch-s for ch in range(s,e) if ch in channel]
+            if len(channel_seg):
+                d_vel = slant_stacking(data[s:e], dx, fs, L=L, slm=slm, sls=sls,
+                                       frqlow=frqlow, frqhigh=frqhigh,
+                                       turning=None, channel=channel_seg)
+                data_vel = np.vstack((data_vel, d_vel))
+    else:
+        data_ex = padding(data, (2 * L, 0))
+        swin = int(max((1 / frqhigh * fs) // 2, 1))
+        data_vel = np.zeros((len(channel), nt))
+        for i, ch in enumerate(channel):
+            p, _ = slowness(data_ex[ch:ch + 2 * L + 1], dx, fs, slm, sls,
+                            swin=swin)
+            data_vel[i] = bandpass(data[ch] / p, fs=fs, freqmin=frqlow,
+                                   freqmax=frqhigh)
+
+    return data_vel

daspy/basic_tools/filter.py

@@ -0,0 +1,257 @@
+# Purpose: Filter the waveform
+# Author: Minzhe Hu
+# Date: 2024.10.16
+# Email: hmz2018@mail.ustc.edu.cn
+# Modified from https://docs.obspy.org/_modules/obspy/signal/filter.html
+import warnings
+import numpy as np
+from scipy.signal import cheb2ord, cheby2, hilbert, iirfilter, zpk2sos, sosfilt
+
+
+def bandpass(data, fs, freqmin, freqmax, corners=4, zi=None, zerophase=True):
+    """
+    Filter data from 'freqmin' to 'freqmax' using Butterworth bandpass filter of
+    'corners' corners.
+
+    :param data: numpy.ndarray. Data to filter.
+    :param fs: Sampling rate in Hz.
+    :param freqmin: Pass band low corner frequency.
+    :param freqmax: Pass band high corner frequency.
+    :param corners: Filter corners / order.
+    :param zi : None, 0, or array_like. Initial conditions for the cascaded
+        filter delays. It is a vector of shape (n_sections, nch, 2). Set to 0 to
+        trigger a output of the final filter delay values.
+    :param zerophase: If True, apply filter once forwards and once backwards.
+        This results in twice the number of corners but zero phase shift in
+        the resulting filtered data. Only valid when zi is None.
+    :return: Filtered data and the final filter delay values (if zi is not
+        None).
+    """
+    if len(data.shape) == 1:
+        data = data[np.newaxis, :]
+    fe = 0.5 * fs
+    low = freqmin / fe
+    high = freqmax / fe
+    # raise for some bad scenarios
+    if high - 1.0 > -1e-6:
+        msg = ('Selected high corner frequency ({}) of bandpass is at or ' +
+               'above Nyquist ({}). Applying a high-pass instead.').format(
+            freqmax, fe)
+        warnings.warn(msg)
+        return highpass(data, freq=freqmin, fs=fs, corners=corners,
+                        zerophase=zerophase)
+    if low > 1:
+        msg = 'Selected low corner frequency is above Nyquist.'
+        raise ValueError(msg)
+    z, p, k = iirfilter(corners, [low, high], btype='band', ftype='butter',
+                        output='zpk')
+    sos = zpk2sos(z, p, k)
+    if zi is None:
+        data_flt = sosfilt(sos, data)
+        if zerophase:
+            data_flt = sosfilt(sos, data_flt[:, ::-1])[:, ::-1]
+        return data_flt
+    elif isinstance(zi, (int, float)):
+        zi = np.ones((sos.shape[0], len(data), 2)) * zi
+    
+    data_flt, zf = sosfilt(sos, data, zi=zi)
+    return data_flt, zf
+
+
+def bandstop(data, fs, freqmin, freqmax, corners=4, zi=None, zerophase=False):
+    """
+    Filter data removing data between frequencies 'freqmin' and 'freqmax' using
+    Butterworth bandstop filter of 'corners' corners.
+
+    :param data: numpy.ndarray. Data to filter.
+    :param fs: Sampling rate in Hz.
+    :param freqmin: Stop band low corner frequency.
+    :param freqmax: Stop band high corner frequency.
+    :param corners: Filter corners / order.
+    :param zi : None, 0, or array_like. Initial conditions for the cascaded
+        filter delays. It is a vector of shape (n_sections, nch, 2). Set to 0 to
+        trigger a output of the final filter delay values.
+    :param zerophase: If True, apply filter once forwards and once backwards.
+        This results in twice the number of corners but zero phase shift in
+        the resulting filtered data. Only valid when zi is None.
+    :return: Filtered data and the final filter delay values (if zi is not
+        None).
+    """
+    if len(data.shape) == 1:
+        data = data[np.newaxis, :]
+    fe = 0.5 * fs
+    low = freqmin / fe
+    high = freqmax / fe
+    # raise for some bad scenarios
+    if high > 1:
+        high = 1.0
+        msg = 'Selected high corner frequency is above Nyquist. Setting ' + \
+              'Nyquist as high corner.'
+        warnings.warn(msg)
+    if low > 1:
+        msg = 'Selected low corner frequency is above Nyquist.'
+        raise ValueError(msg)
+    z, p, k = iirfilter(corners, [low, high],
+                        btype='bandstop', ftype='butter', output='zpk')
+    sos = zpk2sos(z, p, k)
+    if zi is None:
+        data_flt = sosfilt(sos, data)
+        if zerophase:
+            data_flt = sosfilt(sos, data_flt[:, ::-1])[:, ::-1]
+        return data_flt
+    elif isinstance(zi, (int, float)):
+        zi = np.ones((sos.shape[0], len(data), 2)) * zi
+    
+    data_flt, zf = sosfilt(sos, data, zi=zi)
+    return data_flt, zf
+
+
+def lowpass(data, fs, freq, corners=4, zi=None, zerophase=False):
+    """
+    Filter data removing data over certain frequency 'freq' using Butterworth
+    lowpass filter of 'corners' corners.
+
+    :param data: numpy.ndarray. Data to filter.
+    :param fs: Sampling rate in Hz.
+    :param freq: Filter corner frequency.
+    :param corners: Filter corners / order.
+    :param zi : None, 0, or array_like. Initial conditions for the cascaded
+        filter delays. It is a vector of shape (n_sections, nch, 2). Set to 0 to
+        trigger a output of the final filter delay values.
+    :param zerophase: If True, apply filter once forwards and once backwards.
+        This results in twice the number of corners but zero phase shift in
+        the resulting filtered data. Only valid when zi is None.
+    :return: Filtered data and the final filter delay values (if zi is not
+        None).
+    """
+    if len(data.shape) == 1:
+        data = data[np.newaxis, :]
+    fe = 0.5 * fs
+    f = freq / fe
+    # raise for some bad scenarios
+    if f > 1:
+        f = 1.0
+        msg = 'Selected corner frequency is above Nyquist. Setting Nyquist ' + \
+              'as high corner.'
+        warnings.warn(msg)
+    z, p, k = iirfilter(corners, f, btype='lowpass', ftype='butter',
+                        output='zpk')
+    sos = zpk2sos(z, p, k)
+    if zi is None:
+        data_flt = sosfilt(sos, data)
+        if zerophase:
+            data_flt = sosfilt(sos, data_flt[:, ::-1])[:, ::-1]
+        return data_flt
+    elif isinstance(zi, (int, float)):
+        zi = np.ones((sos.shape[0], len(data), 2)) * zi
+    
+    data_flt, zf = sosfilt(sos, data, zi=zi)
+    return data_flt, zf
+
+
+def lowpass_cheby_2(data, fs, freq, maxorder=12, zi=None, ba=False,
+                    freq_passband=False):
+    """
+    Filter data by passing data only below a certain frequency. The main purpose
+    of this cheby2 filter is downsampling. This method will iteratively design a
+    filter, whose pass band frequency is determined dynamically, such that the
+    values above the stop band frequency are lower than -96dB.
+
+    :param data: numpy.ndarray. Data to filter.
+    :param fs: Sampling rate in Hz.
+    :param freq: The frequency above which signals are attenuated with 95 dB.
+    :param maxorder: Maximal order of the designed cheby2 filter.
+    :param zi : None, 0, or array_like. Initial conditions for the cascaded
+        filter delays. It is a vector of shape (n_sections, nch, 2). Set to 0 to
+        trigger a output of the final filter delay values.
+    :param ba: If True return only the filter coefficients (b, a) instead of
+        filtering.
+    :param freq_passband: If True return additionally to the filtered data, the
+        iteratively determined pass band frequency.
+    :return: Filtered data, the final filter delay values (if zi is not None)
+        and the determined pass band frequency (if freq_passband is True).
+    """
+    if data.ndim == 1:
+        data = data[np.newaxis, :]
+
+    nyquist = fs * 0.5
+    # rp - maximum ripple of passband, rs - attenuation of stopband
+    rp, rs, order = 1, 96, 1e99
+    ws = freq / nyquist  # stop band frequency
+    wp = ws  # pass band frequency
+    # raise for some bad scenarios
+    if ws > 1:
+        ws = 1.0
+        warnings.warn('Selected corner frequency is above Nyquist. Setting '
+                      'Nyquist as high corner.')
+    while True:
+        if order <= maxorder:
+            break
+        wp = wp * 0.99
+        order, wn = cheb2ord(wp, ws, rp, rs, analog=0)
+    if ba:
+        return cheby2(order, rs, wn, btype='low', analog=0, output='ba')
+    z, p, k = cheby2(order, rs, wn, btype='low', analog=0, output='zpk')
+    sos = zpk2sos(z, p, k)
+    if zi is None:
+        data_flt = sosfilt(sos, data)
+        if freq_passband:
+            return data_flt, wp * nyquist
+        return data_flt
+    elif isinstance(zi, (int, float)):
+        zi = np.ones((sos.shape[0], len(data), 2)) * zi
+    
+    data_flt, zf = sosfilt(sos, data, zi=zi)
+    if freq_passband:
+        return data_flt, zf, wp * nyquist
+    return data_flt, zf
+
+
+def highpass(data, fs, freq, corners=4, zi=None, zerophase=False):
+    """
+    Filter data removing data below certain frequency 'freq' using Butterworth
+    highpass filter of 'corners' corners.
+
+    :param data: numpy.ndarray. Data to filter.
+    :param fs: Sampling rate in Hz.
+    :param freq: Filter corner frequency.
+    :param corners: Filter corners / order.
+    :param zerophase: If True, apply filter once forwards and once backwards.
+        This results in twice the number of corners but zero phase shift in
+        the resulting filtered data. Only valid when zi is None.
+    :return: Filtered data and the final filter delay values (if zi is not
+        None).
+    """
+    if len(data.shape) == 1:
+        data = data[np.newaxis, :]
+    fe = 0.5 * fs
+    f = freq / fe
+    # raise for some bad scenarios
+    if f > 1:
+        msg = 'Selected corner frequency is above Nyquist.'
+        raise ValueError(msg)
+    z, p, k = iirfilter(corners, f, btype='highpass', ftype='butter',
+                        output='zpk')
+    sos = zpk2sos(z, p, k)
+    if zi is None:
+        data_flt = sosfilt(sos, data)
+        if zerophase:
+            data_flt = sosfilt(sos, data_flt[:, ::-1])[:, ::-1]
+        return data_flt
+    elif isinstance(zi, (int, float)):
+        zi = np.ones((sos.shape[0], len(data), 2)) * zi
+    
+    data_flt, zf = sosfilt(sos, data, zi=zi)
+    return data_flt, zf
+
+def envelope(data):
+    """
+    Computes the envelope of the given data. The envelope is determined by
+    adding the squared amplitudes of the data and it's Hilbert-Transform and
+    then taking the square-root. The envelope at the start/end should not be
+    taken too seriously.
+
+    :param data: numpy.ndarray. Data to make envelope of.
+    :return: Envelope of input data.
+    """
+    return abs(hilbert(data, axis=-1))

daspy/basic_tools/freqattributes.py

@@ -0,0 +1,117 @@
+# Purpose: Analyze frequency attribute and transform in frequency domain
+# Author: Minzhe Hu
+# Date: 2024.6.8
+# Email: hmz2018@mail.ustc.edu.cn
+import numpy as np
+from numpy.fft import rfft, rfft2, fftshift, fftfreq, rfftfreq
+from scipy.signal import stft
+from daspy.basic_tools.preprocessing import demeaning, detrending, cosine_taper
+
+
+def next_pow_2(i):
+    """
+    Find the next power of two.
+
+    :param i: float or int.
+    :return: int. The next power of two for i.
+    """
+    buf = np.ceil(np.log2(i))
+    return np.power(2, buf).astype(int)
+
+
+def spectrum(data, fs, taper=0.05, nfft='default'):
+    """
+    Computes the spectrum of the given data.
+
+    :param data: numpy.ndarray. Data to make spectrum of.
+    :param fs: Sampling rate in Hz.
+    :param taper: Decimal percentage of Tukey taper.
+    :param nfft: Number of points for FFT. None = sampling points, 'default' =
+        next power of 2 of sampling points.
+    :return: Spectrum and frequency sequence.
+    """
+    if len(data.shape) == 1:
+        data = data.reshape(1, len(data))
+    elif len(data.shape) != 2:
+        raise ValueError("Data should be 1-D or 2-D array")
+    data = cosine_taper(data, (0, taper))
+
+    if nfft == 'default':
+        nfft = next_pow_2(len(data[0]))
+    elif nfft is None:
+        nfft = len(data[0])
+
+    spec = rfft(data, n=nfft, axis=1)
+    f = rfftfreq(nfft, d=1 / fs)
+
+    return spec, f
+
+
+def spectrogram(data, fs, nperseg=256, noverlap=None, nfft=None, detrend=False,
+                boundary='zeros'):
+    """
+    Computes the spectrogram of the given data.
+
+    :param data: 1-D or 2-D numpy.ndarray. Data to make spectrogram of.
+    :param fs: Sampling rate in Hz.
+    :param nperseg: int. Length of each segment.
+    :param noverlap: int. Number of points to overlap between segments. If None,
+        noverlap = nperseg // 2.
+    :param nfft: int. Length of the FFT used. None = nperseg.
+    :param detrend : str or bool. Specifies whether and how to detrend each
+        segment.  'linear' or 'detrend' or True = detrend, 'constant' or
+        'demean' = demean.
+    :param boundary: str or None. Specifies whether the input signal is extended
+        at both ends, and how to generate the new values, in order to center the
+        first windowed segment on the first input point. This has the benefit of
+        enabling reconstruction of the first input point when the employed
+        window function starts at zero. Valid options are ['even', 'odd',
+        'constant', 'zeros', None].
+    :return: Spectrogram, frequency sequence and time sequence.
+    """
+    if detrend in [True, 'linear', 'detrend']:
+        detrend = detrending
+    elif detrend in ['constant', 'demean']:
+        detrend = demeaning
+    if data.ndim == 1:
+        f, t, Zxx = stft(data, fs=fs, nperseg=nperseg, noverlap=noverlap,
+                         nfft=nfft, detrend=detrend, boundary=boundary)
+    elif len(data) == 1:
+        f, t, Zxx = stft(data[0], fs=fs, nperseg=nperseg, noverlap=noverlap,
+                         nfft=nfft, detrend=detrend, boundary=boundary)
+    else:
+        Zxx = []
+        for d in data:
+            f, t, Zxxi = stft(d, fs=fs, nperseg=nperseg, noverlap=noverlap,
+                              nfft=nfft, detrend=detrend, boundary=boundary)
+            Zxx.append(abs(Zxxi))
+        Zxx = np.mean(np.array(Zxx), axis=0)
+
+    return Zxx, f, t
+
+
+def fk_transform(data, dx, fs, taper=(0, 0.05), nfft='default'):
+    """
+    Transform the data to the fk domain using 2-D Fourier transform method.
+
+    :param data: numpy.ndarray. Data to do fk transform.
+    :param dx: Channel interval in m.
+    :param fs: Sampling rate in Hz.
+    :param taper: float or sequence of floats. Each float means decimal
+        percentage of Tukey taper for corresponding dimension (ranging from 0 to
+        1). Default is 0.1 which tapers 5% from the beginning and 5% from the
+        end.
+    :param nfft: Number of points for FFT. None means sampling points; 'default'
+        means next power of 2 of sampling points, which makes result smoother.
+    """
+    nch, nt = data.shape
+    data = cosine_taper(data, taper)
+    if nfft == 'default':
+        nfft = (next_pow_2(nch), next_pow_2(nt))
+    elif not nfft:
+        nfft = (nch, nt)
+
+    fk = fftshift(rfft2(data, s=nfft), axes=0)
+    f = rfftfreq(nfft[1], d=1. / fs)
+    k = fftshift(fftfreq(nfft[0], d=dx))
+    return fk, f, k