def EEF_main(set_,pi):
NU = len(set_)
# Initialise output arrays
Ef_av_array = np.zeros((NU,N,N))
Ed_av_array = np.zeros((NU,N,N))
# Now start the loop over each forcing index.
for ui in range(0,NU):
#print(ui)
# Redefine U0 and H0
#sigma = set_[ui]
Umag = set_[ui]
U0, H0 = BG_state.BG_uniform(Umag,Hflat,f0,beta,g,y,N)
U0_nd = U0 / U
H0_nd = H0 / chi
# Solution
a1,a2,a3,a4,b4,c1,c2,c3,c4 = solver.SOLVER_COEFFICIENTS(Ro,Re,K_nd,f_nd,U0_nd,H0_nd,omega_nd,gamma_nd,dy_nd,N)
solution = solver.FREE_SLIP_SOLVER(a1,a2,a3,a4,f_nd,b4,c1,c2,c3,c4,Ro*Ftilde1_nd,Ro*Ftilde2_nd,Ftilde3_nd,N,N2)
utilde_nd, vtilde_nd, etatilde_nd = solver.extractSols(solution,N,N2,BC)
u, v, h = solver.SPEC_TO_PHYS(utilde_nd,vtilde_nd,etatilde_nd,T_nd,dx_nd,omega_nd,N)
# Take real part. Don't normalise by forcing ampltiude, as we want to compare energy input by forcing.
u = np.real(u)
v = np.real(v)
h = np.real(h)
# Calculate full flow quantities.
u_full = diagnostics.fullFlow(u,U0_nd)
h_full = diagnostics.fullFlow(h,H0_nd)
#==
# Energy budget terms.
Ed, Ed_av_array[ui,] = energy.budgetDissipation3(U0_nd,H0_nd,u,v,h,Ro,Re,gamma_nd,dx_nd,dy_nd,T_nd,Nt,N)
Ef, Ef_av_array[ui,] = energy.budgetForcing(u_full,v,h_full,F1_nd,F2_nd,F3_nd,Ro,N,T_nd,omega_nd,Nt)
# End loop.
np.save('Ed_av_array'+str(pi),Ed_av_array)
np.save('Ef_av_array'+str(pi),Ef_av_array)
dy = y[1] - y[0] # Distance between gridpoints (m)
dx = x[1] - x[0]
y_grid, x_grid = np.mgrid[slice(-Ly / 2, Ly / 2 + dy, dy),
slice(-Lx / 2, Lx / 2 + 2. * dx, dx)]
K = np.fft.fftfreq(N, Lx / N) # Array of x-gridpoints in wavenumber space
f = f0 + beta * y # Coriolis frequency (s-1)
# Background flow
#=======================================================
if BG == 'UNIFORM':
U0, H0 = BG_state.BG_uniform(Umag, Hflat, f0, beta, g, y, N)
elif BG == 'SHEAR':
U0, H0 = BG_state.BG_shear(Umag, shear, Hflat, f0, beta, g, y, Ly, N)
elif BG == 'GAUSSIAN':
U0, H0 = BG_state.BG_Gaussian(Umag, sigma, JET_POS, Hflat, f0, beta, g, y,
Ly, N)
elif BG == 'LAPGAUSS':
U0, H0 = BG_state.BG_LapGauss(Umag, sigma, JET_POS, Hflat, f0, beta, g, y,
Ly, N)
elif BG == 'QUAD':
U0, H0 = BG_state.BG_quadratic(Umag, Hflat, f0, beta, g, y, Ly, N)
elif BG == 'ZERO':
U0, H0 = BG_state.BG_zero(Hflat, N)
else:
raise ValueError('Invalid BG flow option selected')
f = f0 + beta * y
# Coriolis frequency (s-1)
# Other physical parameters
#=======================================================
rho1_nd = rho1 / rho2
rho2_nd = (rho2 - rho1) / rho2
# Background flow
#=======================================================
if BG2 == 'NONE':
if BG1 == 'UNIFORM':
U1, U2, H1, H2 = BG_state.BG_uniform_none(Umag1, H1_flat, H2_flat,
rho1_nd, rho2_nd, f0, beta,
g, y, N)
elif BG1 == 'GAUSSIAN':
U1, U2, H1, H2 = BG_state.BG_Gaussian_none(Umag1, sigma, JET_POS,
H1_flat, H2_flat, rho1_nd,
rho2_nd, f0, beta, g, y, Ly,
N)
elif BG1 == 'NONE':
U1, U2, H1, H2 = BG_state.BG_none_none(H1_flat, H2_flat, rho1_nd,
rho2_nd, N)
else:
raise ValueError('Invalid BG flow option selected')
# Code not updated to include nonzero lower layer BG flows yet.
# Forcing
def EEF_main(set_, pi):
NU = len(set_)
# Initialise output arrays
EEF_array = np.zeros((NU, 2))
E_array = np.zeros((NU))
M_array = np.zeros((NU, N, N))
N_array = np.zeros((NU, N, N))
Mu_array = np.zeros((NU, N))
Munorm_array = np.zeros((NU, N))
# Now start the loop over each forcing index.
for ui in range(0, NU):
# Redefine U0 and H0
#sigma = set_[ui]
Umag = set_[ui]
U0, H0 = BG_state.BG_uniform(Umag, Hflat, f0, beta, g, y, N)
U0_nd = U0 / U
H0_nd = H0 / chi
# Solution
a1, a2, a3, a4, b4, c1, c2, c3, c4 = solver.SOLVER_COEFFICIENTS(
Ro, Re, K_nd, f_nd, U0_nd, H0_nd, omega_nd, gamma_nd, dy_nd, N)
solution = solver.FREE_SLIP_SOLVER(a1, a2, a3, a4, f_nd, b4, c1, c2,
c3, c4, Ro * Ftilde1_nd,
Ro * Ftilde2_nd, Ftilde3_nd, N, N2)
utilde_nd, vtilde_nd, etatilde_nd = solver.extractSols(
solution, N, N2, BC)
u, v, h = solver.SPEC_TO_PHYS(utilde_nd, vtilde_nd, etatilde_nd, T_nd,
dx_nd, omega_nd, N)
# Take real part.
u = np.real(u)
v = np.real(v)
h = np.real(h)
# Normalise all solutions by the (non-dimensional) forcing amplitude.
u = u / AmpF_nd
v = v / AmpF_nd
h = h / AmpF_nd
# In order to calculate the vorticities of the system, we require full (i.e. BG + forced response) u and eta.
#h_full = np.zeros((N,N,Nt))
#u_full = np.zeros((N,N,Nt))
#for j in range(0,N):
# h_full[j,:,:] = h[j,:,:] + H0_nd[j]
# u_full[j,:,:] = u[j,:,:] + U0_nd[j]
#==
# Correlations (can do before or after energy)
#M_array[ui,:,:] = corr.M(u,v,T_nd)
#N_array[ui,:,:] = corr.N(u,v,T_nd)
#==
# Energy
#KE_BG, KE_BG_tot, PE_BG, PE_BG_tot = energy.energy_BG(U0_nd,H0_nd,Ro,y_nd,dy_nd,N)
#KE, KE_tot = energy.KE(u_full,v,h_full,x_nd,y_nd,dx_nd,dy_nd,N)
#PE, PE_tot = energy.PE(h_full,Ro,x_nd,y_nd,dx_nd,dy_nd,N)
#E_tot = KE_tot + PE_tot - KE_BG_tot - PE_BG_tot
# Use time-mean KE omitting h for now. Find a better way to do this
KE = u**2 + v**2
KE_tot = np.zeros(Nt)
for ti in range(0, Nt):
KE_tot[ti] = diagnostics.domainInt(KE[:, :, ti], x_nd, dx_nd, y_nd,
dy_nd)
E_array[ui] = diagnostics.timeAverage1D(KE_tot, T_nd, Nt)
# Momentum
#==========
uu, uv, vv = momentum.fluxes(u, v)
Mu, Mv, Mu_array[ui, :], Mv_xav = momentum.footprint(
uu, uv, vv, x_nd, T_nd, dx_nd, dy_nd, N, Nt)
#import matplotlib.pyplot as plt
#plt.plot(KE_tot);plt.show()
# Normalise by energy
#u = u / np.sqrt(E_array[ui]); v = v / np.sqrt(E_array[ui]); h = h / np.sqrt(E_array[ui])
#==
# Calculate PV fields and PV fluxes.
#PV_prime, PV_full, PV_BG = PV.potentialVorticity(u,v,h,u_full,h_full,H0_nd,U0_nd,N,Nt,dx_nd,dy_nd,f_nd,Ro)
#uq, Uq, uQ, UQ, vq, vQ = PV.fluxes(u,v,U0_nd,PV_prime,PV_BG,N,Nt)
#P, P_xav = PV.footprint(uq,Uq,uQ,UQ,vq,vQ,x_nd,T_nd,dx_nd,dy_nd,N,Nt)
#EEF_array[ui,:], l_PV = PV.EEF(P_xav,y_nd,y_nd[y0_index],y0_index,dy_nd,N)
#==
# Repeat process with normalised velocities.
u = u / np.sqrt(E_array[ui])
v = v / np.sqrt(E_array[ui])
h = h / np.sqrt(E_array[ui])
uu, uv, vv = momentum.fluxes(u, v)
Mu, Mv, Munorm_array[ui, :], Mv_xav = momentum.footprint(
uu, uv, vv, x_nd, T_nd, dx_nd, dy_nd, N, Nt)
# In order to calculate the vorticities of the system, we require full (i.e. BG + forced response) u and eta.
#h_full = np.zeros((N,N,Nt))
#u_full = np.zeros((N,N,Nt))
#for j in range(0,N):
# h_full[j,:,:] = h[j,:,:] + H0_nd[j]
# u_full[j,:,:] = u[j,:,:] + U0_nd[j]
#==
# Calculate PV fields and PV fluxes.
#PV_prime, PV_full, PV_BG = PV.potentialVorticity(u,v,h,u_full,h_full,H0_nd,U0_nd,N,Nt,dx_nd,dy_nd,f_nd,Ro)
#uq, Uq, uQ, UQ, vq, vQ = PV.fluxes(u,v,U0_nd,PV_prime,PV_BG,N,Nt)
#P, P_xav = PV.footprint(uq,Uq,uQ,UQ,vq,vQ,x_nd,T_nd,dx_nd,dy_nd,N,Nt)
#EEFnorm_array[ui,:], l_PV = PV.EEF(P_xav,y_nd,y_nd[y0_index],y0_index,dy_nd,N)
#np.save('EEF_array'+str(pi),EEF_array)
#np.save('EEFnorm_array'+str(pi),EEFnorm_array)
#np.save('M_array'+str(pi),M_array)
#np.save('N_array'+str(pi),N_array)
np.save('E_array' + str(pi), E_array)
np.save('Mu_array' + str(pi), Mu_array)
np.save('Munorm_array' + str(pi), Munorm_array)
umag_mesh, y0_mesh = np.mgrid[slice(umag_set.min(),(umag_set.max()+du),du),slice(y0_set.min(),y0_set.max()+dy_nd,dy_nd)]; #slice(sigma_set.min()/1000.0,(sigma_set.max()+ds)/1000.,ds/1000.) #======================================================== Umag_ref = 0.8 sigma_ref = 0.02 * 3840000. Qyu = np.zeros((Ns,nn)) Qys = np.zeros((Ns,nn)) for ui in range(0,Ns): # Umag sigma = sigma_set[ui] U0, H0 = BG_state.BG_Gaussian(Umag_ref,sigma,JET_POS,Hflat,f0,beta,g,y,L,N) U0 = U0 / U; H0 = H0 / chi U0 = U0[y0_index_set]; H0 = H0[y0_index_set] U0y = diff(U0,2,0,dy_nd) Q = (f_nd[y0_index_set] / Ro - U0y) / H0 Qys[ui,:] = diff(Q,2,0,dy_nd) # Sigma Umag = umag_set[ui] U0, H0 = BG_state.BG_Gaussian(Umag,sigma_ref,JET_POS,Hflat,f0,beta,g,y,L,N) U0 = U0 / U; H0 = H0 / chi U0 = U0[y0_index_set]; H0 = H0[y0_index_set] U0y = diff(U0,2,0,dy_nd) Q = (f_nd[y0_index_set] / Ro - U0y) / H0
def EEF_main(set_, pi):
NU = len(set_)
# Initialise output arrays
EEF_array = np.zeros((NU, 2))
# Now start the loop over each forcing index.
for ui in range(0, NU):
# Redefine U0 and H0.
#sigma = set_[ui]
Umag1 = set_[ui]
U1, U2, H1, H2 = BG_state.BG_uniform_none(Umag1, H1_flat, H2_flat,
rho1_nd, rho2_nd, f0, beta,
g, y, N)
U1 = U1 / U
U2 = U2 / U
H1 = H1 / chi
H2 = H2 / chi
# Solver coeffs depend on BG state.
a1, a2, a3, a4, b1, b4, c1, c2, c3, c4, c5, d1, d3, d4, d5, e4, e5, f1, f2, f3, f4 = solver.SOLVER_COEFFICIENTS(
Ro, Re, K, f, U1, U2, H1, H2, rho1_nd, rho2_nd, omega, gamma, dy,
N)
# Solve.
solution = solver.FREE_SLIP_SOLVER(a1, a2, a3, a4, b1, b4, c1, c2, c3,
c4, c5, d1, d3, d4, d5, e4, e5, f1,
f2, f3, f4, Ro * Ftilde1,
Ro * Ftilde2, Ftilde3, Ro * Ftilde4,
Ro * Ftilde5, Ftilde6, N, N2)
# Extract flow components.
utilde, vtilde, htilde = solver.extractSols(solution, N, N2, BC)
u, v, h = solver.SPEC_TO_PHYS(utilde, vtilde, htilde, T, Nt, dx, omega,
N)
# Take real part.
u = np.real(u)
v = np.real(v)
h = np.real(h)
# For use in PV and footprint calculations: the 'full' zonal velocities and interface thicknesses.
u_full = np.zeros((N, N, Nt, 2))
h_full = np.zeros((N, N, Nt, 2))
for j in range(0, N):
u_full[j, :, :, 0] = u[j, :, :, 0] + U1[j]
u_full[j, :, :, 1] = u[j, :, :, 1] + U2[j]
h_full[j, :, :, 0] = h[j, :, :, 0] + H1[j]
h_full[j, :, :, 1] = h[j, :, :, 1] + H2[j]
# Call function calculate PV in upper layer.
q = np.zeros((N, N, Nt, 2))
q_full = np.zeros((N, N, Nt, 2))
Q = np.zeros((N, 2))
q[:, :, :,
0], q_full[:, :, :,
0], Q[:, 0] = PV.vort(u[:, :, :, 0], v[:, :, :, 0],
h[:, :, :, 0], u_full[:, :, :, 0],
h_full[:, :, :,
0], H1, U1, N, Nt, dx, dy, f)
# Calculate footprints using previously calculated PV.
P, P_xav = PV.footprint(u_full[:, :, :, 0], v[:, :, :, 0],
q_full[:, :, :, 0], x, y, dx, dy, T, Nt)
# EEF
EEF_array[ui, :], l = PV.EEF(P_xav, y, y0, y0_index, dy, N)
# End loop.
# Save output from processor pi.
np.save('EEF_array' + str(pi), EEF_array)