%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: FD_1D_DX4_DT4_LW_fast 1-D acoustic Finite-Difference modelling
GNU General Public License v3.0
Author: Florian Wittkamp
Finite-Difference acoustic seismic wave simulation
Discretization of the first-order acoustic wave equation
Temporal second-order accuracy $O(\Delta T^4)$
Spatial fourth-order accuracy  $O(\Delta X^4)$
Theory to Lax-Wendroff method is given in:
Dablain, M. A. (1986).
The application of high-order differencing to the scalar wave equation.
Geophysics, 51(1), 54-66.
Initialisation
End of explanation
"""


# Discretization
c1=20   # Number of grid points per dominant wavelength
c2=0.5  # CFL-Number
nx=2000 # Number of grid points
T=10     # Total propagation time

# Source Signal
f0= 10      # Center frequency Ricker-wavelet
q0= 1       # Maximum amplitude Ricker-Wavelet
xscr = 100  # Source position (in grid points)

# Receiver
xrec1=400  # Position Reciever 1 (in grid points)
xrec2=800  # Position Reciever 2 (in grid points)
xrec3=1800 # Position Reciever 3 (in grid points)

# Velocity and density
modell_v = np.hstack((1000*np.ones((int(nx/2))),1500*np.ones((int(nx/2)))))
rho=np.hstack((1*np.ones((int(nx/2))),1.5*np.ones((int(nx/2)))))

"""
Explanation: Input Parameter
End of explanation
"""


# Init wavefields
vx=np.zeros(nx)
p=np.zeros(nx)

# Calculate first Lame-Paramter
l=rho * modell_v * modell_v

cmin=min(modell_v.flatten())  # Lowest P-wave velocity
cmax=max(modell_v.flatten())  # Highest P-wave velocity
fmax=2*f0                     # Maximum frequency
dx=cmin/(fmax*c1)             # Spatial discretization (in m)
dt=dx/(cmax)*c2               # Temporal discretization (in s)
lampda_min=cmin/fmax          # Smallest wavelength

# Output model parameter:
print("Model size: x:",dx*nx,"in m")
print("Temporal discretization: ",dt," s")
print("Spatial discretization: ",dx," m")
print("Number of gridpoints per minimum wavelength: ",lampda_min/dx)

"""
Explanation: Preparation
End of explanation
"""


x=np.arange(0,dx*nx,dx) # Space vector
t=np.arange(0,T,dt)     # Time vector
nt=np.size(t)           # Number of time steps

# Plotting model
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.subplots_adjust(wspace=0.4,right=1.6)
ax1.plot(x,modell_v)
ax1.set_ylabel('VP in m/s')
ax1.set_xlabel('Depth in m')
ax1.set_title('P-wave velocity')

ax2.plot(x,rho)
ax2.set_ylabel('Density in g/cm^3')
ax2.set_xlabel('Depth in m')
ax2.set_title('Density');

"""
Explanation: Create space and time vector
End of explanation
"""


tau=np.pi*f0*(t-1.5/f0)
q=q0*(1.0-2.0*tau**2.0)*np.exp(-tau**2)

# Plotting source signal
plt.figure(3)
plt.plot(t,q)
plt.title('Source signal Ricker-Wavelet')
plt.ylabel('Amplitude')
plt.xlabel('Time in s')
plt.draw()

"""
Explanation: Source signal - Ricker-wavelet
End of explanation
"""


# Init Seismograms
Seismogramm=np.zeros((3,nt)); # Three seismograms

# Calculation of some coefficients
i_dx=1.0/(dx)
i_dx3=1.0/(dx**3)
c9=dt**3/24.0
kx=np.arange(5,nx-4)

print("Starting time stepping...")
## Time stepping
for n in range(2,nt):

        # Inject source wavelet
        p[xscr]=p[xscr]+q[n]

        # Calculating spatial derivative
        p_x=i_dx*9.0/8.0*(p[kx+1]-p[kx])-i_dx*1.0/24.0*(p[kx+2]-p[kx-1])
        p_xxx=i_dx3*(-3.0)*(p[kx+1]-p[kx])+i_dx3*(1.0)*(p[kx+2]-p[kx-1])

        # Update velocity
        vx[kx]=vx[kx]-dt/rho[kx]*p_x-l[kx]*c9*1/(rho[kx]**2)*(p_xxx)

        # Calculating spatial derivative
        vx_x= i_dx*9.0/8.0*(vx[kx]-vx[kx-1])-i_dx*1.0/24.0*(vx[kx+1]-vx[kx-2])
        vx_xxx=i_dx3*(-3.0)*(vx[kx]-vx[kx-1])+i_dx3*(1.0)*(vx[kx+1]-vx[kx-2])

        # Update pressure
        p[kx]=p[kx]-l[kx]*dt*(vx_x)-l[kx]**2.0*c9*1.0/(rho[kx])*(vx_xxx)

        # Save seismograms
        Seismogramm[0,n]=p[xrec1]
        Seismogramm[1,n]=p[xrec2]
        Seismogramm[2,n]=p[xrec3]
        
print("Finished time stepping!")

"""
Explanation: Time stepping
End of explanation
"""


## Save seismograms
np.save("Seismograms/FD_1D_DX4_DT4_LW_fast",Seismogramm)

## Plot seismograms
fig, (ax1, ax2, ax3) = plt.subplots(3, 1)
fig.subplots_adjust(hspace=0.4,right=1.6, top = 2 )

ax1.plot(t,Seismogramm[0,:])
ax1.set_title('Seismogram 1')
ax1.set_ylabel('Amplitude')
ax1.set_xlabel('Time in s')
ax1.set_xlim(0, T)

ax2.plot(t,Seismogramm[1,:])
ax2.set_title('Seismogram 2')
ax2.set_ylabel('Amplitude')
ax2.set_xlabel('Time in s')
ax2.set_xlim(0, T)

ax3.plot(t,Seismogramm[2,:])
ax3.set_title('Seismogram 3')
ax3.set_ylabel('Amplitude')
ax3.set_xlabel('Time in s')
ax3.set_xlim(0, T);


"""
Explanation: Save seismograms
End of explanation
"""


