def decompose_symasym(array_in):
"""
Function that decompose an array (supports array_in(lat,lon),
array_in(time,lat,lon) and array_in(time,level,lat,lon)) into its symmetric
and asymmetric parts about the equator.
The asymmetric part is stored in the Northern Hemisphere as
SymAsym[lat] = (array_in[lat]-array_in[-lat])/2,
while the symmetric part is stored in the Southern Hemisphere as
SymAsym[-lat] = (array_in[lat]+array_in[-lat])/2
:param array_in:
Input Aray
:type array_in: Numpy Array
:return: SymAsym
:rtype: Numpy Array
"""
#Check the dimensions of the array to find out if it has the structure of
#array_in(lat,lon), array_in(time,lat,lon) or array_in(time,level,lat,lon)
dim_array = array_in.shape
rank_array = len(dim_array)
if (rank_array>=5):
raise InputError("Decompose_SymAsym: currently supports up to 4D: rank = {}D"\
.format(rank_array))
#Copy array with same shape as array_in
#values will be overwritten except when the eq is
#present in the array
SymAsym = np.copy(array_in)
if (rank_array==1):
for i in range(N):
# Save symmetric part in the Southern Hemisphere
SymAsym[i] = 0.5*(array_in[nlat-1-i] + array_in[i])
# Save antisymmetric part in the Northern Hemisphere
SymAsym[nlat-1-i] = 0.5*(array_in[nlat-1-i] - array_in[i])
# Notice that the equator if present is untouched
if (rank_array==2):
SymAsym = sym_asym(array_in)
if (rank_array==3):
ntim = dim_array[0]
for t in range(ntim):
SymAsym[t,:,:] = sym_asym(array_in[t,:,:])
if (rank_array==4):
ntim = dim_array[0]
nlevel = dim_array[1]
for t in range(ntim):
for l in range(nlevel):
SymAsym[t,l,:,:] = sym_asym(array_in[t,l,:,:])
return SymAsym
def import_array(self,array_in,varname):
if (self.data_status=="Empty"):
self.data_status = "Numpy array"
self.file = "array"
self.varname = varname
self.data = array_in
print('An array has been imported')
else:
raise InputError("You have already loaded some data")
return None
def import_analysis(self,file):
if (self.data_status=="Empty"):
self.data_status = "An analysis file has been imported"
self.analysis_status = "Imported"
self.wk_spectra = np.load(file).item()
self.spd = self.wk_spectra['spd']
self.nDayWin = self.wk_spectra['nDayWin']
self.nDaySkip = self.wk_spectra['nDaySkip']
self.max_freq = self.wk_spectra['max_freq']
self.max_wn = self.wk_spectra['max_wn']
self.varname = self.wk_spectra['varname']
print('The analysis stored in {} has been imported.'.format(file))
else:
raise InputError("You have already loaded some data")
return None
def sampling_vars(array_in,spd,nDayWin,nDaySkip):
# Supports only array_in(time,lat,lon)
ntim,nlat,nlon = array_in.shape
if ((ntim%spd)!=0):
raise InputError("Input array must have complete days only ntim%spd = {}".format(ntim%spd))
# Number of days
nDayTot = ntim//spd
# Number of samples per temporal window
nSampWin = nDayWin*spd
# Number of samples to skip between window segments.
# Negative number means overlap
nSampSkip = nDaySkip*spd
return (ntim,nlat,nlon,nDayTot,nSampWin,nSampSkip)
def import_netcdf(self,file,varname,latname='latitude',latBound=15):
if (self.data_status=="Empty"):
self.data_status = "NetCDF File"
self.file = str(file)
self.varname = varname
self.latname = latname
self.latBound = latBound
f = nc.Dataset(file, 'r')
array_out = f.variables[varname][:]
latitudes = f.variables[latname][:]
lat_ind = np.where(np.logical_and(latitudes>=-latBound,latitudes<=latBound))[0]
array_out = array_out[:,lat_ind,:]
self.data = array_out
print('The file {} has been imported'.format(self.file))
else:
raise InputError("You have already loaded some data")
return None
def save_figs(self,name=None,ext='png'):
if (self.plots_available == "Yes"):
g_ext = ['png','pdf','eps']
if ext not in g_ext:
print('The file extension {} is not supported. Using png instead.'\
.format(str(ext)))
ext = 'png'
if name is None:
if hasattr(self, 'name'):
name = self.name
else:
name = str(self.varname)
dirname = name+'_plots'
if not os.path.isdir(dirname):
os.makedirs(dirname)
for key in self.plots:
self.plots[key].savefig(dirname+'/'+name+'_'+key+'.'+ext)
else:
raise InputError("You don't have plots to save")
return None
def wheeler_kiladis_spectra(self,spd,nDayWin,nDaySkip,max_freq=0.5,max_wn=20):
if (self.data_status=="Empty"):
raise InputError("You need to input some data first.")
else:
if (self.analysis_status=="Empty"):
array_in = self.data
self.spd = spd
self.nDayWin = nDayWin
self.nDaySkip = nDaySkip
self.max_freq = max_freq
self.max_wn = max_wn
ntim,nlat,nlon,nDayTot,nSampWin,nSampSkip = sampling_vars(array_in,spd,nDayWin,nDaySkip)
# Remove dominant signals
array_dt = remove_dominant_signals(array_in,spd,nDayWin,nDaySkip)
# Decompose in Symmetric and Antisymmetric components
array_as = decompose_symasym(array_dt)
wavefft,freqfft,peeAS = spectral_coefficients(array_as,spd,nDayWin,nDaySkip)
# Calculate the power spectrum
power = (abs(peeAS))**2 # power[window,freq,lat,lon]
psumanti,psumsym = separate_power(power,nlat,nSampWin,wavefft,freqfft)
psumb = derive_background(power,nlat,nSampWin,wavefft,freqfft)
# Cropping the output
indwave = np.where(np.logical_and(wavefft>=-max_wn,wavefft<=max_wn))[0]
indfreq = np.where(np.logical_and(freqfft>0,freqfft<=max_freq))[0]
wavefft = wavefft[indwave]
freqfft = freqfft[indfreq]
psumanti = psumanti[indfreq,:]
psumanti = psumanti[:,indwave]
psumsym = psumsym[indfreq,:]
psumsym = psumsym[:,indwave]
psumb = psumb[indfreq,:]
psumb = psumb[:,indwave]
# Log10 scaling
psumanti_log = np.log10(psumanti)
psumsym_log = np.log10(psumsym)
psumb_log = np.log10(psumb)
psumanti_r = psumanti/psumb
psumsym_r = psumsym/psumb
self.wk_spectra = {
'psumanti':psumanti,
'psumsym':psumsym,
'psumb':psumb,
'psumanti_log':psumanti_log,
'psumsym_log':psumsym_log,
'psumb_log':psumb_log,
'psumanti_r':psumanti_r,
'psumsym_r':psumsym_r,
'wavefft':wavefft,
'freqfft':freqfft,
'spd':spd,
'nDayWin':nDayWin,
'nDaySkip':nDaySkip,
'max_freq':max_freq,
'max_wn':max_wn,
'varname':self.varname
}
self.analysis_status = "Complete"
print("The Wheeler-Kiladis Analysis is complete.")
else:
raise InputError("You have already done the analysis.")
return None
def spectral_coefficients(array_in,spd,nDayWin,nDaySkip):
ntim,nlat,nlon,nDayTot,nSampWin,nSampSkip = sampling_vars(array_in,spd,nDayWin,nDaySkip)
# Test if there is enought time data for the analysis
if (ntim<nSampWin):
raise InputError("The available number of days is less than the sample window")
else:
# Count the number of available samples
nWindow = (ntim-nSampWin)//(nSampWin+nSampSkip)+1
# Test if longitude and time dimensions are even.
# The fft algorith gives the nyquist frequency one time for even and two times
# for odd dimensions
if (nSampWin%2==0): #if time is even
if (nlon%2==0): # and longitude is even also
peeAS = np.zeros((nWindow,nSampWin+1,nlat,nlon+1),dtype='c16')
else: # but longitude is odd
peeAS = np.zeros((nWindow,nSampWin+1,nlat,nlon),dtype='c16')
else: #if time is odd
if (nlon%2==0): # but longitude is even
peeAS = np.zeros((nWindow,nSampWin,nlat,nlon+1),dtype='c16')
else: # but longitude is odd also
peeAS = np.zeros((nWindow,nSampWin,nlat,nlon),dtype='c16')
# Create a tapering window(nSampWin,nlat,nlon) using the hanning function.
# For more information about the hanning function see:
# http://docs.scipy.org/doc/numpy/reference/generated/numpy.hanning.html
# Note: Hanning function is different from Hamming function.
tapering_window = np.repeat(np.hanning(nSampWin), nlat*nlon).reshape(nSampWin,nlat,nlon)
ntStrt = 0
ntLast = nSampWin
for nw in range(nWindow):
# Detrend temporal window
temp_window = signal.detrend(array_in[ntStrt:ntLast,:,:],axis=0,type='linear')
# Taper temporal window in the time dimension
temp_window = temp_window*tapering_window
# Apply fft to time and longitude
fourier_fft = np.fft.fft2(temp_window,axes=(0,2))
# normalize by # time samples
fourier_fft = fourier_fft/(nlon*nSampWin)
# fourier_fft(nSampWin,nlat,nlon) contains the
# complex space-time spectrum for each latitude
# Special reordering to resolve the Progressive and Retrogressive waves
# based on Hayashi (1971).
fourier_fft = np.fft.fftshift(fourier_fft, axes=(0,2))
if (nSampWin%2==0): #if time is even
if (nlon%2==0): # and longitude is even also
varspacetime = np.zeros((nSampWin+1,nlat,nlon+1),dtype='c16')
varspacetime[:nSampWin,:,:nlon] = fourier_fft
varspacetime[nSampWin,:,:] = varspacetime[0,:,:]
varspacetime[:,:,nlon] = varspacetime[:,:,0]
else: # but longitude is odd
varspacetime = np.zeros((nSampWin+1,nlat,nlon),dtype='c16')
varspacetime[:nSampWin,:,:] = fourier_fft
varspacetime[nSampWin,:,:] = varspacetime[0,:,:]
else: #if time is odd
if (nlon%2==0): # but longitude is even
varspacetime = np.zeros((nSampWin,nlat,nlon+1),dtype='c16')
varspacetime[:,:,:nlon] = fourier_fft
varspacetime[:,:,nlon] = varspacetime[:,:,0]
else: # but longitude is odd also
varspacetime = np.zeros((nSampWin,nlat,nlon),dtype='c16')
varspacetime[:,:,:] = fourier_fft
fourier_fft = varspacetime
# To correct that a positive freq in
# fourier corresponds to a negative wave freq
# i.e. Fourier -> e^i(kx+wt) != Wave -> e^i(kx-wt)
fourier_fft = fourier_fft[:,:,::-1]
# Save the Fourier Coefficients for each window
peeAS[nw,:,:,:] = fourier_fft
# Set index for next temporal window
ntStrt = ntLast+nSampSkip
ntLast = ntStrt+nSampWin
del fourier_fft, temp_window
wavefft = np.arange(-int(nlon/2),int(nlon/2)+1.,1.)
freqfft = np.arange(-1.*int(nDayWin*spd/2),1.*int(nDayWin*spd/2)+1.,1)/(nDayWin)# Calculate the power spectrum
return (wavefft,freqfft,peeAS)