import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 22})

plt.interactive(True)

# Channel flow
data=np.genfromtxt("channel_flow_data.dat", comments="%")

ni=199  # number of grid nodes in x_1 direction
nj=28  # number of grid nodes in x_2 direction
x1=data[:,0] #don't use this array
x2=data[:,1] #don't use this array
v1=data[:,2] #don't use this array
v2=data[:,3] #don't use this array
p=data[:,4]  #don't use this array

# transform the arrays from 1D fields x(n) to 2D fields x(i,j)
# the first index 'i', correponds to the x-direction
# the second index 'j', correponds to the y-direction

x1_2d=np.reshape(x1,(nj,ni)) #this is x_1 (streamwise coordinate)
x2_2d=np.reshape(x2,(nj,ni)) #this is x_2 (wall-normal coordinate)
v1_2d=np.reshape(v1,(nj,ni)) #this is v_1 (streamwise velocity component)
v2_2d=np.reshape(v2,(nj,ni)) #this is v_1 (wall-normal velocity component)
p_2d=np.reshape(p,(ni,nj))   #this is p   (pressure)

x1_2d=np.transpose(x1_2d)
x2_2d=np.transpose(x2_2d)
v1_2d=np.transpose(v1_2d)
v2_2d=np.transpose(v2_2d)
p_2d=np.transpose(p_2d)


#************
# velocity profile plot
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i=169 # plot the velocity profile for i=169
plt.plot(v1_2d[i,:],x2_2d[i,:],'b-')
i=4 # plot the velocity profile for i=4
plt.plot(v1_2d[i,:],x2_2d[i,:],'r--')  #red dashed line
plt.title('Velocity profile')
plt.axis([0,1.5,0,0.01]) # set x & y axis
plt.xlabel('$V_1$') 
plt.ylabel('$x_2$') 
plt.text(0.04,0.004,'$x_1=0.008$ and $0.34$') # show this text at (0.04,0.004)
plt.savefig('velprof.eps')

################################ contour plot of v1
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
plt.contourf(x1_2d,x2_2d,v1_2d, 50)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
plt.clim(0.1,1.)
plt.title("contour v_1 plot")
plt.axis([0,0.1,0,0.011]) # zoom-in on the first 0.1m from the inlet
plt.savefig('v1_contour.eps')

################################ compute the velocity gradient dv_1/dx_2
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
dv1_dx2=np.zeros((ni,nj))
for i in range(0,ni-1):
   for j in range(0,nj-1):
      dx2=x2_2d[i,j+1]-x2_2d[i,j-1]
      dv1_dx2[i,j]=(v1_2d[i,j+1]-v1_2d[i,j-1])/dx2

# fix the derivative at the walls
for i in range(0,ni-1):
# lower wall
   dx2=x2_2d[i,1]-x2_2d[i,0]
   dv1_dx2[i,0]=(v1_2d[i,1]-v1_2d[i,0])/dx2
# upper wall
   dx2=x2_2d[i,nj-1]-x2_2d[i,nj-2]
   dv1_dx2[i,nj-1]=(v1_2d[i,nj-1]-v1_2d[i,nj-2])/dx2

# you can also use the built-in command
x1_1d=x1_2d[:,1] # make 1d array
x2_1d=x2_2d[1,:] # make 1d array
dv1_dx1_built_in, dv1_dx2_bulit_in=np.gradient(v1_2d,x1_1d,x2_1d)


i=169 # plot the velocity gradient for i=169
plt.plot(dv1_dx2[i,:],x2_2d[i,:])
plt.axis([-550,550,0,0.01]) # set x & y axis
plt.title('Velocity gradient')
plt.xlabel('$\partial v_1/\partial x_2$')
plt.ylabel('$x_2$') 
plt.text(-380,0.004,'$x_1=0.52$')
plt.savefig('v1_grad.eps')