#!/usr/bin/env python
# coding: utf-8

# ## Compare receiver function deconvolution methods

# This notebook uses the [rf](https://github.com/trichter/rf) package to compare the results of different deconvolution methods.
# 
# If you want to run this notebook locally, check the repository [readme](https://github.com/hfmark/notebooks) for some instructions on installing dependencies. For this notebook you will need to have obspy and rf installed.

# To start, we load example data into an RFStream instance by calling `read_rf()`. Calling this function without other arguments loads the sample dataset provided with the rf package. The stream consists of 3-component data from 3 earthquakes, 9 traces in total.

# In[1]:


import matplotlib.pyplot as plt
from rf import read_rf, rfstats

stream = read_rf()
# plot the first 3 traces/data from the first event
stream[:3].plot()


# Next, the rfstats function helps us fill in header info for the receiver function calculation, including onset time.

# In[2]:


rfstats(stream)


# The data need a little preprocessing:

# In[3]:


stream.filter('bandpass', freqmin=0.33, freqmax=1)
stream.trim2(-25,75,'onset')
stream[:3].plot(type='relative')


# Next, we make copies of the stream and do the receiver function calculation (rotation, deconvolution, moveout correction) using different deconvolution methods: time- and frequency-domain water level, and time-domain iterative.

# In[16]:


rf_time = stream.copy().rf(deconvolve='time',rotate='NE->RT',solve_toeplitz='scipy').moveout()
rf_freq = stream.copy().rf(deconvolve='waterlevel',rotate='NE->RT').moveout()
rf_iter = stream.copy().rf(deconvolve='iterative',rotate='NE->RT').moveout()
rf_mult = stream.copy().rf(deconvolve='multitaper',rotate='NE->RT').moveout()

rf_time.trim2(-5,20,'onset')
rf_freq.trim2(-5,20,'onset')
rf_iter.trim2(-5,20,'onset')
rf_mult.trim2(-5,20,'onset')


# We can plot and examine the radial components (since these are PRFs) for each method

# In[5]:


rf_time.select(component='R').plot_rf()


# In[6]:


rf_freq.select(component='R').plot_rf()


# In[11]:


rf_iter.select(component='R').plot_rf()


# In[17]:


rf_mult.select(component='R').plot_rf()


# We can also make a comparison plot overlaying the RFs for one event from each deconvolution method

# In[18]:


evt = 0  # choose 0, 1, or 2
rflist = [rf_time,rf_freq,rf_iter,rf_mult]; labels = ['time','waterlevel','iterative','multitaper']
plt.figure()
for i,st in enumerate(rflist):
    sel = st[int(3*evt):int(3*(evt+1))]
    sel = sel.select(component='R')
    plt.plot(sel[0].times(),sel[0].data,label=labels[i])
plt.legend(fontsize=9)


# A similar plot can be made with the stacked RFs for each deconvolution method

# In[19]:


rflist = [rf_time,rf_freq,rf_iter,rf_mult]; labels = ['time','waterlevel','iterative','multitaper']
plt.figure()
for i,st in enumerate(rflist):
    sel = st.copy().select(component='R').stack()
    plt.plot(sel[0].times(),sel[0].data,label=labels[i])
plt.legend(fontsize=9)


# In[ ]: