Implementing FK filter

[Shihao-Yuan]

I've implemented an F-K filter function for the Patch data type. Could any of you please test it and provide feedback?

import sys 
import dascore as dc
from dascore.units import Hz
import numpy as np
import matplotlib.pyplot as plt

def nextpow2(N):
    """ Finding the next power of 2 such that 2**n >= N"""
    n = 1
    while n < N: n *= 2
    return n

def fkfilter(patch, va, vb, vc, vd, symmetry=1, passband=1):
    """
    Frequency-wavenumber(F-K)  filter
    
    Inputs:
    :param patch (PatchType): Input Patch type of data 
    :param va, vb, vc, vd (float): Apparent velocity, va < vb < vc < vd
    :param symmetry (int): 1 for absolute apparent velocity range, 0 for user-defined positive or negative apparent velocity range
    :param passband (int): 1 to pass [vb, vc], 0 to pass (-inf, va) and (vd, +inf)

    Outputs:
    :returns: patch_fk_filtered (PatchType): Filtered patch data of data
    """
    
    # Sample wavenumbers
    dx = patch.get_coord('distance').step
    nx = patch.data.shape[0]
    Nx = nextpow2(nx)
    k = np.fft.fftshift(np.fft.fftfreq(Nx, dx))
    
    # Sample frequencies
    dt = patch.get_coord('time').step / np.timedelta64(1, 's')
    nt = patch.data.shape[1]
    Nt = nextpow2(nt)
    f = np.fft.fftshift(np.fft.fftfreq(Nt, dt))
    
    # Compute FFT
    data_fk = np.fft.fftshift(np.fft.fft2(patch.data, s=(Nx, Nt)))
    
    # Compute apparent velocity
    k[np.abs(k)<=sys.float_info.epsilon] +=  sys.float_info.epsilon
    if symmetry:
        vel = np.abs(f / k[:, np.newaxis])
    else: 
        vel = f / k[:, np.newaxis]
    
    # Apply filter with tapering function 
    fac = np.where((vel >= va) & (vel <= vb), 1.0 - np.sin(0.5 * np.pi * (vel - va) / (vb - va)), 1.0)
    fac = np.where((vel >= vb) & (vel <= vc), 0.0, fac)
    fac = np.where((vel >= vc) & (vel <= vd), np.sin(0.5 * np.pi * (vel - vc) / (vd - vc)), fac)
    
    if passband:
        fac = 1.0 - fac
    
    # Apply filter
    data_fk *= fac
    
    # Compute inverse FFT
    d_filtered = np.fft.ifft2(np.fft.ifftshift(data_fk), s=(Nx, Nt)).real[0:nx,0:nt]
    
    return patch.update(data=d_filtered)

# Example
patch_raw = (dc.get_example_patch('example_event_1')
             .taper(time=0.05)
             .pass_filter(time=(1*Hz, 500*Hz))            
            )
nx = patch_raw.data.shape[0]
nt = patch_raw.data.shape[1]
Nx = nextpow2(nx)
Nt = nextpow2(nt)
dx = patch_raw.get_coord('distance').step
dt = patch_raw.get_coord('time').step / np.timedelta64(1, 's')
k = np.fft.fftshift(np.fft.fftfreq(Nx, dx))
f = np.fft.fftshift(np.fft.fftfreq(Nt, dt))
patch_raw_fk = np.fft.fftshift(np.fft.fft2(patch_raw.data, s=(Nx, Nt)))

# F-K filtered data 
patch_filtered = fkfilter(patch_raw,va=2e3, vb=2.2e3, vc=8e3, vd=2.e4, symmetry=1, passband=1)
patch_filtered_fk = np.fft.fftshift(np.fft.fft2(patch_filtered.data, s=(Nx, Nt)))

# Plot raw waveform and F-K spectrum before and after applyting F-K filter
plt.figure(figsize=(12, 8))

ax1 = plt.subplot(221)
patch_raw.viz.waterfall(ax=ax1,show=False, scale=0.5);
ax1.set_title('Raw')

ax2 = plt.subplot(222)
patch_filtered.viz.waterfall(ax=ax2,show=False, scale=0.5);
ax2.set_title('Filtered')

ax3 = plt.subplot(223)
ax3.imshow(np.abs(patch_raw_fk.T),aspect='auto',extent=[k.min(),k.max(),f.max(),f.min()],
           vmin=0, vmax=5e5,cmap='seismic')
ax3.set_xlim(-0.2,0.2)
ax3.set_ylim(-400,400)
ax3.set_xlabel('Wavenumber (1/m)')
ax3.set_ylabel('Frequency (1/m)')
ax3.set_title('F-K spectrum - Raw')

ax4 = plt.subplot(224,sharex=ax3,sharey=ax3)
ax4.imshow(np.abs(patch_filtered_fk.T),aspect='auto',extent=[k.min(),k.max(),f.max(),f.min()],
           vmin=0, vmax=5e5,cmap='seismic')
ax4.set_xlabel('Wavenumber (1/m)')
ax4.set_ylabel('Frequency (1/m)')
ax4.set_title('F-K spectrum - Filtered')

plt.tight_layout()
plt.show()

[d-chambers]

Hey @Shihao-Yuan, this is a good start, and the examples are compelling.

As far as suggestions, a few things come to mind:

1. Be sure to follow the documentation guidelines in formatting your docstrings. It is better to include the type hints in the function signature, then you don't need to mention the types in the documentation. I am wondering if rather than symmetry=True if it would be more clear for a directional=False.

2. For the flags (symmetry and passband) using 1 or 0 works, but it is better to use True or False since these are really booleans not integers (e.g. a value of 42 makes no sense). I think calling them directional and notch might be more clear, what do you think?

3. Rather than doing the FFT manually, we might as well use the Patch.dft. That way, we can detect if we already have a patch in the FT domain and leave it alone. It also handles the shifting and coordinates for you, so it makes this function simpler.

4. Why not make this a general 2d ft slope filter, which is a superset of an F-K filter? In that case, we might rename the function and we should support dimension names other than just "time" and "distance". Probably a dims parameter would be the way to go.

5. I think it might be cleaner to collapse va, vb, vc, vd into a single parameter. Maybe filt after obspy's pre-filt argument in remove response which I think has identical function.

So implementing these ideas (I haven't tested this, its just to get you started):

import dascore as dc    


@dc.patch_function()
def ft_slope_filter(patch, filt, dims=("time", "distance"), directional=False, notch=False):
    """
    Filter the patch over certain slopes in the 2D Fourier domain.

    Most commonly this used as an F-K (frequency wavenumber) filter to attenuate energy with specified
    apparent velocities.
    
    Parameters
    -----------------
    patch
        The patch to filter.
    filt
         A length 4 sequence of the form (va, vb, vc, vd). If notch is False, {describe filter shape}
    dims
         The dimensions to filter/transform which 
    directional
        If True, the filter should be considered direction. That is to say, the sign of the values in `filt` indicate the 
        direction (towards or away) with increasing coordinate values. This can be used for up-down/left-right
        separation. 
    notch
         If True, the filter represents a notch. 

    Examples
    -----------------
    # need several examples of each feature here, including passing a patch that has already been dft'ed. 

    Notes
    ---------
    - It is a good idea to first pad the transformed dimensions to fast fft lengths. See [Patch.pad](`dascore.Patch.pad`).
    - See the [F-K filter page](`docs/processing/fk.qmd`). 

    """
    # Perform dft if needed.
    dft_patch = patch.dft(dims)
    transformed = dft_patch is not patch  # if dft didn't create a new patch it was already in dft.   

    # Get slope array for specified dimensions
    assert len(dims) == 2, "dims must be a length 2 tuple"
    coord1 = dft_patch.get_array(f"ft_{dims[0]}")
    coord2 = dft_patch.get_array(f"ft_{dims[1]}")
    vel = coord1 / coord2[:, None]
    if not directional:
        vel = np.abs(vel)

    # Apply filter with tapering function 
    fac = np.where((vel >= va) & (vel <= vb), 1.0 - np.sin(0.5 * np.pi * (vel - va) / (vb - va)), 1.0)
    fac = np.where((vel >= vb) & (vel <= vc), 0.0, fac)
    fac = np.where((vel >= vc) & (vel <= vd), np.sin(0.5 * np.pi * (vel - vc) / (vd - vc)), fac)
    if passband:
        fac = 1.0 - fac
    
    # Apply filter
    new_data = fft_patch.data * fac
    out = fft_patch.update(data=new_data)

    # Undo transformation if an un-transformed patch was passed in.
    if transformed:
        out = fft_patch.idft()
    
    return out

Since this is a really powerful filter, I suggest we create a recipe or processing page called f-k filtering. In that page we could show:

• An example of a directional filter to highlight the down-going S-reflection (I think Stanek did the same in his paper)
• An example separating P and S arrivals with a tighter filter, maybe notching out the S waves to see if we see any weak P arrivals?

Does anyone else have any thoughts on this function design?

[Shihao-Yuan]

Hi @d-chambers ,

1. Why do we need this?

transformed = dft_patch is not patch  # if dft didn't create a new patch it was already in dft.  

2. What does this 'transformation' and 'un-transformation' mean?

# Undo transformation if an un-transformed patch was passed in.
if transformed:
  out = fft_patch.idft()

[d-chambers]

Hey @Shihao-Yuan

The idea is that the 2d fft can be expensive, especially on a large patch. You may also be doing multiple processing steps in the FK domain, so it might be useful to leave the patch in the FK domain if so. These lines test if the patch is already in the FK domain and do not transform it if so. If not, it will perform both the dft and the idft to return the patch to its original domain.

[d-chambers]

closed by #413