<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 To make a better wedge
This notebook is an update to the notebook entitled "To make a wedge" featured in the blog post, To make a wedge, on December 12, 2013.
Start by importing Numpy and Matplotlib's pyplot module in the usual way
Step1:  Import the ricker wavelet function from bruges
Step2:  Make a wedge
Step3:  Let's make a more generic wedge that will handle any 3 layer case we want to make.
Step4:  Plotting the synthetic
Step5:  We can make use of the awesome apply_along_axis in Numpy to avoid looping over all the traces. https
<ASSISTANT_TASK:>
Python Code:
import numpy as np
% matplotlib inline
import matplotlib.pyplot as plt
Explanation: To make a better wedge
This notebook is an update to the notebook entitled "To make a wedge" featured in the blog post, To make a wedge, on December 12, 2013.
Start by importing Numpy and Matplotlib's pyplot module in the usual way:
End of explanation
from bruges.filters import ricker
Explanation: Import the ricker wavelet function from bruges:
End of explanation
from IPython.display import Image
Explanation: Make a wedge
End of explanation
Image('images/generic_wedge.png', width=600)
defaults = {'ta1':150, 'tb1':30, 'dta':50, 'dtb':50, 
            'xa1':100, 'xa2':100, 'dx':1,
            'mint':0, 'maxt': 600, 'dt':1,
            'minx':0, 'maxx': 500}
def make_upper_boundary(**kw):
    x = kw['maxx']-kw['minx']
    t0 = kw['ta1']
    x2 = np.arange(1, x-(kw['xa2']+kw['xa1']), kw['dx'])
    m2 = kw['dta']/x2[-1]
    seg1 = np.ones(int(kw['xa1']/kw['dx']))
    seg3 = np.ones(int(kw['xa2']/kw['dx']))
    seg2 =  x2 * m2
    interface = t0 + np.concatenate((seg1, seg2, kw['dta']+seg3))
    return interface
def make_lower_boundary(**kw):
    x = kw['maxx']-kw['minx']
    t1 = kw['ta1'] + kw['tb1']
    x2 = np.arange(1, x-(kw['xa2']+kw['xa1']), kw['dx'])
    m2 = (kw['dta']+kw['dtb'])/x2[-1]
    seg1 = np.ones(int(kw['xa1']/kw['dx']))
    seg3 = np.ones(int(kw['xa2']/kw['dx']))
    seg2 =  x2 * m2
    interface = t1 + np.concatenate((seg1, seg2, seg2[-1]+seg3))
    return interface
def make_wedge(kwargs):
    upper_interface = make_upper_boundary(**kwargs)
    lower_interface = make_lower_boundary(**kwargs)
    return upper_interface, lower_interface
def plot_interfaces(ax, upper, lower, **kw):
    ax.plot(upper,'-r')
    ax.plot(lower,'-b')
    ax.set_ylim(0,600)
    ax.set_xlim(kw['minx'],kw['maxx'])
    ax.invert_yaxis()
upper, lower = make_wedge(defaults)
f = plt.figure()
ax = f.add_subplot(111)
plot_interfaces(ax, upper, lower, **defaults)
def make_meshgrid(**kw):
    upper, lower = make_wedge(defaults)
    t = np.arange(kw['mint'], kw['maxt']-1, kw['dt'])
    x = np.arange(kw['minx'], kw['maxx']-1, kw['dx'])
    xv, yv = np.meshgrid(x, t, sparse=False, indexing='ij')
    return xv, yv
xv, yv = make_meshgrid(**defaults)
conditions = {'upper': yv.T < upper, 
              'middle': (yv.T >= upper) & (yv.T <= lower),
              'lower': yv.T > lower
             }
labels = {'upper': 1, 'middle':2, 'lower': 3}
d = yv.T.copy()
for name, cond in conditions.items():
    d[cond] = labels[name]
plt.imshow(d, cmap='copper')
vp = np.array([3300., 3200., 3300.]) 
rho = np.array([2600., 2550., 2650.]) 
AI = vp*rho
AI
model = d.copy()
model[model == 1] = AI[0]
model[model == 2] = AI[1]
model[model == 3] = AI[2]
def wvlt(f): 
    return ricker(0.512, 0.001, f)
def conv(a):
    return np.convolve(wvlt(f), a, mode='same')
plt.imshow(model, cmap='Spectral')
plt.colorbar()
plt.title('Impedances')
Explanation: Let's make a more generic wedge that will handle any 3 layer case we want to make.
End of explanation
# These are just some plotting parameters
rc_params = {'cmap':'RdBu', 
             'vmax':0.05, 
             'vmin':-0.05, 
             'aspect':0.75}
txt_params = {'fontsize':12, 'color':'black',
              'horizontalalignment':'center',
              'verticalalignment':'center'}
tx = [0.85*defaults['maxx'],0.85*defaults['maxx'],0.85*defaults['maxx']]
ty = [(defaults['ta1'] + defaults['dta'])/2,
      defaults['ta1'] + defaults['dta'] + (defaults['dtb']/1.33),
      defaults['maxt']-(defaults['maxt'] - defaults['ta1'] - defaults['dta'] - defaults['dtb'])/2]
rock_names = ['shale1', 'sand', 'shale2']
defaults['ta1'], defaults['dta'], defaults['dtb']/1.25
rc = (model[1:] - model[:-1]) / (model[1:] + model[:-1])
Explanation: Plotting the synthetic
End of explanation
freqs = np.array([7,14,21])
f, axs = plt.subplots(1,len(freqs), figsize=(len(freqs)*5,6))
for i, f in enumerate(freqs):
    axs[i].imshow(np.apply_along_axis(conv, 0, rc), **rc_params)
    [axs[i].text(tx[j], ty[j], rock_names[j], **txt_params) for j in range(3)]
    plot_interfaces(axs[i], upper, lower, **defaults)
    axs[i].set_ylim(defaults['maxt'],defaults['mint'])
    axs[i].set_title( f'{f} Hz wavelet' )
    axs[i].grid(alpha=0.5)
Explanation: We can make use of the awesome apply_along_axis in Numpy to avoid looping over all the traces. https://docs.scipy.org/doc/numpy/reference/generated/numpy.apply_along_axis.html
End of explanation
<END_TASK>