source: dynamico_lmdz/guided/Data/vertical_interpolation_camelot.py @ 4159

from __future__ import division
import argparse
import numpy as np
from netCDF4 import Dataset
import matplotlib.pyplot as plt 
from mpl_toolkits.basemap import Basemap
import matplotlib.mlab as mlab
import matplotlib.colors
import sys 
import os
import math
from scipy.interpolate import interp1d
from matplotlib.ticker import AutoMinorLocator
import netCDF4 as cdf

dw1 = Dataset('2001.nc','r')
dtfull = dw1.variables['t'][:,:,:,:]
dpfulloned = dw1.variables['level'][:]
dpfull = np.zeros(((len(dpfulloned),dtfull.shape[2],dtfull.shape[3])))
for l in xrange(0, len(dpfulloned)):
    dpfull[l,:,:]= dpfulloned[l]
        
print(dpfull)
dtfull = dw1.variables['t'][:,:,:,:]
#dPSfull = dw1.variables['PS'][:,:,:]
time_counterfull = dw1.variables['time'][:]
lat = dw1.variables['lat'][:]
lon = dw1.variables['lon'][:]
levels = [100., 200., 300., 500., 700., 1000., 2000., 3000., 5000., 7000., 10000., 12500., 15000., 17500., 20000., 22500.,25000., 30000., 35000., 40000., 
    45000., 50000., 55000., 60000., 65000., 70000., 75000., 77500., 80000., 82500., 85000., 87500., 90000., 92500., 95000., 97500., 100000.]

levels = levels[::-1]
print(levels)

dim0=dtfull.shape[0]
dim1=dtfull.shape[1]
dim2=dtfull.shape[2]
dim3=dtfull.shape[3]
out3 = np.zeros((((dim0,len(levels),dim2,dim3))))

def vertical_int2(levels, pfull, dataset):
    nl, ni, nj = len(levels), dataset.shape[1], dataset.shape[2]
    out = np.zeros((nl,ni,nj))
    for i in range(dataset.shape[1]):
        for j in range(dataset.shape[2]):
            f = interp1d(pfull[:,i,j],dataset[:,i,j],kind='linear',fill_value="extrapolate")
            out[:,i,j] = f(levels)
    return out

pfull = dpfull

for t in range(dim1):
    print("time step=",t)
    tfull = dtfull[t,:,:,:]
    out2 = vertical_int2(levels, pfull, tfull)
    out3[t,:,:,:]=out2[:,:,:]

del tfull,dtfull

dufull = dw1.variables['u'][:,:,:,:]
out4 = np.zeros((((dim0,len(levels),dim2,dim3))))
for t in range(dim1):
    print("time step=",t)
    ufull = dufull[t,:,:,:]
    out2 = vertical_int2(levels, pfull, ufull)
    out4[t,:,:,:]=out2[:,:,:]
del ufull, dufull 

dvfull = dw1.variables['v'][:,:,:,:]
out5 = np.zeros((((dim0,len(levels),dim2,dim3))))
for t in range(dim1):
    print("time step=",t)
    vfull = dvfull[t,:,:,:]
    out2 = vertical_int2(levels, pfull, vfull)
    out5[t,:,:,:]=out2[:,:,:]
del vfull,dvfull

print("-----------------------------writing NetCDF--------------------------------")
print(out3)
#sys.exit(0)
try:
    f = cdf.Dataset('interpolated_data.nc', 'w', format='NETCDF4')
except:
    print("Error occurred while opening new netCDF file, Error: ", sys.exc_info()[0])
levels2 = [levels[i] / 100. for i in range(len(levels))]
#levels2 = [int(i) for i in levels2]
#print(levels2)
f.createDimension('lon', len(lon))
f.createDimension('lat', len(lat))
f.createDimension('lev', len(levels2))
f.createDimension('time_counter', len(time_counterfull))
vlon = f.createVariable('lon', 'f4', 'lon')
vlat = f.createVariable('lat', 'f4', 'lat')  
vlev = f.createVariable('lev', 'f4', 'lev')
#vlev = f.createVariable('lev', 'i4', 'lev')
vtime = f.createVariable('time_counter', 'f8', 'time_counter')
vt = f.createVariable('t', 'f4',('time_counter', 'lev', 'lat', 'lon'))
vu = f.createVariable('u', 'f4',('time_counter', 'lev', 'lat', 'lon'))
vv = f.createVariable('v', 'f4',('time_counter', 'lev', 'lat', 'lon'))
#vPS = f.createVariable('PS', 'f4',('time_counter', 'lat', 'lon'))
vlon[:] = lon
vlat[:] = lat
vlev[:] = levels2
vtime[:] = time_counterfull
vt[:] = out3
vu[:] = out4
vv[:] = out5
#vPS[:] = dPSfull
f.close()
print("-----------------------done--------------------")