CEEMDAN and 写nc文件（UVCDAT）

# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
import cdms2
import MV2
import numpy as np
from pyeemd import ceemdan

target_level = 925
f=cdms2.open('~/zy/data/jra55/monthly/uwnd.mon.mean.1958-2014.nc')
data_in = f('uwnd',time=('1980-01','2014-12'),latitude=(-20.,50),longitude=(30,130),lev=target_level)
data_shape = data_in.squeeze().shape
ntime = data_shape[0]
nlat = data_shape[1]
nlon = data_shape[2]
data = data_in.reshape(ntime,nlat*nlon)

ceemdan_filter=lambda x:ceemdan(x,ensemble_size=5)

imfs0 = map(ceemdan_filter,data)
nimfs = imfs0[0].shape[0]
imfs = MV2.masked_array(imfs0).reshape(ntime,nimfs,nlat,nlon)
imfs_mask=np.empty(imfs.shape,dtype=bool)
imfs_mask[:,:,:,:] = data_in.mask
imfs.mask = imfs_mask

def zero_period(x):
    omega=0
    x_len=len(x)
    for i in range(x_len-1):
        if set(x.mask)=={False}:
            if x[i]*x[i+1]<0:
                omega=omega+1
    if omega==0:
        return np.nan
    if omega!=0:
        return float(x_len)*2/omega
    
period=np.empty_like(imfs[0,:nimfs-1,:,:])*np.nan
for ilat in range(nlat):
    for ilon in range(nlon):
        for iimf in range(nimfs-1):
            period[iimf,ilat,ilon]=zero_period(imfs[:,iimf,ilat,ilon])
           
nimf = np.arange(0,nimfs)+1
imfs.id='imfs'
imfs.missing=1.e+20
imfs.setAxis(0,data_in.getAxis(0))
imfs.setAxis(1,cdms2.createAxis(nimf,bounds=None,id='nimf'))
imfs.setAxis(2,data_in.getAxis(2))
imfs.setAxis(3,data_in.getAxis(3))

period.id = 'period'
nimf_1 = np.arange(0,nimfs-1)+1
period.setAxis(0,cdms2.createAxis(nimf_1,bounds=None,id='nimf_1'))
period.setAxis(1,data_in.getAxis(2))
period.setAxis(2,data_in.getAxis(3))

fil_out = cdms2.open('~/zy/project/io_tp_mam/'+'uwnd'+'.1980-2014.'+str(target_level)+'.'+'ceemdan.nc','w')
fil_out.write(imfs)
fil_out.write(period)
fil_out.close()