''' NAME NetCDF with Python PURPOSE To demonstrate how to read and write data with NetCDF files using a NetCDF file from the NCEP/NCAR Reanalysis. Plotting using Matplotlib and Basemap is also shown. PROGRAMMER(S) Chris Slocum REVISION HISTORY 20140320 -- Initial version created and posted online 20140722 -- Added basic error handling to ncdump Thanks to K.-Michael Aye for highlighting the issue REFERENCES netcdf4-python -- http://code.google.com/p/netcdf4-python/ NCEP/NCAR Reanalysis -- Kalnay et al. 1996 http://dx.doi.org/10.1175/1520-0477(1996)077<0437:TNYRP>2.0.CO;2 ''' import os import errno import sys import datetime as dt # Python standard library datetime module import numpy as np from netCDF4 import Dataset # http://code.google.com/p/netcdf4-python/ import matplotlib.pyplot as plt import matplotlib as mpl from mpl_toolkits.basemap import Basemap, addcyclic, shiftgrid from matplotlib.ticker import MultipleLocator from matplotlib import ticker import time def ncdump(nc_fid, verb=True): ''' ncdump outputs dimensions, variables and their attribute information. The information is similar to that of NCAR's ncdump utility. ncdump requires a valid instance of Dataset. Parameters ---------- nc_fid : netCDF4.Dataset A netCDF4 dateset object verb : Boolean whether or not nc_attrs, nc_dims, and nc_vars are printed Returns ------- nc_attrs : list A Python list of the NetCDF file global attributes nc_dims : list A Python list of the NetCDF file dimensions nc_vars : list A Python list of the NetCDF file variables ''' def print_ncattr(key): """ Prints the NetCDF file attributes for a given key Parameters ---------- key : unicode a valid netCDF4.Dataset.variables key """ try: print "\t\ttype:", repr(nc_fid.variables[key].dtype) for ncattr in nc_fid.variables[key].ncattrs(): print '\t\t%s:' % ncattr,\ repr(nc_fid.variables[key].getncattr(ncattr)) except KeyError: print "\t\tWARNING: %s does not contain variable attributes" % key # NetCDF global attributes nc_attrs = nc_fid.ncattrs() if verb: print "NetCDF Global Attributes:" for nc_attr in nc_attrs: print '\t%s:' % nc_attr, repr(nc_fid.getncattr(nc_attr)) nc_dims = [dim for dim in nc_fid.dimensions] # list of nc dimensions # Dimension shape information. if verb: print "NetCDF dimension information:" for dim in nc_dims: print "\tName:", dim print "\t\tsize:", len(nc_fid.dimensions[dim]) print_ncattr(dim) # Variable information. nc_vars = [var for var in nc_fid.variables] # list of nc variables if verb: print "NetCDF variable information:" for var in nc_vars: if var not in nc_dims: print '\tName:', var print "\t\tdimensions:", nc_fid.variables[var].dimensions print "\t\tsize:", nc_fid.variables[var].size print_ncattr(var) return nc_attrs, nc_dims, nc_vars def createFig(prefix, nc_fid, air, lons, lats, figType, depth_idx): tStart_f = time.time() path = prefix+'/'+figType+'/' if figType == "P" or figType == "P_anomaly": air = air/98000 fig = plt.figure(figsize=(20,20)) fig.subplots_adjust(left=None, bottom=None, right=None, top=None, \ wspace=None, hspace=None) axes = fig.add_subplot(1,1,1) tEnd_f = time.time() print "creatFig Cal1 cost %f sec" % (tEnd_f - tStart_f) tStart_f = time.time() ll_lat = 0 ll_lon = 100 ur_lat = 50 ur_lon = 150 m = Basemap(projection='merc', llcrnrlat=ll_lat,urcrnrlat=ur_lat,\ llcrnrlon=ll_lon,urcrnrlon=ur_lon, resolution='c') tEnd_f = time.time() print "creatFig Cal2 cost %f sec" % (tEnd_f - tStart_f) tStart_f = time.time() m.drawcoastlines() #m.drawmapboundary(fill_color='aqua') m.drawmapboundary() #air_cyclic, lons_cyclic = addcyclic(air[:, :], lons) #print air #print air_cyclic #print lons_cyclic #print lons tEnd_f = time.time() print "creatFig Cal3 cost %f sec" % (tEnd_f - tStart_f) tStart_f = time.time() # Create 2D lat/lon arrays for Basemap #lon2d, lat2d = np.meshgrid(lons_cyclic, lats) lon2d, lat2d = np.meshgrid(lons, lats) # Transforms lat/lon into plotting coordinates for projection #x, y = m(lons,lats) x, y = m(lon2d, lat2d) #air = air_cyclic # Plot of air temperature with 11 contour intervals #cs = m.contourf(x, y, air_cyclic, 21, cmap=plt.cm.Spectral_r) #norm = mpl.colors.Normalize(vmin=air.min(), vmax=((air.max()*0.8, air.max()) [air.max()*0.8 < air.min()])) #cs = m.contourf(x, y, air_cyclic, 21, cmap=plt.cm.jet, norm=norm) #cs = m.contourf(x, y, air, 21, cmap=plt.cm.jet, norm=norm) color_num = 101 ''' norm = mpl.colors.Normalize(vmin=air.min(), vmax=air.max()) cmap=plt.cm.seismic if figType == "T" or figType == "S": cs = m.contourf(x, y, air, color_num, cmap=plt.cm.jet) elif figType == "P_anomaly": norm = mpl.colors.Normalize(vmin=-0.5, vmax=0.5) cs = m.contourf(x, y, air, color_num, cmap=plt.cm.seismic, norm=norm) elif figType == "T_anomaly": norm = mpl.colors.Normalize(vmin=-4, vmax=4) cs = m.contourf(x, y, air, color_num, cmap=plt.cm.seismic, norm=norm) elif figType == "S_anomaly": norm = mpl.colors.Normalize(vmin=-1, vmax=1) cs = m.contourf(x, y, air, color_num, cmap=plt.cm.seismic, norm=norm) else : norm = mpl.colors.Normalize(vmin=air.max()-0.8*(air.max()-air.min()), vmax=0.8*(air.max()-air.min())+air.min()) cs = m.contourf(x, y, air, color_num, cmap=plt.cm.jet, norm=norm) ''' norm = mpl.colors.Normalize(vmin=air.min(), vmax=air.max()) cmap=plt.cm.jet if figType == "P_anomaly": norm = mpl.colors.Normalize(vmin=-0.5, vmax=0.5) cmap=plt.cm.seismic elif figType == "T_anomaly": norm = mpl.colors.Normalize(vmin=-4, vmax=4) cmap=plt.cm.seismic elif figType == "S_anomaly": norm = mpl.colors.Normalize(vmin=-1, vmax=1) cmap=plt.cm.seismic elif figType == "currents": norm = mpl.colors.Normalize(vmin=air.max()-0.8*(air.max()-air.min()), vmax=0.8*(air.max()-air.min())+air.min()) cmap=plt.cm.jet elif figType == "taux" or figType == "tauy": maxv = np.abs(air).max() norm = mpl.colors.Normalize(vmin=-maxv, vmax=maxv) cmap=plt.cm.jet cs = m.contourf(x, y, air, color_num, cmap=cmap, norm=norm) fig.tight_layout() fig.savefig(path+figType+'_'+depth_idx+'.png', transparent=True, format='png', pad_inches=0, bbox_inches='tight') fig = 0 if figType == "P": fig = plt.figure(figsize=(2, 4.0), frameon=False, facecolor=None) else : fig = plt.figure(figsize=(1.5, 4.0), frameon=False, facecolor=None) ax = fig.add_axes([0.2, 0.1, 0.2, 0.8], axisbg='g') ##ax.set_alpha(1) #cb = fig.colorbar(cs, cax=ax) #cb.set_label('Mean Dynamic Topography [m]', rotation=-90, color='k', labelpad=20) #if 'units' in nc_fid.variables[filename].ncattrs(): cb = mpl.colorbar.ColorbarBase(ax, cmap=cmap, norm=norm) cb.set_label('', color='k') if figType == "P" or figType == "P_anomaly": plt.xlabel("%s (%s)" % ("", "m")) #labels = np.arange(-0.5, 0.5, 0.1) #cb.set_ticklabels(labels) #tick_locator = ticker.MaxNLocator(nbins=5) #cb.locator = tick_locator elif figType == "currents": #cb.set_label("%s (%s)" % ("", nc_fid.variables["U"].units), rotation=0, color='k', labelpad=15) #cb.locator = MultipleLocator(30) #cb.set_ticks(np.arange(0, 150)) #cb.set_ticklabels(np.arange(0, 150, 30)) plt.xlabel("%s (%s)" % ("", nc_fid.variables["U"].units), x=2) elif figType == "heat_up": #cb.set_label("%s (%s)" % ("", ""), rotation=0, color='k', labelpad=15) plt.xlabel("%s (%s)" % ("", "W/m^2")) elif figType == "swr_down": #cb.set_label("%s (%s)" % ("", ""), rotation=0, color='k', labelpad=15) plt.xlabel("%s (%s)" % ("", "W/m^2")) elif figType == "taux": #cb.set_label("%s (%s)" % ("", ""), rotation=0, color='k', labelpad=15) plt.xlabel("%s (%s)" % ("", "dyn/cm^2")) elif figType == "tauy": #cb.set_label("%s (%s)" % ("", ""), rotation=0, color='k', labelpad=15) plt.xlabel("%s (%s)" % ("", "dyn/cm^2")) elif figType == "e_p": #cb.set_label("%s (%s)" % ("", ""), rotation=0, color='k', labelpad=15) plt.xlabel("%s (%s)" % ("", "g/cm2/s")) elif figType == "T_anomaly": plt.xlabel("%s (%s)" % ("", "degree C")) elif figType == "S_anomaly": plt.xlabel("%s (%s)" % ("", "ppt")) elif 'units' in nc_fid.variables[figType].ncattrs() : #cb.set_label("%s (%s)" % ("", nc_fid.variables[figType].units), rotation=0, color='k', labelpad=15, y=0, fontsize='small') plt.xlabel("%s (%s)" % ("", nc_fid.variables[figType].units)) tEnd_f = time.time() print "creatFig Cal4 cost %f sec" % (tEnd_f - tStart_f) tStart_f = time.time() fig.savefig(path+figType+'_'+depth_idx+'_cbar.png', transparent=False, format='png', facecolor='white') tEnd_f = time.time() print "createFig Save step cost %f sec" % (tEnd_f - tStart_f) #plt.show() plt.close("all") def smooth(INs, a, depth): end_lon = 406 end_lat = 488 IN = INs[int(depth)] AVG = ( \ IN[1:end_lat-1,1:end_lon-1] + \ IN[2:end_lat,1:end_lon-1] + \ IN[1:end_lat-1,2:end_lon] + \ IN[2:end_lat,2:end_lon] ) for i in range(1, 2, 1): a = \ 0.5 * a[1:end_lat-1,1:end_lon-1] * IN[1:end_lat-1,1:end_lon-1] + \ 0.5/AVG * (a[:end_lat-2,1:end_lon-1] * IN[:end_lat-2,1:end_lon-1] + \ a[2:end_lat,1:end_lon-1] * IN[2:end_lat,1:end_lon-1] + \ a[1:end_lat-1,:end_lon-2] * IN[1:end_lat-1,:end_lon-2] + \ a[1:end_lat-1,2:end_lon] * IN[1:end_lat-1,2:end_lon]) return a def standardizeSeq(u, v): vmax_u = np.max(abs(u)) vmax_v = np.max(abs(v)) if vmax_u > vmax_v: vmax = vmax_u; else: vmax = vmax_v; level = 25 level_diff = vmax/level end_lon = 406 end_lat = 488 u = u[1:end_lat-1, 1:end_lon-1] u = np.array([np.ceil(x/level_diff) for x in u]) v = u[1:end_lat-1, 1:end_lon-1] v = np.array([np.ceil(x/level_diff) for x in v]) return u, v def standardizeSeq(u): vmax_u = np.max(abs(u)) level = 25 level_diff = vmax_u/level end_lon = 406 end_lat = 488 #u = u[1:end_lat-1, 1:end_lon-1] u = np.array([np.ceil(x/level_diff) for x in u]) return u def drawVector(prefix, INs, U, V, lons, lats, figType, depth, zoom, feq, scale, arrow_width): tStart_f = time.time() print U.shape print V.shape U = smooth(INs, U, depth) V = smooth(INs, V, depth) U = standardizeSeq(U) V = standardizeSeq(V) #U, V = standardizeSeq(U, V) print U.shape print V.shape path = prefix+'/'+figType+'/' ll_lat = 0 ll_lon = 100 ur_lat = 50 ur_lon = 150 fig = plt.figure(figsize=(20,20)) # KML friendly image. If using basemap try: `fix_aspect=False`. #ax = fig.add_axes([0, 0, 1, 1]) fig.subplots_adjust(left=None, bottom=None, right=None, top=None, \ wspace=None, hspace=None) ax = fig.add_subplot(1,1,1) ax.set_xlim(ll_lon, ur_lon) ax.set_ylim(ll_lat, ur_lat) ax.set_xbound(ll_lon, ur_lon) ax.set_ybound(ll_lat, ur_lat) #u_cyclic, lons_cyclic = addcyclic(U[:, :], lons) #v_cyclic, lons_cyclic = addcyclic(V[:, :], lons) #lon2d, lat2d = np.meshgrid(lons_cyclic, lats) lon2d, lat2d = np.meshgrid(lons, lats) #u2d, v2d = np.meshgrid(u_cyclic, v_cyclic) #u, v = m(u, v) m = Basemap(projection='merc', llcrnrlat=ll_lat,urcrnrlat=ur_lat,\ llcrnrlon=ll_lon,urcrnrlon=ur_lon, resolution='i') #m.drawcoastlines() x, y = m(lon2d, lat2d) m.drawmapboundary(fill_color='aqua') m.quiver(x[::feq, ::feq], y[::feq, ::feq], U[::feq, ::feq], V[::feq, ::feq], scale=scale, headlength=5, headwidth=6, width=arrow_width, minshaft=1, minlength=1, scale_units='inches', units='inches', color='w') #m.quiver(x[::feq, ::feq], y[::feq, ::feq], U[::feq, ::feq], V[::feq, ::feq], scale=1, headlength=5, headwidth=6, width=0.02, minshaft=1, minlength=1, scale_units='dots', units='dots', color='b') #m.quiver(x[::feq, ::feq], y[::feq, ::feq], U[::feq, ::feq], V[::feq, ::feq], scale=180, headlength=vectorsize/2, headwidth=vectorsize*2/3, width=vectorsize, headaxislength=vectorsize/2, minlength=0, units='inches', minshaft=0.5, color='b') #U_cyclic, lons_cyclic = addcyclic(U[:, :], lons) #V_cyclic, lons_cyclic = addcyclic(V[:, :], lons) # Create 2D lat/lon arrays for Basemap #lon2d, lat2d = np.meshgrid(lons_cyclic, lats) plt.axis('off') fig.tight_layout() #plt.tight_layout() tEnd_f = time.time() print "drawVector Cal step cost %f sec" % (tEnd_f - tStart_f) tStart_f = time.time() fig.savefig(path+figType+'_'+depth+'_'+str(zoom)+'_v.png', transparent=True, format='png', pad_inches=0, bbox_inches='tight', dpi=250) tEnd_f = time.time() print "drawVector Save step cost %f sec" % (tEnd_f - tStart_f) plt.close("all") def checkDir(prefix): dir = ['P', 'S', 'T', 'currents', 'e_p', 'heat_up', 'swr_down', 'taux', 'tauy', 'P_anomaly', 'T_anomaly', 'S_anomaly'] for d in dir : try: os.stat(prefix) except: os.mkdir(prefix) try: os.stat(prefix + '/' +d) except: os.mkdir(prefix + '/' + d) ''' # Plot of global temperature on our random day fig = plt.figure(figsize=(1, 16)) ax1 = fig.add_axes([0.05, 0.80, 0.9, 0.15]) norm = mpl.colors.Normalize(vmin=np.nanmin(a=air_cyclic), vmax=np.nanmax(a=air_cyclic)) cbar = mpl.colorbar.ColorbarBase(ax1, orientation='vertical', cmap=plt.cm.Spectral_r, norm=norm) cbar.set_label("%s (%s)" % ("var_desc", nc_fid.variables['U'].units)) #plt.title("%s on %s" % (nc_fid.variables['air'].var_desc, cur_time)) fig.savefig(filename+'_cbar.png', transparent=True, format='png', pad_inches=0, bbox_inches='tight') ''' def initReader(): inputFile = '' inputFileFluxes = '' inputFileAnomaly = '' prefix = '' testFlag = False if "-t" not in sys.argv: inputFile = sys.argv[1] inputFileFluxes = sys.argv[2] inputFileAnomaly = sys.argv[3] try: os.stat(inputFile) except: print 'file ' + inputFile + ' is not exist' exit try: os.stat(inputFileFluxes) except: print 'file ' + inputFileFluxes + ' is not exist' exit try: os.stat(inputFileAnomaly) except: print 'file ' + inputFileAnomaly + ' is not exist' exit prefix = inputFile.replace('.nc', '') else : print 'test mode' if sys.argv[0] == "test.py": inputFile = 'member001_T1_00031.nc' inputFileFluxes = 'member001_fluxes_00031.nc' inputFileAnomaly = 'member001_T1_00031_anomaly.nc' testFlag = True prefix = inputFile.replace('.nc', '') else : exit return testFlag, prefix, inputFile, inputFileFluxes, inputFileAnomaly def readMainNc(fileName): try: os.stat(fileName) except: print 'file ' + fileName + ' is not exist' exit nc_fid = Dataset(fileName, 'r') # Dataset is the class behavior to open the file # and create an instance of the ncCDF4 class nc_attrs, nc_dims, nc_vars = ncdump(nc_fid) # Extract data from NetCDF file lats_ = nc_fid.variables['Latitude'][:][1:487] lons_ = nc_fid.variables['Longitude'][:][1:405] U = nc_fid.variables['U'][0][:] # shape is time, lat, lon as shown above V = nc_fid.variables['V'][0][:] # shape is time, lat, lon as shown above P = nc_fid.variables['P'][0][:] # shape is time, lat, lon as shown above T = nc_fid.variables['T'][0][:] # shape is time, lat, lon as shown above S = nc_fid.variables['S'][0][:] # shape is time, lat, lon as shown above mask = np.ma.getmask(P[0][1:487, 1:405]) return nc_fid, lats_, lons_, U, V, P, T, S, mask def readFlux(fileName, mask): try: os.stat(fileName) except: print 'file ' + fileName + ' is not exist' exit nc_fid = Dataset(fileName, 'r') # Dataset is the class behavior to open the file # and create an instance of the ncCDF4 class nc_attrs, nc_dims, nc_vars = ncdump(nc_fid) #mask = np.ma.getmask(P[1:487, 1:405]) heat_up = np.ma.MaskedArray(nc_fid.variables['heat_up'], mask=mask) swr_down = np.ma.MaskedArray(nc_fid.variables['swr_down'], mask=mask) taux = np.ma.MaskedArray(nc_fid.variables['taux'], mask=mask) tauy = np.ma.MaskedArray(nc_fid.variables['tauy'], mask=mask) e_p = np.ma.MaskedArray(nc_fid.variables['e_p'], mask=mask) return nc_fid, heat_up, swr_down, taux, tauy, e_p def readAnomaly(fileName, mask): try: os.stAnomalyat(fileName) except: print 'file ' + fileName + ' is not exist' exit nc_fid = Dataset(fileName, 'r') nc_attrs, nc_dims, nc_vars = ncdump(nc_fid) P = np.ma.MaskedArray(nc_fid.variables['P'][0], mask=mask) T = np.ma.MaskedArray(nc_fid.variables['T'][0], mask=mask) S = np.ma.MaskedArray(nc_fid.variables['S'][0], mask=mask) ''' for i in range(1, 25, 1): P = np.ma.MaskedArray(nc_fid.variables['P'][0][i][:], mask=mask) T = np.ma.MaskedArray(nc_fid.variables['T'][0][i][:], mask=mask) S = np.ma.MaskedArray(nc_fid.variables['S'][0][i][:], mask=mask) ''' return nc_fid, P, T, S def main(argv) : tStart = time.time() testFlag, prefix, inputFile, inputFileFluxes, inputFileAnomaly = initReader() checkDir(prefix) tEnd = time.time() print "Init step cost %f sec" % (tEnd - tStart) #exit nc = Dataset('IN.nc', 'r') # Dataset is the class behavior to open the file # and create an instance of the ncCDF4 class nc_attrs, nc_dims, nc_vars = ncdump(nc) INs = nc.variables['IN'][:] nc.close() tStart = time.time() nc_fid_main, lats_, lons_, U, V, P, T, S, mask = readMainNc(inputFile) print mask.shape tEnd = time.time() print "read file step cost %f sec" % (tEnd - tStart) tStart = time.time() if testFlag == True: d_size = 1 else : d_size = 25 feq = 10 scale = 25 for i in range(0, d_size, 1): Pi = P[i] # shape is time, lat, lon as shown above Ti = T[i] # shape is time, lat, lon as shown above Si = S[i] # shape is time, lat, lon as shown above Ui = U[i] # shape is time, lat, lon as shown above Vi = V[i] # shape is time, lat, lon as shown above currents = np.sqrt(Ui*Ui+Vi*Vi) print Pi.shape print lats_.shape print lons_.shape createFig(prefix, nc_fid_main, Pi[1:487, 1:405], lons_, lats_, 'P', str(i)) createFig(prefix, nc_fid_main, Ti[1:487, 1:405], lons_, lats_, 'T', str(i)) createFig(prefix, nc_fid_main, Si[1:487, 1:405], lons_, lats_, 'S', str(i)) createFig(prefix, nc_fid_main, currents[1:487, 1:405], lons_, lats_, 'currents', str(i)) #drawVector(U[1:487, 1:405], V[1:487, 1:405], lons_, lats_, 'currents', str(i)) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 4, 8, 25, 0.01) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 5, 8, 25, 0.01) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 6, 6, 25, 0.01) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 6, 6, 50, 0.01) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 7, 4, 50, 0.0075) drawVector(prefix, INs, Ui, Vi, lons_, lats_, 'currents', str(i), 8, 2, 100, 0.005) #drawVector(prefix, INs, U, V, lons_, lats_, 'currents', str(i), 9, 1, 100) tEnd = time.time() print "createFig total cost %f sec" % (tEnd - tStart) tStart = time.time() nc_fid_flux, heat_up, swr_down, taux, tauy, e_p = readFlux(inputFileFluxes, mask) tEnd = time.time() print "read file Flux total cost %f sec" % (tEnd - tStart) tStart = time.time() createFig(prefix, nc_fid_flux, heat_up, lons_, lats_, 'heat_up', '0') createFig(prefix, nc_fid_flux, swr_down, lons_, lats_, 'swr_down', '0') createFig(prefix, nc_fid_flux, taux, lons_, lats_, 'taux', '0') createFig(prefix, nc_fid_flux, tauy, lons_, lats_, 'tauy', '0') createFig(prefix, nc_fid_flux, e_p, lons_, lats_, 'e_p', '0') tEnd = time.time() print "createFig Flux total cost %f sec" % (tEnd - tStart) nc_fid_anomaly, P_anomaly, T_anomaly, S_anomaly = readAnomaly(inputFileAnomaly, np.ma.getmask(P)) feq = 10 scale = 25 for i in range(0, d_size, 1): createFig(prefix, nc_fid_anomaly, P_anomaly[i][1:487, 1:405], lons_, lats_, 'P_anomaly', str(i)) createFig(prefix, nc_fid_anomaly, T_anomaly[i][1:487, 1:405], lons_, lats_, 'T_anomaly', str(i)) createFig(prefix, nc_fid_anomaly, S_anomaly[i][1:487, 1:405], lons_, lats_, 'S_anomaly', str(i)) #plt.show() nc_fid_main.close() nc_fid_flux.close() nc_fid_anomaly.close() # Close original NetCDF file. if __name__ == "__main__": main(sys.argv)