Skip to content

Commit

Permalink
Update to release 0.2.1. Some bugs fixed.
Browse files Browse the repository at this point in the history
  • Loading branch information
csyhuang committed Feb 26, 2018
1 parent d79ee7a commit af8c013
Show file tree
Hide file tree
Showing 4 changed files with 95 additions and 62 deletions.
127 changes: 74 additions & 53 deletions hn2016_falwa/beta_version.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@ def ufunclike(xs):


def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
zm_PT, Input_B0, Input_B1, use_real_Data=True):
zm_PT, Input_B0, Input_B1, use_real_Data=True,
plot_all_ref_quan=False):
"""
Compute equivalent latitude and wave activity on a barotropic sphere.
Expand All @@ -48,6 +49,9 @@ def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
Zonal-mean surface wave activity for the second lowest layer (k=1). Part of the lower-boundary condition.
use_real_Data : boolean
Whether to use input data to compute the reference states. By detault True. If false, randomly generated arrays will be used.
plot_all_ref_quan : boolean
Whether to plot the solved reference states using matplotlib library. By default False. For debugging.
Returns
Expand Down Expand Up @@ -82,6 +86,8 @@ def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
from copy import copy
import numpy as np
import itertools
if plot_all_ref_quan:
import matplotlib.pyplot as plt

# === Parameters (should be input externally. To be modified) ===
dz = 1000. # vertical z spacing (m)
Expand All @@ -94,7 +100,7 @@ def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,

# === These changes with input variables' dimensions ===
nlat = FAWA_cos.shape[-1]
jmax1 = nlat
jmax1 = nlat//4
dm = 1./float(jmax1+1) # gaussian latitude spacing
gl = np.array([(j+1)*dm for j in range(jmax1)]) # This is sin / mu
gl_2 = np.array([j*dm for j in range(jmax1+2)]) # This is sin / mu
Expand Down Expand Up @@ -142,15 +148,21 @@ def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
ephalf[:,:] = f_ep_toGaussian(alat[:])

# --- Interpolation of Delta_PT ---
f_DT_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Delta_PT[:], kind='linear') ) # This is txt in Noboru's code
#f_DT_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Delta_PT[:], kind='linear') ) # This is txt in Noboru's code
f_DT_toGaussian = interpolate.interp1d(ylat[:],Delta_PT[:],
kind='linear',fill_value='extrapolate')
Delta_PT1[:] = f_DT_toGaussian(alat_2[:])

# --- Interpolation of Input_B0_1 ---
f_B0_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B0[:], kind='linear') ) # This is txt in Noboru's code
#f_B0_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B0[:], kind='linear') ) # This is txt in Noboru's code
f_B0_toGaussian = interpolate.interp1d(ylat[:],Input_B0[:],
kind='linear',fill_value='extrapolate') # This is txt in Noboru's code
Input_B0_1[:] = f_B0_toGaussian(alat_2[:])

# --- Interpolation of Input_B1_1 ---
f_B1_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B1[:], kind='linear') ) # This is txt in Noboru's code
# f_B1_toGaussian = extrap1d( interpolate.interp1d(ylat[:],Input_B1[:], kind='linear') ) # This is txt in Noboru's code
f_B1_toGaussian = interpolate.interp1d(ylat[:],Input_B1[:],
kind='linear',fill_value='extrapolate') # This is txt in Noboru's code
Input_B1_1[:] = f_B1_toGaussian(alat_2[:])

else:
Expand Down Expand Up @@ -351,65 +363,74 @@ def solve_uref_both_bc(tstamp, zmum, FAWA_cos, ylat, ephalf2, Delta_PT,
T_MassCorr[j,k] = T_MassCorr[j-2,k] - (2.*om*gl[j-1])*aa*hh*dmdz / (r0 * cosl[j-1]) * (u_MassCorr[j-1,k+1]-u_MassCorr[j-1,k-1])
# ---- First do interpolation (gl is regular grid) ----
# f_Todd = interpolate.interp1d(gl[:-1:2],T_MassCorr[1:-1:2,k]) #[jmax x kmax]
f_Todd = interpolate.interp1d(gl_2[::2],T_MassCorr[::2,k]) #[jmax x kmax]
f_Todd_ex = extrap1d(f_Todd)
T_MassCorr[:,k] = f_Todd_ex(gl_2[:]) # Get all the points interpolated
#f_Todd = interpolate.interp1d(gl_2[::2],T_MassCorr[::2,k]) #[jmax x kmax]
#f_Todd_ex = extrap1d(f_Todd)

f_Todd = interpolate.interp1d(gl_2[::2],T_MassCorr[::2,k],
kind='linear',fill_value='extrapolate')
T_MassCorr[:,k] = f_Todd(gl_2[:])



# T_MassCorr[:,k] = f_Todd_ex(gl_2[:]) # Get all the points interpolated

# ---- Then do domain average ----
T_MC_mean = np.mean(T_MassCorr[:,k])
T_MassCorr[:,k] -= T_MC_mean

# --- First, interpolate MassCorr back to regular grid first ---
f_u_MassCorr = interpolate.interp1d(alat_2,u_MassCorr,axis=0, kind='linear') #[jmax x kmax]
u_MassCorr_regular[:,-nlat/2:] = f_u_MassCorr(ylat[-nlat/2:]).T
u_MassCorr_regular[:,-nlat//2:] = f_u_MassCorr(ylat[-nlat//2:]).T
f_T_MassCorr = interpolate.interp1d(alat_2,T_MassCorr,axis=0, kind='linear') #[jmax x kmax]
T_MassCorr_regular[:,-nlat/2:] = f_T_MassCorr(ylat[-nlat/2:]).T
T_MassCorr_regular[:,-nlat//2:] = f_T_MassCorr(ylat[-nlat//2:]).T

u_Ref = zmum[:,-nlat/2:] - u_MassCorr_regular[:,-nlat/2:]
T_ref = zm_PT[:,-nlat/2:] * np.exp(-np.arange(kmax)/7. * rkappa)[:,np.newaxis] - T_MassCorr_regular[:,-nlat/2:]
u_Ref = zmum[:,-nlat//2:] - u_MassCorr_regular[:,-nlat//2:]
T_ref = zm_PT[:,-nlat//2:] * np.exp(-np.arange(kmax)/7. * rkappa)[:,np.newaxis] - T_MassCorr_regular[:,-nlat//2:]

u_Ref_regular[:,-nlat/2:] = u_Ref
T_Ref_regular[:,-nlat/2:] = T_ref
u_Ref_regular[:,-nlat//2:] = u_Ref
T_Ref_regular[:,-nlat//2:] = T_ref
#
# plot_all_ref_quan = False
# if plot_all_ref_quan:
# # --- Colorbar scale ---
# contour_int = np.arange(-120,145,5)
# dT_contour_int = np.arange(-120,81,5)
# T_contour_int = np.arange(160,321,5)
# # --- Start plotting figure ---
# fig = plt.subplots(figsize=(12,12))
# plt.subplot(221)
# plt.contourf(ylat[-nlat/2:],height[:-2],u_MassCorr_regular[:-2,-nlat/2:],contour_int)
# plt.colorbar()
# c1=plt.contour(ylat[-nlat/2:],height[:-2],u_MassCorr_regular[:-2,-nlat/2:],contour_int[::2],colors='k')
# plt.clabel(c1,c1.levels,inline=True, fmt='%d', fontsize=10)
# plt.title('$\Delta$ u '+tstamp)
# plt.ylabel('height (km)')
# plt.subplot(222)
# plt.contourf(ylat[-nlat/2:],height[:-2],u_Ref[:-2,:],contour_int)
# plt.colorbar()
# c2=plt.contour(ylat[-nlat/2:],height[:-2],u_Ref[:-2,:],contour_int[::2],colors='k')
# plt.clabel(c2,c2.levels,inline=True, fmt='%d', fontsize=10)
# plt.title('$u_{REF}$ ('+BCstring+' BC)')
# plt.subplot(223)
# plt.contourf(ylat[-nlat/2:],height[:-2],T_MassCorr_regular[:-2,-nlat/2:],dT_contour_int)
# plt.colorbar()
# c3=plt.contour(ylat[-nlat/2:],height[:-2],T_MassCorr_regular[:-2,-nlat/2:],dT_contour_int,colors='k')
# plt.clabel(c3,c3.levels,inline=True, fmt='%d', fontsize=10)
# plt.title('$\Delta$ T')
# plt.ylabel('height (km)')
# plt.subplot(224)
# plt.contourf(ylat[-nlat/2:],height[:-2],T_ref[:-2,:],T_contour_int)
# plt.colorbar()
# c4=plt.contour(ylat[-nlat/2:],height[:-2],T_ref[:-2,:],T_contour_int[::2],colors='k')
# plt.clabel(c4,c4.levels,inline=True, fmt='%d', fontsize=10)
# plt.title('$T_{REF}$')
# plt.ylabel('height (km)')
# plt.tight_layout()
# plt.show()
# #plt.savefig('/home/csyhuang/Dropbox/Research-code/Sep12_test3_'+BCstring+'_'+tstamp+'.png')
# plt.close()
#plot_all_ref_quan = False
if plot_all_ref_quan:
# --- height coordinate ---
height = np.array([i for i in range(kmax)]) # in [km]
# --- Colorbar scale ---
contour_int = np.arange(-120,145,5)
dT_contour_int = np.arange(-120,81,5)
T_contour_int = np.arange(160,321,5)
# --- Start plotting figure ---
fig = plt.subplots(figsize=(12,12))
plt.subplot(221)
plt.contourf(ylat[-nlat//2:],height[:-2],u_MassCorr_regular[:-2,-nlat//2:],contour_int)
plt.colorbar()
c1=plt.contour(ylat[-nlat//2:],height[:-2],u_MassCorr_regular[:-2,-nlat//2:],contour_int[::2],colors='k')
plt.clabel(c1,c1.levels,inline=True, fmt='%d', fontsize=10)
plt.title('$\Delta$ u '+tstamp)
plt.ylabel('height (km)')
plt.subplot(222)
plt.contourf(ylat[-nlat//2:],height[:-2],u_Ref[:-2,:],contour_int)
plt.colorbar()
c2=plt.contour(ylat[-nlat//2:],height[:-2],u_Ref[:-2,:],contour_int[::2],colors='k')
plt.clabel(c2,c2.levels,inline=True, fmt='%d', fontsize=10)
plt.title('$u_{REF}$ ('+BCstring+' BC)')
plt.subplot(223)
plt.contourf(ylat[-nlat//2:],height[:-2],T_MassCorr_regular[:-2,-nlat//2:],dT_contour_int)
plt.colorbar()
c3=plt.contour(ylat[-nlat//2:],height[:-2],T_MassCorr_regular[:-2,-nlat//2:],dT_contour_int,colors='k')
plt.clabel(c3,c3.levels,inline=True, fmt='%d', fontsize=10)
plt.title('$\Delta$ T')
plt.ylabel('height (km)')
plt.subplot(224)
plt.contourf(ylat[-nlat//2:],height[:-2],T_ref[:-2,:],T_contour_int)
plt.colorbar()
c4=plt.contour(ylat[-nlat//2:],height[:-2],T_ref[:-2,:],T_contour_int[::2],colors='k')
plt.clabel(c4,c4.levels,inline=True, fmt='%d', fontsize=10)
plt.title('$T_{REF}$')
plt.ylabel('height (km)')
plt.tight_layout()
plt.show()
#plt.savefig('/home/csyhuang/Dropbox/Research-code/Sep12_test3_'+BCstring+'_'+tstamp+'.png')
plt.close()

# This is for only outputing Delta_u and Uref for no-slip and adiabatic boundary conditions.
return u_MassCorr_regular_noslip,u_Ref_regular_noslip,T_MassCorr_regular_noslip,T_Ref_regular_noslip, u_MassCorr_regular_adiab,u_Ref_regular_adiab,T_MassCorr_regular_adiab,T_Ref_regular_adiab
Expand Down
13 changes: 8 additions & 5 deletions hn2016_falwa/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,9 @@ def static_stability(height,area,theta,s_et=None,n_et=None):

nlat = theta.shape[1]
if s_et==None:
s_et = nlat/2
s_et = nlat//2
if n_et==None:
n_et = nlat/2
n_et = nlat//2

stat_n = np.zeros(theta.shape[0])
stat_s = np.zeros(theta.shape[0])
Expand Down Expand Up @@ -85,7 +85,7 @@ def static_stability(height,area,theta,s_et=None,n_et=None):


def compute_qgpv_givenvort(omega,nlat,nlon,kmax,unih,ylat,avort,potential_temp,
t0_cn,t0_cs,stat_cn,stat_cs,scale_height=7000.):
t0_cn,t0_cs,stat_cn,stat_cs,nlat_s=None,scale_height=7000.):
"""
The function "compute_qgpv_givenvort" computes the quasi-geostrophic potential
vorticity based on the absolute vorticity, potential temperature and static
Expand Down Expand Up @@ -141,6 +141,9 @@ def compute_qgpv_givenvort(omega,nlat,nlon,kmax,unih,ylat,avort,potential_temp,
"""

if nlat_s==None:
nlat_s=nlat//2

clat = np.cos(ylat*pi/180.)
clat = np.abs(clat) # Just to avoid the negative value at poles

Expand All @@ -155,8 +158,8 @@ def compute_qgpv_givenvort(omega,nlat,nlon,kmax,unih,ylat,avort,potential_temp,
zdiv = np.empty_like(potential_temp)
dzdiv = np.empty_like(potential_temp)
for kk in range(kmax): # This is more efficient
zdiv[kk,:60,:] = exp(-unih[kk]/scale_height)*(potential_temp[kk,:60,:]-t0_cs[kk])/stat_cs[kk]
zdiv[kk,60:,:] = exp(-unih[kk]/scale_height)*(potential_temp[kk,60:,:]-t0_cn[kk])/stat_cn[kk]
zdiv[kk,:nlat_s,:] = exp(-unih[kk]/scale_height)*(potential_temp[kk,:nlat_s,:]-t0_cs[kk])/stat_cs[kk]
zdiv[kk,-nlat_s:,:] = exp(-unih[kk]/scale_height)*(potential_temp[kk,-nlat_s:,:]-t0_cn[kk])/stat_cn[kk]

dzdiv[1:kmax-1,:,:] = np.exp(unih[1:kmax-1,np.newaxis,np.newaxis]/scale_height)* \
(zdiv[2:kmax,:,:]-zdiv[0:kmax-2,:,:]) \
Expand Down
10 changes: 7 additions & 3 deletions hn2016_falwa/wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ def barotropic_eqlat_lwa(ylat,vort,area,dmu,n_points,planet_radius = 6.378e+6):
2-d numpy array of local wave activity values;
dimension = [nlat_s x nlon]
"""

from hn2016_falwa import basis

nlat = vort.shape[0]
Expand Down Expand Up @@ -508,7 +509,7 @@ def qgpv_input_qref_to_compute_lwa(ylat, qref, vort, area, dmu, nlat_s=None,



def theta_lwa(ylat,theta,area,dmu,nlat_s=None,planet_radius=6.378e+6):
def theta_lwa(ylat,theta,area,dmu,nlat_s=None,n_points=None,planet_radius=6.378e+6):
"""
Compute the surface wave activity *B* based on surface potential temperature.
See Nakamura and Solomon (2010a) for details.
Expand Down Expand Up @@ -542,13 +543,16 @@ def theta_lwa(ylat,theta,area,dmu,nlat_s=None,planet_radius=6.378e+6):
nlon = theta.shape[1]
if nlat_s == None:
nlat_s = nlat/2
if n_points == None:
n_points = nlat_s

qref = np.zeros(nlat)
lwa_result = np.zeros((nlat,nlon))

# --- southern Hemisphere ---
qref[:nlat_s], brac = basis.eqvlat(ylat[:nlat_s], theta[:nlat_s, :], area[:nlat_s, :],
n_points, planet_radius=planet_radius)
qref[:nlat_s], brac = basis.eqvlat(ylat[:nlat_s], theta[:nlat_s, :],
area[:nlat_s, :], n_points,
planet_radius=planet_radius)
#qref1 = eqvlat(ylat[:nlat_s],theta[:nlat_s,:],area[:nlat_s,:],nlat_s,planet_radius=6.378e+6)
# qref[:nlat_s] = qref1
# lwa_south = lwa(nlon,nlat_s,theta[:nlat_s,:],qref1,dmu[:nlat_s])
Expand Down
7 changes: 6 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from setuptools import setup, find_packages

VERSION='0.2.0'
VERSION='0.2.1'
DISTNAME='hn2016_falwa'
URL='https://github.com/csyhuang/hn2016_falwa' # how can we make download_url automatically get the right version?
DOWNLOAD_URL='https://github.com/csyhuang/hn2016_falwa/'
Expand All @@ -26,6 +26,11 @@
v0.2.0:
- Functions are structred in 4 different modules: basis, wrapper, utilities and beta_version. See documentation on Github for details.
v0.2.1 (bug fixing):
- hn2016_falwa/beta_version.py: In *solve_uref_both_bc*, a plotting option is added.
- hn2016_falwa/wrapper.py: In all functions, when n_points are not specified, it is taken to be nlat_s (input). Also fixed a bug of missing argument n_points in *theta_lwa*.
- hn2016_falwa/utilities.py: In *static_stability*, make s_et and n_et an integer if they are not input. In *compute_qgpv_givenvort*, remove the bug that nlat_s being hard-coded by mistake.
Links:
-----
Expand Down

0 comments on commit af8c013

Please sign in to comment.