read_and_map_merra_aerosol.py

#!/usr/bin/python3
'''
Module: read_and_map_merra_aerosol.py
==========================================================================================
Disclaimer: The code is for demonstration purposes only. Users are responsible to check for accuracy and revise to fit their objective.

Author: Justin Roberts-Pierel, 2015 
Organization: NASA ARSET
Purpose: To extract various SDS from a MERRA NetCDF4 file and create a map of the resulting data

See the README associated with this module for more information.
==========================================================================================
'''

#import necessary modules
import numpy as np
import netcdf4 as nc4
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap


#This finds the user's current path so that all hdf4 files can be found
try:
	fileList=open('fileList.txt','r')
except:
	print('Did not find a text file containing file names (perhaps name does not match)')
   

#loops through all files listed in the text file
for FILE_NAME in fileList:
	FILE_NAME=FILE_NAME.strip()
	user_input=input('\nWould you like to process\n' + FILE_NAME + '\n\n(Y/N)')
	if(user_input == 'N' or user_input == 'n'):
		continue
	else:
		#read in the data 
		merraData = nc4.Dataset(FILE_NAME, 'r')
		variables=set(merraData.variables)
		desiredVariables=set({'AOD','RH','ps','PBLH','TA','US','VS'})
		desiredVariables=set([x.lower() for x in desiredVariables])
		var1=variables.intersection(desiredVariables)
		desiredVariables=set([x.upper() for x in desiredVariables])
		var2=variables.intersection(desiredVariables)
		fileVars=list(var1.union(var2))
		if len(fileVars)==0:
			print('This file contains none of the selected SDS. Skipping...')
			continue
		print('Please select an SDS from the current file to map: ') 
		[print('(' + str(fileVars.index(x)) + ')',x) for x in fileVars]
		user_var=input()
		SDS_NAME=fileVars[int(user_var)]
		try:
			#gets the fill value 
			fv=float((merraData.variables[SDS_NAME])._FillValue)
		except:
			print('There is an issue with your MERRA file (might be the wrong MERRA file type). Skipping...')
			continue
		#gets the SDS data from merraData (multidimensional)
		data = merraData.variables[SDS_NAME][:]
		if len(data.shape) ==4:
			data=data[0,:,:,:]
			level=data.shape[0]-1
		elif len(data.shape) == 3:
			level=data.shape[0]-1
		#replace fill values with nan 
		data[data==fv]=np.nan

		#asks the user what type of map they would like 
		user_input=float(input('For a world map, enter (1). To enter a latitude and longitude, press (2)'))
		while user_input != 1 and user_input != 2:
			user_input=float(input('Please enter a valid choice'))
		if user_input == 1: #then set up a world map
			# set up cylindrical map
			user_min_lat=-90
			user_max_lat=90
			user_min_lon=-180
			user_max_lon=180
			xMax=180.0
			yMax=90.1
		else: #get more info for a specific location map 
			user_min_lat=float(input('Enter a starting (min) latitude (Deg N): '))
			user_max_lat=float(input('Enter an ending (max) latitude (Deg N): '))
			while user_max_lat<=user_min_lat:
				user_min_lat=float(input('Your ending latitude is smaller than your starting latitude. Please enter a new starting (min) latitude (Deg N): '))
				user_max_lat=float(input('Enter an ending (max) latitude (Deg N): '))
			user_min_lon=float(input('Enter a starting (min) longitude (Deg E): '))
			user_max_lon=float(input('Enter an ending (max) longitude (Deg E): '))
			while user_max_lon<=user_min_lon:
				user_min_lon=float(input('Your ending longitude is smaller than your starting longitude. Please enter a new starting (min) longitude (Deg N): '))
				user_max_lon=float(input('Enter an ending (max) longitude (Deg N): '))

			#extract lat and lon info 
			tempLat=merraData.variables['lat'][:]
			min_lat=np.unravel_index(abs(tempLat-user_min_lat).argmin(),abs(tempLat-user_min_lat).shape)
			max_lat=np.unravel_index(abs(tempLat-user_max_lat).argmin(),abs(tempLat-user_max_lat).shape)
			tempLon=merraData.variables['lon'][:]
			min_lon=np.unravel_index(abs(tempLon-user_min_lon).argmin(),abs(tempLon-user_min_lon).shape)
			max_lon=np.unravel_index(abs(tempLon-user_max_lon).argmin(),abs(tempLon-user_max_lon).shape)
			data = merraData.variables[SDS_NAME][0,:,min_lat[0]-1:max_lat[0],min_lon[0]-1:max_lon[0]]
			xMax=user_max_lon+.1
			yMax=user_max_lat+.1
		
		#start plotting 
		m = Basemap(
				projection='cyl',
				llcrnrlat=user_min_lat, urcrnrlat=user_max_lat,
				llcrnrlon=user_min_lon, urcrnrlon=user_max_lon,
				resolution='c'
				)
		m.drawcoastlines(linewidth=0.5)
		m.drawmapboundary()
		# plot pcolormesh
		if len(data.shape)>2:
			data=data[level,:,:]
		data[data==np.nanmin(data)]=np.nan
		data = np.ma.masked_array(data, np.isnan(data))
		X = np.arange(user_min_lon, xMax, .625)
		Y = np.arange(user_min_lat, yMax, .5) # 90 is the last element
		my_cmap = plt.cm.get_cmap('Oranges')
		my_cmap.set_under('w')
		cp=m.pcolormesh(X, Y, data, latlon=True, vmin=np.nanmin(data), vmax=np.nanmax(data),cmap=my_cmap)
		plt.autoscale()
		
		#title the plot
		plotTitle=plotTitle=FILE_NAME[:-4]
		if SDS_NAME != 'PBLH':
			plt.title('{0}\n {1}'.format(plotTitle, SDS_NAME + ' at the surface'))
		else:
			plt.title('{0}\n {1}'.format(plotTitle, SDS_NAME))
		fig = plt.gcf()
		cb = m.colorbar()
		cb.set_label(SDS_NAME)
		plt.show()
		
		is_save=str(input('\n Would you like to save this map? Please enter Y or N \n'))
		if is_save == 'Y' or is_save == 'y':
			#saves as a png if the user would like
			pngfile = '{0}.png'.format(FILE_NAME[:-4])
			fig.savefig(pngfile)