def EEF_U0(set_, pi):
# For every U0 value in U0_set,
un = len(set_)
EEF_array = np.zeros((un, 6, 2))
for ui in range(
0,
un): # ui indexes the local EEF_array (i.e. computational domain)
# Redefine U0 and H0 in each run.
for j in range(0, N):
U0[j] = set_[ui]
H0[j] = -(U0[j] / g) * (f0 * y[j] + beta * y[j]**2 / 2) + Hflat
U0_nd = U0 / U
H0_nd = H0 / chi
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]
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_array[ui, :] = PV.EEF(P_xav, y_nd, y0_nd, y0_index,
dy_nd, N)
filename = 'EEF_array_' + str(pi)
np.save(filename, EEF_array)
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)
def EIG_DECOMP_main(U0_set, pi):
# Need to loop over a set of background states.
NU = len(U0_set)
theta = np.zeros(
(NU, N, dim),
dtype=complex) # Initialise the set of weights; these will be complex.
val = np.zeros((NU, N, dim), dtype=complex)
count = np.zeros((NU, N, dim), dtype=int)
count2 = np.zeros((NU, N, dim), dtype=int)
ratio = np.zeros((NU, N, dim))
# Start the loop.
for ui in range(0, NU):
# Start by defining the background state.
for j in range(0, N):
U0[j] = U0_set[ui]
H0[j] = -(U0[j] / g) * (f0 * y[j] + beta * y[j]**2 / 2) + Hflat
U0_nd = U0 / U
H0_nd = H0 / chi
# The 1L SW solution
#====================================================
I = np.complex(0.0, 1.0)
# Coefficients
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)
# Solver (free-slip)
solution = solver.FREE_SLIP_SOLVER4(a1, a2, a3, a4, f_nd, b4, c1, c2,
c3, c4, Ro * Ftilde1_nd,
Ro * Ftilde2_nd, Ftilde3_nd, N, N2)
solution = eigDiagnostics.switchKsign(solution, N)
# Analysis
#====================================================
# Loop over desired wavenumbers (for tests, this may not be the full r'ange of wavenumbers)
# ii indexes arrays storing information at ALL wavenumbers k
# i indexes arrays storing information ONLY at wavenumbers used in the decomposition.
a1, a2, a3, a4, b1, b4, c1, c2, c3, c4 = eigSolver.EIG_COEFFICIENTS2(
Ro, Re, K_nd, f_nd, U0_nd, H0_nd, gamma_nd, dy_nd, N)
for i in range(0, N):
k = K_nd[i]
#print('k = ' + str(k))
# Eigenmodes, eigenvalues and count.
#====================================================
# Use free-slip solver
val[ui,
i, :], vec = eigSolver.FREE_SLIP_EIG(a1, a2, a3, a4, b1, b4,
c1, c2, c3, c4, N, N2, i,
False)
# Each eigenmode is currently a unit vector, but we normalise so that each mode contains unit energy.
#==
# Extract the three components.
u_vec, v_vec, h_vec = eigDiagnostics.vec2vecs(vec, N, dim, BC)
# Calculate the contained in each component.
E = np.zeros(dim)
for wi in range(0, dim):
EE = energy.E_anomaly_EIG(u_vec[:, wi], v_vec[:, wi],
h_vec[:, wi], H0_nd, U0_nd, Ro, y_nd,
dy_nd)
# Normalise each vector by the square root of the energy.
u_vec[:,
wi], v_vec[:, wi], h_vec[:, wi] = u_vec[:, wi] / np.sqrt(
EE), v_vec[:,
wi] / np.sqrt(EE), h_vec[:,
wi] / np.sqrt(EE)
# Rebuild the vector. This should have unit energy perturbation.
# (There are more direct ways of executing this normalisation, but this method is the safest.)
vec = eigDiagnostics.vecs2vec(u_vec, v_vec, h_vec, N, dim, BC)
# Order modes by meridional pseudo wavenumber (count).
count[ui, i, :], count2[ui, i, :], ratio[
ui, i, :], i_count = eigDiagnostics.orderEigenmodes(
u_vec, x_nd, k, N, dim)
#count, i_count = eigDiagnostics.orderEigenmodes2(vec,val,N,False)
# Comment out this line, depending on which EIG_COEFFICIENTS function is being called.
#val = val / (2. * np.pi * I * Ro);
#====================================================
Phi = solution[:, i]
# 1. Assign the solution corresponding to wavenumber k=K_nd[ii].
theta[ui, i, :] = np.linalg.solve(vec, Phi)
# 2.
np.save('theta_array' + str(pi), theta)
np.save('val_array' + str(pi), val)
np.save('count_array' + str(pi), count)
np.save('count2_array' + str(pi), count2)
np.save('ratio_array' + str(pi), count)
# Return the weights, eigenvalues, and meridional wavenumber.
#return theta, val, count
return
#====================================================
#====================================================
u = np.zeros((N,N,Nt),dtype=complex);
v = np.zeros((N,N,Nt),dtype=complex);
h = np.zeros((N,N,Nt),dtype=complex);
S = np.load('time_series1.npy');
#Om = np.fft.fftfreq(Nt,dt_nd);
S_tilde = np.fft.fft(S);
for wi in range(1,Nt):
print(wi);
# Coefficients
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)
# Solver
if BC == 'NO-SLIP':
solution = solver.NO_SLIP_SOLVER(a1,a2,a3,a4,f_nd,b4,c1,c2,c3,c4,S_tilde[wi]*Ro*Ftilde1_nd,S_tilde[wi]*Ro*Ftilde2_nd,S_tilde[wi]*Ftilde3_nd,N,N2);
if BC == 'FREE-SLIP':
solution = solver.FREE_SLIP_SOLVER(a1,a2,a3,a4,f_nd,b4,c1,c2,c3,c4,S_tilde[wi]*Ro*Ftilde1_nd,S_tilde[wi]*Ro*Ftilde2_nd,S_tilde[wi]*Ftilde3_nd,N,N2)
u[:,:,wi], v[:,:,wi], h[:,:,wi] = solver.extractSols(solution,N,N2,BC);
u, v, h = solver.SPEC_TO_PHYS_STOCH(u,v,h,dx_nd,N);
# Normalise all solutions by the (non-dimensional) forcing amplitude.
u = u / AmpF_nd;
v = v / AmpF_nd;
h = h / AmpF_nd;
y0_min = y[0] + r0
# We want to keep the forcing at least one gridpoint away from the boundary
y0_max = y[N - 1] - r0
y0_set = []
# Initialise an empty set of forcing latitudes
y0_index_set = []
for j in range(0, N):
if y0_min <= y[j] <= y0_max:
y0_set.append(
y[j]
) # Build the set of forcing locations, all at least 1 gridpoint away from the boundary.
y0_index_set.append(j)
y0_set = np.array(y0_set)
y0_index_set = np.array(y0_index_set)
nn = np.shape(y0_set)[0]
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)
# Initialise the array which stores the EEF values.
if os.path.isfile(filename + '.npy'):
EEF_PV = np.load(filename + '.npy')
else:
EEF_PV = np.zeros((nn, 2))
l_PV = np.zeros((nn, 2))
P_xav = np.zeros((nn, N))
# Correlation
corr_vy_uy = np.zeros((nn, Nt))
corr_v_uyy = np.zeros((nn, Nt))
cs = N / 4
ce = N - N / 4
def RSW_main():
# Forcing
#plotting.forcingPlot_save(x_grid,y_grid,F3_nd[:,0:N],FORCE,BG,Fpos,N);
#F1_nd, F2_nd, F3_nd = forcing.forcingInv(Ftilde1_nd,Ftilde2_nd,Ftilde3_nd,x_nd,y_nd,dx_nd,N);
#F12, F22 = forcing.F12_from_F3(F3_nd,f_nd,dx_nd,dy_nd,N,N);
#plotting.forcingPlots(x_nd[0:N],y_nd,Ro*F1_nd,Ro*F2_nd,F3_nd,Ftilde1_nd,Ftilde2_nd,Ftilde3_nd,N);
# Coefficients
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)
# Solver
if BC == 'NO-SLIP':
solution = solver.NO_SLIP_SOLVER(a1, a2, a3, a4, f_nd, b4, c1, c2, c3,
c4, Ro * Ftilde1_nd, Ro * Ftilde2_nd,
Ftilde3_nd, N, N2)
if BC == 'FREE-SLIP':
#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)
solution = solver.FREE_SLIP_SOLVER4(a1, a2, a3, a4, f_nd, b4, c1, c2,
c3, c4, Ro * Ftilde1_nd,
Ro * Ftilde2_nd, Ro * Ftilde3_nd,
N, N2)
utilde_nd, vtilde_nd, etatilde_nd = solver.extractSols(solution, N, N2, BC)
#utilde_nd, vtilde_nd, etatilde_nd = diagnostics.selectModes(utilde_nd,vtilde_nd,etatilde_nd,6,False,N)
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
#np.save('u.npy',u)
#np.save('v.npy',v)
#np.save('h.npy',h)
#sys.exit()
# In order to calculate the vorticities/energies of the system, we require full (i.e. BG + forced response) u and eta.
u_full = diagnostics.fullFlow(u, U0_nd)
h_full = diagnostics.fullFlow(h, H0_nd)
#====================================================
# Energy
if doEnergy:
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 = KE + PE
#E_tot = KE_tot + PE_tot
Ef, Ef_av = energy.budgetForcing(u_full, v, h_full, F1_nd, F2_nd,
F3_nd, Ro, N, T_nd, omega_nd, Nt)
#Ef, Ef2_av = energy.budgetForcing2(U0_nd,H0_nd,u,v,h,F1_nd,F2_nd,F3_nd,Ro,N,T_nd,omega_nd,Nt)
#Ed, Ed_av = energy.budgetDissipation(u_full,v,h_full,Ro,Re,gamma_nd,dx_nd,dy_nd,T_nd,Nt)
Ed, Ed_av = energy.budgetDissipation2(U0_nd, H0_nd, u, v, h, Ro, Re,
gamma_nd, dx_nd, dy_nd, T_nd, Nt,
N)
Eflux, Eflux_av, uEflux_av, vEflux_av = energy.budgetFlux(
u_full, v, h_full, Ro, dx_nd, dy_nd, T_nd, Nt)
print(np.sum(Ef_av))
print(np.sum(Ed_av))
plt.subplot(221)
plt.contourf(Ef_av)
plt.grid()
plt.colorbar()
plt.subplot(222)
plt.contourf(Ed_av)
plt.grid()
plt.colorbar()
plt.subplot(223)
plt.contourf(vEflux_av)
plt.grid()
plt.colorbar()
plt.subplot(224)
plt.contourf(-Ed_av - Ef_av)
plt.grid()
plt.colorbar()
plt.show()
#uE, vE = energy.flux(KE,u,v)
#Econv, Econv_xav = energy.conv(uE,vE,T_nd,Nt,x_nd,dx_nd,y_nd,dy_nd)
#vE_av = diagnostics.timeAverage(vE,T_nd,Nt)
#plt.contourf(vE_av); plt.colorbar(); plt.show()
#plt.subplot(121); plt.contourf(Econv); plt.colorbar();
#plt.subplot(122); plt.plot(Econv_xav); plt.show()
#quit()
#====================================================
if doCorr:
M = corr.M(u, v, T_nd)
N_ = corr.N(u, v, T_nd)
K = corr.K(u, v, T_nd)
D = corr.D(u, v, 1, dx_nd, dy_nd)
Curl_uD = corr.Curl_uD(u, v, D, T, dx_nd, dy_nd)
theta = corr.orientation(M, N_)
Mnorm = corr.Mnorm(u, v, T_nd)
Nnorm = corr.Nnorm(u, v, T_nd)
Dv, Du = corr.Curl_uD_components(u, v, D, T, dx_nd, dy_nd)
#corr.plotComponents(x_nd,y_nd,M,N_,K,Du)
plotting_bulk.plotKMN(K, Mnorm, Nnorm, x_grid, y_grid, N, 0, 2, '')
plt.show()
#N_ /= diagnostics.domainInt(K,x_nd,dx_nd,y_nd,dy_nd)
N_av = np.trapz(diagnostics.extend(N_), x_nd, dx_nd, axis=1)
Nyy = diagnostics.diff(diagnostics.diff(N_, 0, 0, dy_nd), 0, 0, dy_nd)
Nyy_av = np.trapz(diagnostics.extend(Nyy), x_nd, dx_nd, axis=1)
Curl_uD_av = np.trapz(diagnostics.extend(Curl_uD), x_nd, dx_nd, axis=1)
#np.save('M',M); np.save('N',N_); np.save('Nyy',Nyy)
lim = np.max(np.abs(N_)) / 2.
plt.figure(figsize=[12, 6])
plt.subplot(121)
plt.pcolor(x_grid, y_grid, N_, cmap='bwr', vmin=-lim, vmax=lim)
plt.xlim(-0.1, 0.1)
plt.ylim(-.1, 0.1)
plt.text(-0.08, 0.08, 'N', fontsize=18)
plt.xlabel('x', fontsize=18)
plt.ylabel('y', fontsize=18)
plt.grid()
plt.colorbar()
plt.subplot(122)
plt.plot(N_av, y_nd)
plt.ylim(-.1, .1)
plt.xlabel('<N>', fontsize=18)
plt.grid()
plt.tight_layout()
plt.show()
#plt.figure(figsize=[12,6])
#plt.subplot(121)
#plt.contourf(x_nd[0:N],y_nd,Nyy)
#plt.xlim(-0.1,0.1)
#plt.ylim(-0.1,0.1)
#plt.colorbar()
#plt.subplot(122)
#plt.plot(Nyy_av,y_nd)
#plt.show()
#corr.plotOrientation(theta,K,x_nd,y_nd)
#uav = np.trapz(diagnostics.extend(Du),x_nd,dx_nd,axis=1)
#vav = np.trapz(diagnostics.extend(Dv),x_nd,dx_nd,axis=1)
#plt.plot(uav,label='u')
#plt.plot(vav,label='v')
#plt.plot(Curl_uD_av,label='full')
#plt.legend()
#plt.show()
# Correlation.
# Central half?
#cs = N / 4;
#ce = N - N / 4;
#corr = corr.arrayCorrTime(u[cs:ce,cs:ce,:],v[cs:ce,cs:ce,:]);
#print corr
#quit()
#====================================================
# Error - if calculated, should be done before real part of solution is taken
if errorPhys:
e1, e2, e3 = diagnostics.error(u, v, h, dx_nd, dy_nd, dt_nd, U0_nd,
H0_nd, Ro, gamma_nd, Re, f_nd, F1_nd,
F2_nd, F3_nd, T_nd, ts, omega_nd, N)
e = np.sqrt((e1**2 + e2**2 + e3**2) / 3.0)
print 'Error = ' + str(e) + '. Error split = ' + str(e1) + ', ' + str(
e2) + ', ' + str(e3)
if errorSpec:
error_spec = np.zeros((3, N))
# An array to save the spectral error at each wavenumber for each equation.
for i in range(0, N):
error_spec[:,i] = diagnostics.specError(utilde_nd[:,i],vtilde_nd[:,i],etatilde_nd[:,i],Ftilde1_nd[:,i],Ftilde2_nd[:,i],Ftilde3_nd[:,i],a1[:,i],a2,a3,a4[i],\
b4,c1[:,i],c2,c3,c4[:,i],f_nd,Ro,K_nd[i],H0_nd,y_nd,dy_nd,N)
for eq in range(0, 3):
error = sum(error_spec[eq, :]) / N
print('Error' + str(int(eq + 1)) + '=' + str(error))
#====================================================
# Momentum footprints
#====================================================
if doMomentum:
uu, uv, vv = momentum.fluxes(u, v)
Mu, Mv, Mu_xav, Mv_xav = momentum.footprint(uu, uv, vv, x_nd, T_nd,
dx_nd, dy_nd, N, Nt)
#plotting.MomFootprints(Mu,Mv,Mu_xav,Mv_xav);
Mumax = np.max(Mu_xav)
plt.plot(Mu_xav / Mumax, y_nd, linewidth=2., color='k')
plt.text(-0.4, 0.4, str(Mumax))
plt.xlabel('Zonal mom. flux convergence', fontsize=18)
plt.ylabel('y', fontsize=18)
plt.ylim([-.5, .5])
plt.yticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.grid()
plt.show()
if False:
plt.subplot(121)
plt.pcolor(x_grid, y_grid, Mu, cmap='bwr', vmin=-.5, vmax=.5)
plt.xticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.yticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.xlabel('x', fontsize=16)
plt.ylabel('y', fontsize=16)
plt.axis([x_grid.min(),
x_grid.max(),
y_grid.min(),
y_grid.max()])
plt.grid(b=True, which='both', color='0.65', linestyle='--')
plt.colorbar()
plt.subplot(122)
plt.plot(Mu_xav, y_nd)
plt.xlabel('y')
plt.show()
plt.subplot(121)
plt.pcolor(x_grid, y_grid, Mv, cmap='bwr', vmin=-1., vmax=1.)
plt.xticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.yticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.xlabel('x', fontsize=16)
plt.ylabel('y', fontsize=16)
plt.axis([x_grid.min(),
x_grid.max(),
y_grid.min(),
y_grid.max()])
plt.grid(b=True, which='both', color='0.65', linestyle='--')
plt.subplot(122)
plt.plot(Mv_xav, y_nd)
plt.tight_layout()
plt.show()
EEF_u, EEF_v = momentum.EEF_mom(Mu_xav, Mv_xav, y_nd, y0_nd, y0_index,
dy_nd, omega_nd, N)
#print(EEF_u, EEF_v);
# PV and PV footprints
#====================================================
# Calculate PV fields, footprints and equivalent eddy fluxes (EEFs)
if doPV:
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)
#PV_prime1, PV_prime2, PV_prime3 = PV.potentialVorticity_linear(u,v,h,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)
# Keep these next two lines commented out unless testing effects of normalisation.
# uq, Uq, uQ, UQ, vq, vQ = uq/AmpF_nd**2, Uq/AmpF_nd**2, uQ/AmpF_nd**2, UQ/AmpF_nd**2, vq/AmpF_nd**2, vQ/AmpF_nd**2
# PV_prime, PV_full = PV_prime/AmpF_nd, PV_full/AmpF_nd
if doFootprints:
if footprintComponents:
P, P_uq, P_uQ, P_Uq, P_vq, P_vQ, P_xav, P_uq_xav, P_uQ_xav, P_Uq_xav, P_vq_xav, P_vQ_xav = PV.footprintComponents(
uq, Uq, uQ, vq, vQ, x_nd, T_nd, dx_nd, dy_nd, N, Nt)
#plotting.footprintComponentsPlot(uq,Uq,uQ,vq,vQ,P,P_uq,P_Uq,P_uQ,P_vq,P_vQ,P_xav,P_uq_xav,P_uQ_xav,P_Uq_xav,P_vq_xav,P_vQ_xav,x_nd,y_nd,N,Nt);
#plotting.plotPrimaryComponents(P_uq,P_vq,P_uq_xav,P_vq_xav,x_nd,y_nd,FORCE,BG,Fpos,N);
else:
P, P_xav = PV.footprint(uq, Uq, uQ, UQ, vq, vQ, x_nd, T_nd,
dx_nd, dy_nd, N, Nt)
if doEEFs:
from scipy.ndimage.measurements import center_of_mass
iii = center_of_mass(np.abs(P_xav))[0]
i1 = int(iii)
i2 = int(i1 + 1)
r = iii - i1
#print(iii,i1,i2,r)
com = y_nd[int(iii)]
#print(y0_index-iii)
if footprintComponents:
EEF_array = PV.EEF_components(P_xav, P_uq_xav, P_uQ_xav,
P_Uq_xav, P_vq_xav, P_vQ_xav,
y_nd, y0_nd, y0_index, dy_nd,
omega_nd, N)
# This returns EEF_array, an array with the following structure:
# EEF_array = ([EEF_north,EEF_south],[uq_north,uq_south],[Uq_north,Uq_south],[uQ_north,uQ_south],[vq_north,vq_south],[vQ_north,vQ_south]).
EEF_north = EEF_array[0, 0]
EEF_south = EEF_array[0, 1]
else:
EEF, l = PV.EEF(P_xav, y_nd, com, int(iii), dy_nd, N)
# These lines for Gaussian EEFs, when com jumps from 1 grid point to next, need to smooth EEF.
#EEF1, l = PV.EEF(P_xav,y_nd,y_nd[i1],i1,dy_nd,N)
#EEF2, l = PV.EEF(P_xav,y_nd,y_nd[i2],i2,dy_nd,N)
#EEF_ = (1 - r) * EEF1 + r * EEF2
EEF_north = EEF[0]
EEF_south = EEF[1]
EEF = EEF_north - EEF_south
print(EEF)
Pmax = np.max(abs(P_xav))
plt.plot(P_xav / Pmax, y_nd, linewidth=2., color='k')
plt.text(-0.4, 0.4, str(Pmax))
plt.xlabel('PV flux convergence', fontsize=18)
plt.ylabel('y', fontsize=18)
plt.ylim([-.5, .5])
plt.yticks((-1. / 2, -1. / 4, 0, 1. / 4, 1. / 2))
plt.grid()
plt.show()
# Buoyancy footprints
#====================================================
if doThickness:
# Should these be zero, according to conservation of mass?
Pb, Pb_xav = thickness.footprint(u_full, v, h_full, x_nd, y_nd, T_nd,
dx_nd, dy_nd, dt_nd, N, Nt)
#output.ncSave(utilde_nd,vtilde_nd,etatilde_nd,u,v,h,x_nd,y_nd,K_nd,T_nd,PV_full,PV_prime,PV_BG,Pq,EEFq,N,Nt);
#sys.exit()
#====================================================
# Plots
#====================================================
#====================================================
# Call the function that plots the forcing in physical and physical-spectral space.
if plotForcing:
plotting.forcingPlots(x_nd, y_nd, F1_nd, F2_nd, F3_nd, Ftilde1_nd,
Ftilde2_nd, Ftilde3_nd, N)
#forcing_1L.forcingInv(Ftilde1_nd,Ftilde2_nd,Ftilde3_nd,x_nd,y_nd,dx_nd,N); # For diagnostic purposes
# Background state plots (inc. BG SSH, BG flow, BG PV)
if plotBG:
plotting.bgPlots(y_nd, H0_nd, U0_nd, PV_BG)
# Soltuion Plots
if plotSol:
plotting.solutionPlots(x_nd, y_nd, u, v, h, ts, FORCE, BG, Fpos, N,
x_grid, y_grid, True)
#plotting.solutionPlots_save(x_nd,y_nd,u,v,h,ts,FORCE,BG,Fpos,N,x_grid,y_grid,True);
#plotting.solutionPlotsDim(x,y,u,v,eta,ts,L,FORCE,BG,Fpos,N);
# Plots of PV and zonally averaged PV
if doPV:
if plotPV:
#plotting.pvPlots(PV_full,PV_prime,x_nd,y_nd);
plotting.pvPlots_save(PV_full, PV_prime, P, P_xav, x_nd, y_nd, ts,
FORCE, BG, Fpos, N, U0_str, x_grid, y_grid,
True)
if plotPV_av:
plotting.PV_avPlots(x_nd, y_nd, PV_prime, PV_BG, PV_full, ts,
FORCE, BG, Fpos, N)
if doFootprints:
if plotFootprint:
plotting.footprintPlots(x_nd, y_nd, P, P_xav, Fpos, BG, FORCE,
nu, r0, period_days, U0_nd, U, N)
# Phase and amplitude
if plotPhaseAmp:
plotting.solutionPlotsAmp(x_nd, y_nd, u, v, h, ts, FORCE, BG, Fpos, N)
plotting.solutionPlotsPhase(x_nd, y_nd, u, v, h, ts, FORCE, BG, Fpos,
N)
def EEF_main(set_, pi):
NU = len(set_)
# Initialise output arrays
EEF_array = np.zeros((NU, nn, 2))
com = np.zeros((NU, nn))
E_array = np.zeros((NU))
M_array = np.zeros((NU, N, N))
N_array = np.zeros((NU, N, 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_Gaussian(Umag,sigma,JET_POS,Hflat,f0,beta,g,y,L,N)
#U0_nd = U0 / U;
#H0_nd = H0 / chi;
# r0
#r0 = set_[ui]
#r0_nd = r0 / L
# period
period_days = set_[ui]
period = 3600. * 24. * period_days
omega = 1. / (period)
T = np.linspace(0, period, Nt + 1)
dt = T[1] - T[0]
t = T[ts]
omega_nd = omega * T_adv
t_nd = t / T_adv
T_nd = T / T_adv
dt_nd = dt / T_adv
# k
#nu = set_[ui]
#Re = L * U / nu
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)
for yi in range(0, nn):
y0 = y0_set[yi]
# Redefine y0 and the forcing in each run.
y0_index = y0_index_set[yi]
y0_nd = y0 / L
# Forcing
F1_nd, F2_nd, F3_nd, Ftilde1_nd, Ftilde2_nd, Ftilde3_nd = forcing.forcing_cts2(
x_nd, y_nd, K_nd, y0_nd, r0_nd, N, FORCE, AmpF_nd, f_nd, f0_nd,
bh, dx_nd, dy_nd)
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)
#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)
com[ui, yi] = center_of_mass(np.abs(P_xav))[0]
i1 = int(com[ui, yi])
i2 = int(i1 + 1)
r = com[ui, yi] - i1
# For Gaussian flows, need to calculate EEF about new center of mass.
# This requires calculation of EEF at two grid points to aid continuity.
#EEF1, l_PV = PV.EEF(P_xav,y_nd,y_nd[i1],i1,dy_nd,N)
#EEF2, l_PV = PV.EEF(P_xav,y_nd,y_nd[i2],i2,dy_nd,N)
#EEF_array[ui,yi,:] = (1 - r) * EEF1 + r * EEF2
EEF_array[ui, yi, :], l_PV = PV.EEF(P_xav, y_nd, y_nd[y0_index],
y0_index, dy_nd, N)
filename = 'EEF_array' + str(pi)
np.save(filename, EEF_array)
filename_com = 'com' + str(pi)
np.save(filename_com, com)
def EIG_PROJ_main(dim):
'''
This function projects the SW solution onto eigenmodes as outputted by EIG.py.
The aim is to look at the number of modes required for certain levels of accuracy.
To do this we can look at the error in the whole solution (x,y), or look at the
error at specific wavenumbers (e.g. wavenumbers where there is a lot of weight).
'''
# Dimensions
Nm = dim # How many modes to use in the decomposition at each wavenumber (dim is maximum).
Nk = N
Ke = 10 # Wavenumber to look at error in spectral space.
Np = 50 # Number of modes to keep in decomposition.
# The 1L SW solution
#====================================================
I = np.complex(0.0, 1.0)
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)
# Define the solution in (k,y)-space - can be from FILE or a NEW run.
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) / AmpF_nd
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)
u = np.real(u)
v = np.real(v)
h = np.real(h)
print('solved')
#====================================================
# Perform decomposition and build projection.
#====================================================
theta = np.zeros((dim, N), dtype=complex)
theta_abs = np.zeros((dim, N))
proj = np.zeros((dim, N), dtype=complex)
projk = np.zeros((dim), dtype=complex)
for i in range(233, 234):
print(i)
k = K_nd[i]
# Load modes
path = '/home/mike/Documents/GulfStream/RSW/DATA/1L/EIG/256/04/'
#path = ''
ncFile = path + 'RSW1L_Eigenmodes_k' + str(int(k)) + '_N257.nc'
val, vec, count = output_read.ncReadEigenmodes(ncFile)
Phi = solution[:,
i] # 1. Assign the solution corresponding to wavenumber k=K_nd[ii].
theta_tmp = np.linalg.solve(vec, Phi) # 2.
theta_abs_tmp = np.abs(theta_tmp)
dom_index = np.argsort(
-theta_abs_tmp
) # 3. The indices of the modes, ordered by 'dominance'.
theta[:, i] = theta_tmp[dom_index]
theta_abs[:, i] = np.abs(theta[:, i])
vec = vec[:, dom_index]
# Now loop over each mode (at wavenumber k)
plt.plot(np.abs(theta[:, i]))
plt.show()
np.save('mode1', vec[:, 0])
np.save('mode2', vec[:, 1])
np.save('mode3', vec[:, 2])
np.save('mode4', vec[:, 3])
theta0123 = theta[0:4, i]
np.save('theta0123', theta0123)
for mi in range(0, 100):
proj[:, i] = proj[:, i] + theta[mi, i] * vec[:, mi]
theta_abs_tot = np.sum(theta_abs, axis=0)
utilde_proj = np.zeros((N, N), dtype=complex)
vtilde_proj = np.zeros((N, N), dtype=complex)
etatilde_proj = np.zeros((N, N), dtype=complex)
# Separate projection into flow components.
for j in range(0, N):
utilde_proj[j, :] = proj[j, :]
etatilde_proj[j, :] = proj[N + N2 + j, :]
for j in range(0, N2):
vtilde_proj[j + 1, :] = proj[N + j, :]
u_proj, v_proj, eta_proj = solver.SPEC_TO_PHYS(utilde_proj, vtilde_proj,
etatilde_proj, T_nd, dx_nd,
omega_nd, N)
u_proj = np.real(u_proj)
v_proj = np.real(v_proj)
eta_proj = np.real(eta_proj)
c = np.corrcoef(u[:, :, ts].reshape(N**2), u_proj[:, :,
ts].reshape(N**2))[0, 1]
print(c)
#====================================================
# Plot
#====================================================
if True:
plt.plot(utilde_proj[:, 3])
plt.plot(utilde_nd[:, 3])
plt.show()
plt.subplot(121)
plt.contourf(u[:, :, ts])
plt.colorbar()
plt.subplot(122)
plt.contourf(u_proj[:, :, ts])
plt.colorbar()
plt.show()
# Compare projection with solution.
#====================================================
corrk = np.zeros(N)
#for i in range(1,N):
# corrk[i] = np.corrcoef(utilde_nd[:,i],utilde_proj[:,i])[0,1]
return u, v, h, u_proj, v_proj, eta_proj
def EIG_DECOMP_main(U0_nd, H0_nd, dim):
'''
Load outputted theta and val and post-process. Function does multiple things.
(1) For a given U0-value calculate the SW solution. Rebuild solution with Nm-most dominant modes
to see how many are required for p % accuracy. Does this depend on k, U0?
(2) Plots of weights at specific wavenumber (and how quickly they decay), and plots of weights over all wavenumbers and BG flows.
(3) Phase speeds and 'average' phase speeds.
'''
# Dimensions
Nm = dim
# How many modes to use in the decomposition at each wavenumber (dim is maximum).
Nk = N
# The 1L SW solution
#====================================================
I = np.complex(0.0, 1.0)
SOL = 'NEW'
# Define the solution in (k,y)-space - can be from FILE or a NEW run.
if SOL == 'FILE':
solution = np.load(
'/home/mike/Documents/GulfStream/Code/DATA/1L/REF/solution_NORTH256.npy'
)
if SOL == 'NEW':
# Forcing
#if FORCE_TYPE == 'CTS':
# F1_nd, F2_nd, F3_nd, Ftilde1_nd, Ftilde2_nd, Ftilde3_nd = forcing.forcing_cts(x_nd,y_nd,K_nd,y0_nd,r0_nd,N,FORCE,AmpF_nd,f_nd,f0_nd,dx_nd,dy_nd);
#elif FORCE_TYPE == 'DCTS':
# F1_nd, F2_nd, F3_nd, Ftilde1_nd, Ftilde2_nd, Ftilde3_nd = forcing.forcing_dcts(x_nd,y_nd,K_nd,y0_nd,r0_nd,N,FORCE,AmpF_nd,f_nd,f0_nd,dx_nd,dy_nd)
#elif FORCE_TYPE == 'DELTA':
# F1_nd, F2_nd, F3_nd, Ftilde1_nd, Ftilde2_nd, Ftilde3_nd = forcing.forcing_delta(AmpF_nd,y0_index,dx_nd,N);
# Coefficients
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)
# Solver
if BC == 'NO-SLIP':
solution = solver.NO_SLIP_SOLVER(a1, a2, a3, a4, f_nd, b4, c1, c2,
c3, c4, Ro * Ftilde1_nd,
Ro * Ftilde2_nd, Ftilde3_nd, N,
N2)
elif BC == 'FREE-SLIP':
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)
else:
sys.exit('ERROR: choose valid BC')
solution = solution / AmpF_nd
print('solved')
#====================================================
# Eigenmode analysis
#====================================================
#====================================================
# Initisialastion steps
#====================================================
VEC = 'FILE'
# From FILE, requires pre-saved vectors which take up lots of memory.
LOOP = 'FULL'
# FULL, PART
if LOOP == 'FULL':
loop = range(0, N)
Nk_neg = 9
Nk_pos = 8
elif LOOP == 'PART':
Nk_neg = 6
Nk_pos = 6
# How many positive/negative wavenumbers to perform this decomposition at.
loop = it.chain(range(0, Nk_pos + 1), range(N - Nk_neg, N))
Nk = Nk_neg + Nk_pos + 1
else:
sys.exit('ERROR: LOOP must be FULL or PART.')
theta = np.zeros((Nm, Nk), dtype=complex)
# Initialise the set of weights; these will be complex.
proj = np.zeros((dim, N), dtype=complex)
# The projection. Sums the Nm most dominant modes, each of length dim, for Nk i-values.
dom_index = np.zeros((Nm, Nk), dtype=int)
# To store the indices of the Nm-most dominant modes.
var = np.zeros((Nk))
mean = np.zeros((Nk))
scatter_k = np.zeros(Nm * Nk)
# An empty array for saving k-values, for use in the scatter plot of dominant modes.
scatter_l = np.zeros(Nm * Nk)
# An empty array for storing the count, psuedo-wavenumber l.
scatter_p = np.zeros(Nm * Nk)
# An empty array for storing periods of the dominant wavenumbers.
theta_abs = np.zeros((Nm, Nk))
# For storing the absolute value of each weight.
theta_abs_tot = np.zeros(Nk)
# For storing sum of absolute values of each set of decomposition weights.
cx = np.zeros((Nm, Nk))
# For storing the zonal phase speed of each mode,
cy = np.zeros((Nm, Nk))
# and the meridional phase speed.
p = np.zeros((Nk, 2))
# For storing weighted phase speeds at each wavenumber. (i=0 => y)
# Analysis
#====================================================
# Loop over desired wavenumbers (for tests, this may not be the full range of wavenumbers)
# ii indexes arrays storing information at ALL wavenumbers k
# i indexes arrays storing information ONLY at wavenumbers used in the decomposition.
for ii in loop:
print(' ')
print('ii = ' + str(ii))
k = K_nd[ii]
print('k = ' + str(k))
i = k % Nk
print('i = ' + str(i))
i = int(i + 0.01)
# Eigenmodes, eigenvalues and count.
#====================================================
# This section returns three arrays: 1. val, 2. vec, 3. count
# 1.) val = val[0:dim] stores the eigenvalues/frequencies.
# 2.) vec = vec[]
# 3.) count = count[]
# Run the solver for the current k-value.
if VEC == 'NEW': # Solve the eigenmode problem anew.
a1, a2, a3, a4, b1, b4, c1, c2, c3, c4 = eigSolver.EIG_COEFFICIENTS2(
Ro, Re, K_nd, f_nd, U0_nd, H0_nd, gamma_nd, dy_nd, N)
if BC == 'NO-SLIP':
val, u_vec, v_vec, h_vec = eigSolver.NO_SLIP_EIG(
a1, a2, a3, a4, b1, b4, c1, c2, c3, c4, N, N2, ii, True)
if BC == 'FREE-SLIP':
val, vec = eigSolver.FREE_SLIP_EIG(a1, a2, a3, a4, b1, b4, c1,
c2, c3, c4, N, N2, ii,
False)
# Order modes by meridional pseudo wavenumber (count).
count, i_count = eigDiagnostics.orderEigenmodes(
vec, val, x_nd, k, T_nd[ts], N, dim, BC)
#count, i_count = eigDiagnostics.orderEigenmodes2(vec,val,N,False)
count = count[i_count]
vec = vec[:, i_count]
val = val[i_count]
# Each eigenmode is currently a unit vector, but we normalise so that each mode contains unit energy.
#==
# Extract the three components.
u_vec, v_vec, h_vec = eigDiagnostics.vec2vecs(vec, N, dim, BC)
# Calculate the contained in each component.
E = np.zeros(dim)
for wi in range(0, dim):
EE = energy.E_anomaly_EIG(u_vec[:, wi], v_vec[:, wi],
h_vec[:, wi], H0_nd, U0_nd, Ro, y_nd,
dy_nd)
# Normalise each vector by the square root of the energy.
u_vec[:,
wi], v_vec[:, wi], h_vec[:, wi] = u_vec[:, wi] / np.sqrt(
EE), v_vec[:,
wi] / np.sqrt(EE), h_vec[:,
wi] / np.sqrt(EE)
# Rebuild the vector. This should have unit energy perturbation.
# (There are more direct ways of executing this normalisation, but this method is the safest.)
vec = eigDiagnostics.vecs2vec(u_vec, v_vec, h_vec, N, dim, BC)
# Comment out this line, depending on which EIG_COEFFICIENTS function is being called.
#val = val / (2. * np.pi * I * Ro);
elif VEC == 'FILE': # Load eigenmodes and eigenvalues from file.
# Low-res
path = '/home/mike/Documents/GulfStream/RSW/DATA/1L/EIG/128/west/'
#path = '/home/mike/Documents/GulfStream/RSW/DATA/1L/EIG/128/nu='+str(int(nu))+'/';
ncFile = path + 'RSW1L_Eigenmodes_k' + str(int(k)) + '_N129.nc'
# High-res
#path = '/media/mike/Seagate Expansion Drive/Documents/GulfStream/RSW/DATA/1L/EIG/256/16/'
#path = '/home/mike/Documents/GulfStream/RSW/DATA/1L/EIG/256/west/';
#ncFile = path + 'RSW1L_Eigenmodes_k' + str(int(k)) + '_N257.nc';
print('Reading from ' + ncFile + '...')
val, vec, count = output_read.ncReadEigenmodes(ncFile)
else:
sys.exit('VEC must be FILE or NEW')
# Expresses the eigenvalues (frequencies) in terms of periods (units days).
freq = np.real(val)
period_days = T_adv / (freq * 24. * 3600.)
dim = np.size(val)
#====================================================
# Now we have the solution and the eigenmodes.
# The decomposition follows the following steps:
# 1. Define the solution to be decomposed as Phi.
# 2. Decompose Phi into vec using a linear solver; theta_tmp stores the weights.
# 3. Arrange the weights in descending order according to their complex amplitude.
# 4. Sum the Nm-most dominant weights.
Phi = solution[:, ii]
# 1. Assign the solution corresponding to wavenumber k=K_nd[ii].
theta_tmp = np.linalg.solve(vec, Phi)
# 2.
theta_abs_tmp = np.abs(theta_tmp)
dom_index_tmp = np.argsort(-theta_abs_tmp)
# 3. The indices of the modes, ordered by 'dominance'.
theta_abs_tot[i] = sum(theta_abs_tmp[dom_index_tmp[0:dim]])
# Change dim to Nm if necessary
# Now loop over each mode (at wavenumber k)
for mi in range(0, Nm):
#print(dom_index_tmp[mi]);
#print('count = ' + str(count[dom_index_tmp[mi]]));
#print(np.abs(theta_tmp[dom_index_tmp[mi]]));
dom_index[mi, i] = dom_index_tmp[mi]
theta[mi, i] = theta_tmp[dom_index_tmp[mi]]
# All weights are now ordered in terms of their absolute value.
# Absolute value of each mode
theta_abs[mi, i] = np.abs(theta[mi, i])
# Zonal & meridional phase speed of each mode
cx[mi, i] = freq[dom_index[mi, i]] / k
if count[dom_index[mi, i]] != 0:
cy[mi, i] = freq[dom_index[mi, i]] / count[dom_index[mi, i]]
if mi < 95:
# The projection.
proj[:, ii] = proj[:, ii] + theta_tmp[
dom_index_tmp[mi]] * vec[:, dom_index_tmp[mi]]
# 4.
# Scatter plot arrays.
scatter_k[i * Nm + mi] = k
scatter_l[i * Nm + mi] = count[dom_index[mi, i]]
scatter_p[i * Nm + mi] = period_days[dom_index[mi, i]]
#plt.plot(vec[0:N,dom_index_tmp[mi]],y_nd);
#plt.ylim(-0.5,0.5);
#plt.show();
#plt.plot(theta_abs[:,i]);
#plt.show();
# Statistics: mean, variance, weighted average
# Should normalise so that all have the same mean (i.e. mean = 1/Nm);
mean[i] = theta_abs_tot[i] / Nm
var[i] = sum((theta_abs[:, i] - mean[i])**2) / (Nm * mean[i]**2)
p[i, 1] = sum((theta_abs[:, i] * cx[:, i])) / theta_abs_tot[i]
p[i, 0] = sum((theta_abs[:, i] * cy[:, i])) / theta_abs_tot[i]
print(cy[10, i])
return theta, mean, var, p, proj, solution, scatter_k, scatter_l, scatter_p
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)
def RSW_main():
# Coefficients
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)
# Solver
if BC == 'NO-SLIP':
solution = solver.NO_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)
if BC == 'FREE-SLIP':
solution = solver.FREE_SLIP_SOLVER4(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)
#===================================================
utilde, vtilde, htilde = solver.extractSols(solution, N, N2, BC)
u, v, h = solver.SPEC_TO_PHYS(utilde, vtilde, htilde, T, Nt, dx, omega, N)
# Before taking real part, can define an error calculator to call here.
u = np.real(u)
v = np.real(v)
h = np.real(h)
#u = u / AmpF_nd
#v = v / AmpF_nd
#h = h / AmpF_nd
# 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 each 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)
q[:, :, :,
1], q_full[:, :, :,
1], Q[:, 1] = PV.vort(u[:, :, :, 1], v[:, :, :, 1], h[:, :, :,
1],
u_full[:, :, :, 1], h_full[:, :, :, 1],
H2, U2, N, Nt, dx, dy, f)
# Calculate footprints using previously calculated PV. Most interseted in the upper layer.
P, P_xav = PV.footprint(u_full[:, :, :, 0], v[:, :, :, 0],
q_full[:, :, :, 0], x, y, dx, dy, T, Nt)
# PLOTS
#====================================================
plt.contourf(u[:, :, 0, 1])
plt.colorbar()
plt.show()
quit()
plotting.solutionPlots(x, y, x_grid, y_grid, u, v, h, ts, N, False)
plotting.footprintPlots(x, y, P, P_xav)
def RSW_main():
# Coefficients
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_nd, f_nd, U1_nd, U2_nd, H1_nd, H2_nd, rho1_nd, rho2_nd,
omega_nd, gamma_nd, dy_nd, N)
# Solver
if BC == 'NO-SLIP':
u1tilde_nd, u2tilde_nd, v1tilde_nd, v2tilde_nd, eta0tilde_nd, eta1tilde_nd = (
solver.NO_SLIP_SOLVER(a1, a2, a3, a4, b1, b4, c1, c2, c3, c4, c5,
d1, d3, d4, d5, e4, e5, f1, f2, f3, f4,
Ftilde1_nd, Ftilde2_nd, Ftilde3_nd,
Ftilde4_nd, Ftilde5_nd, Ftilde6_nd, N, N2))
if BC == 'FREE-SLIP':
u1tilde_nd, u2tilde_nd, v1tilde_nd, v2tilde_nd, eta0tilde_nd, eta1tilde_nd = (
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,
Ftilde1_nd, Ftilde2_nd, Ftilde3_nd,
Ftilde4_nd, Ftilde5_nd, Ftilde6_nd, N, N2))
#===================================================
u1_nd, u2_nd, v1_nd, v2_nd, eta0_nd, eta1_nd = solver.SPEC_TO_PHYS(
u1tilde_nd, u2tilde_nd, v1tilde_nd, v2tilde_nd, eta0tilde_nd,
eta1tilde_nd, T_nd, Nt, dx_nd, omega_nd, N)
# Before taking real part, can define an error calculator to call here.
u1_nd = np.real(u1_nd)
u2_nd = np.real(u2_nd)
v1_nd = np.real(v1_nd)
v2_nd = np.real(v2_nd)
eta0_nd = np.real(eta0_nd)
eta1_nd = np.real(eta1_nd)
# The interface thicknesses defined via the interface heights.
h1_nd = eta0_nd - eta1_nd
h2_nd = eta1_nd
# For use in PV and footprint calculations: the 'full' zonal velocities and interface thicknesses.
u1_full = np.zeros((N, N, Nt))
u2_full = np.zeros((N, N, Nt))
h1_full = np.zeros((N, N, Nt))
h2_full = np.zeros((N, N, Nt))
for j in range(0, N):
u1_full[:, j, :] = u1_nd[:, j, :] + U1_nd[j]
u2_full[:, j, :] = u2_nd[:, j, :] + U2_nd[j]
h1_full[:, j, :] = h1_nd[:, j, :] + H1_nd[j]
h2_full[:, j, :] = h2_nd[:, j, :] + H2_nd[j]
# Call function calculate PV in each layer.
#PV1_prime, PV1_full, PV1_BG = PV.vort(u1_nd,v1_nd,h1_nd,u1_full,h1_full,H1_nd,U1_nd,N,Nt,dx_nd,dy_nd,f_nd);
#PV2_prime, PV2_full, PV2_BG = PV.vort(u2_nd,v2_nd,h2_nd,u2_full,h2_full,H2_nd,U2_nd,N,Nt,dx_nd,dy_nd,f_nd);
# Calculate footprints using previously calculated PV. Most interseted in the upper layer.
#P, P_xav = PV.footprint(u1_full,v1_nd,PV1_full,U1_nd,U1,x_nd,y_nd,dx_nd,dy_nd,AmpF_nd,FORCE1,r0,nu,BG1,Fpos,ts,period_days,N,Nt,GAUSS);
# PLOTS
#====================================================
plotting.solutionPlots(x_nd, y_nd, x_grid, y_grid, u1_nd, u2_nd, v1_nd,
v2_nd, h1_nd, h2_nd, ts, N, True)
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)
def EIG_DECOMP_main(y0_set, pi):
I = np.complex(0.0, 1.0)
# Need to loop over a set of background states.
Ny = len(y0_set)
theta = np.zeros(
(Ny, N, dim),
dtype=complex) # Initialise the set of weights; these will be complex.
# Only theta will have Ny dimension. Only need one set of eigenmodes, so no dependence on forcing latitude.
val = np.zeros((N, dim), dtype=complex)
count = np.zeros((N, dim), dtype=int)
count2 = np.zeros((N, dim), dtype=int)
ratio = np.zeros((N, dim))
# Find solution at every y0-value and for all i. Store it.
#========================================================
# Coefficients - no dependence on y0
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 = np.zeros((Ny, dim, N), dtype=complex)
# Start the loop.
for yi in range(0, Ny):
# Get forcing latitude.
y0_nd = y0_set[yi]
# Redefine forcing, moving one gridpoint each iteration.
F1_nd, F2_nd, F3_nd, Ftilde1_nd, Ftilde2_nd, Ftilde3_nd = forcing.forcing_cts2(
x_nd, y_nd, K_nd, y0_nd, r0_nd, N, FORCE, AmpF_nd, f_nd, f0_nd, bh,
dx_nd, dy_nd)
# The 1L SW solution.
# Solver (free-slip)
solution[yi, ] = solver.FREE_SLIP_SOLVER4(a1, a2, a3, a4, f_nd, b4, c1,
c2, c3, c4, Ro * Ftilde1_nd,
Ro * Ftilde2_nd, Ftilde3_nd,
N, N2)
solution[yi, ] = eigDiagnostics.switchKsign(solution[yi, ], N)
# Done finding solution.
#====================================================
# Find eigenmodes and decompose solution.
#====================================================
# Loop over desired wavenumbers (for tests, this may not be the full r'ange of wavenumbers)
# Eig coefficients don't change. Define them just once.
a1, a2, a3, a4, b1, b4, c1, c2, c3, c4 = eigSolver.EIG_COEFFICIENTS2(
Ro, Re, K_nd, f_nd, U0_nd, H0_nd, gamma_nd, dy_nd, N)
for i in range(0, N):
k = K_nd[i]
#print('k = ' + str(k))
# Eigenmodes, eigenvalues and count.
#====================================================
# Use free-slip solver
val[i, :], vec = eigSolver.FREE_SLIP_EIG(a1, a2, a3, a4, b1, b4, c1,
c2, c3, c4, N, N2, i, False)
# Each eigenmode is currently a unit vector, but we normalise so that each mode contains unit energy.
#==
# Extract the three components.
u_vec, v_vec, h_vec = eigDiagnostics.vec2vecs(vec, N, dim, BC)
# Calculate the contained in each component.
E = np.zeros(dim)
for wi in range(0, dim):
EE = energy.E_anomaly_EIG(u_vec[:, wi], v_vec[:, wi], h_vec[:, wi],
H0_nd, U0_nd, Ro, y_nd, dy_nd)
# Normalise each vector by the square root of the energy.
u_vec[:, wi], v_vec[:, wi], h_vec[:, wi] = u_vec[:, wi] / np.sqrt(
EE), v_vec[:, wi] / np.sqrt(EE), h_vec[:, wi] / np.sqrt(EE)
# Rebuild the vector. This should have unit energy perturbation.
# (There are more direct ways of executing this normalisation, but this method is the safest.)
vec = eigDiagnostics.vecs2vec(u_vec, v_vec, h_vec, N, dim, BC)
# Eigenmodes now have unit time-mean energy.
#==
# Meridional pseudo wavenumber (count).
count[i, :], count2[i, :], ratio[
i, :], i_count = eigDiagnostics.orderEigenmodes(
u_vec, x_nd, k, N, dim)
# Now decompose for each forcing latitude.
for yi in range(0, Ny):
# Decompose.
theta[yi, i, :] = np.linalg.solve(vec, solution[yi, :, i])
# End of eigenmode calculation and decomposition.
#========================================================
np.save('theta_array' + str(pi), theta)
np.save('val_array', val)
np.save('count_array', count)
np.save('count2_array', count2)
np.save('ratio_array', count)
# Don't return anything, arrays are saved instead.
return
def RSW_main(): # Forcing #plotting.forcingPlot_save(x_grid,y_grid,F3_nd[:,0:N],FORCE,BG,Fpos,N); #F1_nd, F2_nd, F3_nd = forcing.forcingInv(Ftilde1_nd,Ftilde2_nd,Ftilde3_nd,x_nd,y_nd,dx_nd,N); #F1_nd, F2_nd = forcing.F12_from_F3(F3_nd,f_nd,dx_nd,dy_nd,N,N); #plotting.forcingPlots(x_nd[0:N],y_nd,Ro*F1_nd,Ro*F2_nd,F3_nd,Ftilde1_nd,Ftilde2_nd,Ftilde3_nd,N); # sys.exit(); # Coefficients 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) # Solver if BC == 'NO-SLIP': solution = solver.NO_SLIP_SOLVER(a1,a2,a3,a4,f_nd,b4,c1,c2,c3,c4,Ro*Ftilde1_nd,Ro*Ftilde2_nd,Ftilde3_nd,N,N2) if BC == 'FREE-SLIP': #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) solution = solver.FREE_SLIP_SOLVER4(a1,a2,a3,a4,f_nd,b4,c1,c2,c3,c4,Ro*Ftilde1_nd,Ro*Ftilde2_nd,Ro*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); PP = np.zeros(N) EEF = 0. for iji in range(0,N): print(iji) ut = np.zeros((N,N),dtype=complex) vt = np.zeros((N,N),dtype=complex) ht = np.zeros((N,N),dtype=complex) ut[:,iji] = utilde_nd[:,iji] vt[:,iji] = vtilde_nd[:,iji] ht[:,iji] = etatilde_nd[:,iji] u, v, h = solver.SPEC_TO_PHYS(ut,vt,ht,T_nd,dx_nd,omega_nd,N); 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/energies 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] 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) PP += P_xav from scipy.ndimage.measurements import center_of_mass iii = center_of_mass(np.abs(P_xav))[0] com = y_nd[int(iii)] EEF, l = PV.EEF(P_xav,y_nd,com,int(iii),dy_nd,N) EEF_north = EEF[0]; EEF_south = EEF[1]; EEF_tmp = EEF_north - EEF_south; EEF += EEF_tmp plt.plot(PP) plt.show() print(EEF)