def calc_ykt_daily_VQ_eddy_aqua(dir_in, dim):
"""
Calculates zonally and vertically integrated moisture transport
as a function of zonal wavenumber for all eddies (transient + stw).
This is a wrapper for the main transport calculation function
"""
import numpy as np
import xarray as xr
from ds21grl import const
from ds21grl.misc import tic, toc, get_aqua_timestamp
# Initialize
VQ = np.zeros(
(dim.years.size, dim.nday, dim.lat.size, dim.wavenumber.size))
# loop through all data chunks
for iyear in range(0, dim.years.size):
for ichunk in range(0, int(365 / dim.chunk)):
print('year: ' + str(dim.years[iyear]) + ', chunk: ' + str(ichunk))
tic()
# read data
timestamp = get_aqua_timestamp(dim.years[iyear],
ichunk,
branch_flag=0)
filename1 = dir_in + 'cam.h1.' + timestamp + '.nc'
filename2 = dir_in + 'cam.h3.' + timestamp + '.nc'
filename3 = dir_in + 'cam.h3.' + timestamp + '.nc'
ds1 = xr.open_dataset(filename1, decode_times=False)
ds2 = xr.open_dataset(filename2, decode_times=False)
ds3 = xr.open_dataset(filename3, decode_times=False)
V = ds1.V.values
Q = ds2.Q.values
PS = ds3.PS.values
dim.hyai = ds1.hyai.values
dim.hybi = ds1.hybi.values
dim.P0 = ds1.P0.values
dim.lev = ds1.lev.values
dim.ilev = ds1.ilev.values
ds1.close()
ds2.close()
ds3.close()
# define energy + mass flux terms
Q = const.Lv * Q
dp = compute_dp_ml(PS, dim)
dp = np.swapaxes(dp, 0, 1) # shift dims of dp same as V
V = V * dp / const.go
# calc transport
VQ_chunk = calc_ykt_transport(V, Q, dim)
# dump into long time series array and convert to PW
daysofyear = np.arange(ichunk * dim.chunk,
ichunk * dim.chunk + dim.chunk)
VQ[iyear, daysofyear, :, :] = VQ_chunk[:, :, :] / 10**15
toc()
return VQ
def calc_ykt_daily_VQ_eddy_erai(dir_in, dim):
"""
Calculates zonally and vertically integrated moisture transport
as a function of zonal wavenumber for all eddies (transient + stw).
This is a wrapper for the main transport calculation function
"""
import numpy as np
import xarray as xr
from ds21grl import const
from ds21grl.misc import tic, toc, leap_year_test, month_to_dayofyear
# Initialize
VQ = np.zeros(
(dim.years.size, dim.nday, dim.lat.size, dim.wavenumber.size))
for iyear in range(0, dim.years.size):
for imonth in range(0, dim.months.size):
tic()
# read data for each month
timestamp = str(dim.years[iyear]) + '-' + format(
dim.months[imonth], "02")
filename1 = dir_in + '/model_level/v/v_' + timestamp + '.nc'
filename2 = dir_in + '/model_level/q/q_' + timestamp + '.nc'
filename3 = dir_in + '/sfc/lnsp/lnsp_' + timestamp + '.nc'
ds1 = xr.open_dataset(filename1)
ds2 = xr.open_dataset(filename2)
ds3 = xr.open_dataset(filename3)
day = ds1.day.values
dim.P0 = ds1.P0.values
dim.hyai = ds1.hyai.values / dim.P0 # set units same as model
dim.hybi = ds1.hybi.values
dim.lev = ds1.lev.values
dim.ilev = ds1.ilev.values
V = ds1.v.values
Q = ds2.q.values
PS = np.exp(ds3.lnsp.values)
ds1.close()
ds2.close()
ds3.close()
# remove leap year day
if (leap_year_test(dim.years[iyear]) == 1) & (dim.months[imonth]
== 2):
V = V[:28, :, :, :]
Q = Q[:28, :, :, :]
PS = PS[:28, :, :]
day = day[:28]
# define energy + mass flux terms
Q = const.Lv * Q
dp = compute_dp_ml(PS, dim)
dp = np.swapaxes(dp, 0, 1) # shift dims of dp same as
V = V * dp / const.go
# calc transport
VQ_month = calc_ykt_transport(V, Q, dim)
# dump into long time series array and convert to PW
daysofyear = np.arange(month_to_dayofyear(imonth - 1),
month_to_dayofyear(imonth))
VQ[iyear, daysofyear, :, :] = VQ_month[:, :, :] / 10**15
toc()
return VQ
def calc_ykt_monthly_VQ_stw_aqua(dir_in, dim):
"""
Calculates zonally and vertically integrated moisture transport
as a function of zonal wavenumber for stationary waves.
This is a wrapper for the main transport calculation function
where the input data is the monthly climatological mass flux and
specific humidity.
"""
import numpy as np
import xarray as xr
from ds21grl import const
from ds21grl.misc import tic, toc, get_aqua_timestamp, daysinmonths
# initialize
V = np.zeros((dim.months.size, dim.lev.size, dim.lat.size, dim.lon.size))
Q = np.zeros((dim.months.size, dim.lev.size, dim.lat.size, dim.lon.size))
# calculate monthly mean mass flux and specific humidity
for iyear in range(0, dim.years.size):
for ichunk in range(0, int(365 / dim.chunk)):
print('year: ' + str(dim.years[iyear]) + ', chunk: ' + str(ichunk))
tic()
# read data
timestamp = get_aqua_timestamp(dim.years[iyear],
ichunk,
branch_flag=0)
filename1 = dir_in + 'cam.h1.' + timestamp + '.nc'
filename2 = dir_in + 'cam.h3.' + timestamp + '.nc'
filename3 = dir_in + 'cam.h3.' + timestamp + '.nc'
ds1 = xr.open_dataset(filename1)
ds2 = xr.open_dataset(filename2)
ds3 = xr.open_dataset(filename3)
V_chunk = ds1.V.values
Q_chunk = ds2.Q.values
PS_chunk = ds3.PS.values
dim.hyai = ds1.hyai.values
dim.hybi = ds1.hybi.values
dim.lev = ds1.lev.values
dim.ilev = ds1.ilev.values
month = ds1['time.month'].values - 1
ds1.close()
ds2.close()
ds3.close()
# define energy + mass flux terms
Q_chunk = const.Lv * Q_chunk
dp = compute_dp_ml(PS_chunk, dim)
dp = np.swapaxes(dp, 0, 1) # shift dims of dp same as V
V_chunk = V_chunk * dp / const.go
# calculate monthly mean stationary wave components
for t in range(0, dim.chunk):
V[month[t], :, :, :] = V[month[t], :, :, :] + V_chunk[
t, :, :, :] / dim.years.size / daysinmonths(month[t])
Q[month[t], :, :, :] = Q[month[t], :, :, :] + Q_chunk[
t, :, :, :] / dim.years.size / daysinmonths(month[t])
toc()
# calc transport using climatological monthly mean mass flux and latent energy
VQ = calc_ykt_transport(V, Q, dim)
VQ = VQ / 10**15 # convert to PW
return VQ
ds = xr.open_dataset(filename).sel(lat=ilat,wavenumber=waves,method='nearest').sum(dim='wavenumber')
data1 = ds['VQ'].values
ds.close()
# read zonal-mean data
filename = dir_in + 'yt_zm_daily_' + var + '_' + dim.timestamp + '.nc'
data2 = xr.open_dataarray(filename).values
# get anomalies relative to daily-mean seasonal cycle
data1 = get_anomaly_daily_seasonal_cycle(data1,dim)
data2 = get_anomaly_daily_seasonal_cycle(data2,dim)
# calculate correlation + significance for all lags at each latitude
data1 = np.squeeze(np.reshape(data1,(dim.nday*dim.years.size,1)))
data2 = np.reshape(data2,(dim.nday*dim.years.size,dim.lat.size))
corr = np.zeros([2*maxlag+1,dim.lat.size])
sig = np.zeros([2*maxlag+1,dim.lat.size])
lag = np.zeros([2*maxlag+1])
for j in range(0,dim.lat.size):
print('latitude: ' + str(dim.lat[j]))
tic()
[corr[:,j],sig[:,j],lag] = cross_correlate_ndim_sig(data1,data2[:,j],maxlag,nbs,sigthresh,ax=0)
toc()
# write 2 file
string1 = 'VQ_k' + str(waves[0]) + '-' + str(waves[-1]) + '_' + str(ilat) + 'N'
string2 = 'with_zm_' + var
filename = 'yt_corr_nbs' + str(nbs) + '_' + string1 + '_' + string2 + '_' + dim.timestamp + '.nc'
write_yt_daily(corr,sig,lag,filename,dir_out,dim,write2file)