rcmet/src/main/python/rcmes/toolkit/plots.py - climate - Git at Google

 #
 #  Licensed to the Apache Software Foundation (ASF) under one or more
 #  contributor license agreements.  See the NOTICE file distributed with
 #  this work for additional information regarding copyright ownership.
 #  The ASF licenses this file to You under the Apache License, Version 2.0
 #  (the "License"); you may not use this file except in compliance with
 #  the License.  You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
 """Module that handles the generation of data plots"""


 # Import Statements

 from math import floor, log
 from matplotlib import pyplot as plt
 from mpl_toolkits.basemap import Basemap
 import matplotlib as mpl
 import numpy as np
 import os

 def pow_round(x):
     '''
      Function to round x to the nearest power of 10
     '''
     return 10 ** (floor(log(x, 10) - log(0.5, 10)))

 def calc_nice_color_bar_values(mymin, mymax, target_nlevs):
     '''
      Function to help make nicer plots.

      Calculates an appropriate min, max and number of intervals to use in a color bar
      such that the labels come out as round numbers.

      i.e. often, the color bar labels will come out as  0.1234  0.2343 0.35747 0.57546
      when in fact you just want  0.1, 0.2, 0.3, 0.4, 0.5 etc


      Method::
          Adjusts the max,min and nlevels slightly so as to provide nice round numbers.

      Input::
         mymin        - minimum of data range (or first guess at minimum color bar value)
         mymax        - maximum of data range (or first guess at maximum color bar value)
         target_nlevs - approximate number of levels/color bar intervals you would like to have

      Output::
         newmin       - minimum value of color bar to use
         newmax       - maximum value of color bar to use
         new_nlevs    - number of intervals in color bar to use
         * when all of the above are used, the color bar should have nice round number labels.
     '''
     myrange = mymax - mymin
     # Find target color bar label interval, given target number of levels.
     #  NB. this is likely to be not a nice rounded number.
     target_interval = myrange / float(target_nlevs)

     # Find power of 10 that the target interval lies in
     nearest_ten = pow_round(target_interval)

     # Possible interval levels,
     #  i.e.  labels of 1,2,3,4,5 etc are OK,
     #        labels of 2,4,6,8,10 etc are OK too
     #        labels of 3,6,9,12 etc are NOT OK (as defined below)
     #  NB.  this is also true for any multiple of 10 of these values
     #    i.e.  0.01,0.02,0.03,0.04 etc are OK too.
     pos_interval_levels = np.array([1, 2, 5])

     # Find possible intervals to use within this power of 10 range
     candidate_intervals = (pos_interval_levels * nearest_ten)

     # Find which of the candidate levels is closest to the target level
     absdiff = abs(target_interval - candidate_intervals)

     rounded_interval = candidate_intervals[np.where(absdiff == min(absdiff))]

     # Define actual nlevels to use in colorbar
     nlevels = myrange / rounded_interval

     # Define the color bar labels
     newmin = mymin - mymin % rounded_interval

     all_labels = np.arange(newmin, mymax + rounded_interval, rounded_interval)

     newmin = all_labels.min()
     newmax = all_labels.max()

     new_nlevs = int(len(all_labels)) - 1

     return newmin, newmax, new_nlevs

 def draw_cntr_map_single(pVar, lats, lons, mnLvl, mxLvl, pTitle, pName, pType = 'png', cMap = None):
     '''
     Purpose::
         Plots a filled contour map.

     Input::
         pVar - 2d array of the field to be plotted with shape (nLon, nLat)
         lon - array of longitudes
         lat - array of latitudes
         mnLvl - an integer specifying the minimum contour level
         mxLvl - an integer specifying the maximum contour level
         pTitle - a string specifying plot title
         pName  - a string specifying the filename of the plot
         pType  - an optional string specifying the filetype, default is .png
         cMap - an optional matplotlib.LinearSegmentedColormap object denoting the colormap,
                default is matplotlib.pyplot.cm.jet

         TODO: Let user specify map projection, whether to mask bodies of water??

     '''
     if cMap is None:
         cMap = plt.cm.jet

     # Set up the figure
     fig = plt.figure()
     ax = fig.gca()

     # Determine the map boundaries and construct a Basemap object
     lonMin = lons.min()
     lonMax = lons.max()
     latMin = lats.min()
     latMax = lats.max()
     m = Basemap(projection = 'cyl', llcrnrlat = latMin, urcrnrlat = latMax,
             llcrnrlon = lonMin, urcrnrlon = lonMax, resolution = 'l', ax = ax)

     # Draw the borders for coastlines and countries
     m.drawcoastlines(linewidth = 1)
     m.drawcountries(linewidth = .75)

     # Draw 6 parallels / meridians.
     m.drawmeridians(np.linspace(lonMin, lonMax, 5), labels = [0, 0, 0, 1])
     m.drawparallels(np.linspace(latMin, latMax, 5), labels = [1, 0, 0, 1])

     # Convert lats and lons to projection coordinates
     if lats.ndim == 1 and lons.ndim == 1:
         lons, lats = np.meshgrid(lons, lats)
     x, y = m(lons, lats)

     # Plot data with filled contours
     nsteps = 24
     mnLvl, mxLvl, nsteps = calc_nice_color_bar_values(mnLvl, mxLvl, nsteps)
     spLvl = (mxLvl - mnLvl) / nsteps
     clevs = np.arange(mnLvl, mxLvl, spLvl)
     cs = m.contourf(x, y, pVar, cmap = cMap)

     # Add a colorbar and save the figure
     cbar = m.colorbar(cs, ax = ax, pad = .05)
     plt.title(pTitle)
     fig.savefig('%s.%s' %(pName, pType))

 def draw_time_series_plot(data, times, myfilename, myworkdir, data2='', mytitle='', ytitle='Y', xtitle='time', year_labels=True):
     '''
      Purpose::
          Function to draw a time series plot

      Input::
          data - a masked numpy array of data masked by missing values
          times - a list of python datetime objects
          myfilename - stub of png file created e.g. 'myfile' -> myfile.png
          myworkdir - directory to save images in
          data2 - (optional) second data line to plot assumes same time values)
          mytitle - (optional) chart title
     	 xtitle - (optional) y-axis title
     	 ytitle - (optional) y-axis title

      Output::
          no data returned from function
          Image file produced with name {filename}.png
     '''
     print 'Producing time series plot'

     fig = plt.figure()
     ax = fig.gca()

     if year_labels == False:
         xfmt = mpl.dates.DateFormatter('%b')
         ax.xaxis.set_major_formatter(xfmt)

     # x-axis title
     plt.xlabel(xtitle)

     # y-axis title
     plt.ylabel(ytitle)

     # Main title
     fig.suptitle(mytitle, fontsize=12)

     # Set y-range to sensible values
     # NB. if plotting two lines, then make sure range appropriate for both datasets
     ymin = data.min()
     ymax = data.max()

     # If data2 has been passed in, then set plot range to fit both lines.
     # NB. if data2 has been passed in, then it is an array, otherwise it defaults to an empty string
     if type(data2) != str:
         ymin = min(data.min(), data2.min())
         ymax = max(data.max(), data2.max())

     # add a bit of padding so lines don't touch bottom and top of the plot
     ymin = ymin - ((ymax - ymin) * 0.1)
     ymax = ymax + ((ymax - ymin) * 0.1)

     # Set y-axis range
     plt.ylim((ymin, ymax))

     # Make plot, specifying marker style ('x'), linestyle ('-'), linewidth and line color
     line1 = ax.plot_date(times, data, 'bo-', markersize=6, linewidth=2, color='#AAAAFF')
     # Make second line, if data2 has been passed in.
     # TODO:  Handle the optional second dataset better.  Maybe set the Default to None instead
     # of an empty string
     if type(data2) != str:
         line2 = ax.plot_date(times, data2, 'rx-', markersize=6, linewidth=2, color='#FFAAAA')
         lines = []
         lines.extend(line1)
         lines.extend(line2)
         fig.legend((lines), ('model', 'obs'), loc='upper right')

     fig.savefig(myworkdir + '/' + myfilename + '.png')
	#
	# Licensed to the Apache Software Foundation (ASF) under one or more
	# contributor license agreements. See the NOTICE file distributed with
	# this work for additional information regarding copyright ownership.
	# The ASF licenses this file to You under the Apache License, Version 2.0
	# (the "License"); you may not use this file except in compliance with
	# the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	#
	"""Module that handles the generation of data plots"""


	# Import Statements

	from math import floor, log
	from matplotlib import pyplot as plt
	from mpl_toolkits.basemap import Basemap
	import matplotlib as mpl
	import numpy as np
	import os

	def pow_round(x):
	'''
	Function to round x to the nearest power of 10
	'''
	return 10 ** (floor(log(x, 10) - log(0.5, 10)))

	def calc_nice_color_bar_values(mymin, mymax, target_nlevs):
	'''
	Function to help make nicer plots.

	Calculates an appropriate min, max and number of intervals to use in a color bar
	such that the labels come out as round numbers.

	i.e. often, the color bar labels will come out as 0.1234 0.2343 0.35747 0.57546
	when in fact you just want 0.1, 0.2, 0.3, 0.4, 0.5 etc


	Method::
	Adjusts the max,min and nlevels slightly so as to provide nice round numbers.

	Input::
	mymin - minimum of data range (or first guess at minimum color bar value)
	mymax - maximum of data range (or first guess at maximum color bar value)
	target_nlevs - approximate number of levels/color bar intervals you would like to have

	Output::
	newmin - minimum value of color bar to use
	newmax - maximum value of color bar to use
	new_nlevs - number of intervals in color bar to use
	* when all of the above are used, the color bar should have nice round number labels.
	'''
	myrange = mymax - mymin
	# Find target color bar label interval, given target number of levels.
	# NB. this is likely to be not a nice rounded number.
	target_interval = myrange / float(target_nlevs)

	# Find power of 10 that the target interval lies in
	nearest_ten = pow_round(target_interval)

	# Possible interval levels,
	# i.e. labels of 1,2,3,4,5 etc are OK,
	# labels of 2,4,6,8,10 etc are OK too
	# labels of 3,6,9,12 etc are NOT OK (as defined below)
	# NB. this is also true for any multiple of 10 of these values
	# i.e. 0.01,0.02,0.03,0.04 etc are OK too.
	pos_interval_levels = np.array([1, 2, 5])

	# Find possible intervals to use within this power of 10 range
	candidate_intervals = (pos_interval_levels * nearest_ten)

	# Find which of the candidate levels is closest to the target level
	absdiff = abs(target_interval - candidate_intervals)

	rounded_interval = candidate_intervals[np.where(absdiff == min(absdiff))]

	# Define actual nlevels to use in colorbar
	nlevels = myrange / rounded_interval

	# Define the color bar labels
	newmin = mymin - mymin % rounded_interval

	all_labels = np.arange(newmin, mymax + rounded_interval, rounded_interval)

	newmin = all_labels.min()
	newmax = all_labels.max()

	new_nlevs = int(len(all_labels)) - 1

	return newmin, newmax, new_nlevs

	def draw_cntr_map_single(pVar, lats, lons, mnLvl, mxLvl, pTitle, pName, pType = 'png', cMap = None):
	'''
	Purpose::
	Plots a filled contour map.

	Input::
	pVar - 2d array of the field to be plotted with shape (nLon, nLat)
	lon - array of longitudes
	lat - array of latitudes
	mnLvl - an integer specifying the minimum contour level
	mxLvl - an integer specifying the maximum contour level
	pTitle - a string specifying plot title
	pName - a string specifying the filename of the plot
	pType - an optional string specifying the filetype, default is .png
	cMap - an optional matplotlib.LinearSegmentedColormap object denoting the colormap,
	default is matplotlib.pyplot.cm.jet

	TODO: Let user specify map projection, whether to mask bodies of water??

	'''
	if cMap is None:
	cMap = plt.cm.jet

	# Set up the figure
	fig = plt.figure()
	ax = fig.gca()

	# Determine the map boundaries and construct a Basemap object
	lonMin = lons.min()
	lonMax = lons.max()
	latMin = lats.min()
	latMax = lats.max()
	m = Basemap(projection = 'cyl', llcrnrlat = latMin, urcrnrlat = latMax,
	llcrnrlon = lonMin, urcrnrlon = lonMax, resolution = 'l', ax = ax)

	# Draw the borders for coastlines and countries
	m.drawcoastlines(linewidth = 1)
	m.drawcountries(linewidth = .75)

	# Draw 6 parallels / meridians.
	m.drawmeridians(np.linspace(lonMin, lonMax, 5), labels = [0, 0, 0, 1])
	m.drawparallels(np.linspace(latMin, latMax, 5), labels = [1, 0, 0, 1])

	# Convert lats and lons to projection coordinates
	if lats.ndim == 1 and lons.ndim == 1:
	lons, lats = np.meshgrid(lons, lats)
	x, y = m(lons, lats)

	# Plot data with filled contours
	nsteps = 24
	mnLvl, mxLvl, nsteps = calc_nice_color_bar_values(mnLvl, mxLvl, nsteps)
	spLvl = (mxLvl - mnLvl) / nsteps
	clevs = np.arange(mnLvl, mxLvl, spLvl)
	cs = m.contourf(x, y, pVar, cmap = cMap)

	# Add a colorbar and save the figure
	cbar = m.colorbar(cs, ax = ax, pad = .05)
	plt.title(pTitle)
	fig.savefig('%s.%s' %(pName, pType))

	def draw_time_series_plot(data, times, myfilename, myworkdir, data2='', mytitle='', ytitle='Y', xtitle='time', year_labels=True):
	'''
	Purpose::
	Function to draw a time series plot

	Input::
	data - a masked numpy array of data masked by missing values
	times - a list of python datetime objects
	myfilename - stub of png file created e.g. 'myfile' -> myfile.png
	myworkdir - directory to save images in
	data2 - (optional) second data line to plot assumes same time values)
	mytitle - (optional) chart title
	xtitle - (optional) y-axis title
	ytitle - (optional) y-axis title

	Output::
	no data returned from function
	Image file produced with name {filename}.png
	'''
	print 'Producing time series plot'

	fig = plt.figure()
	ax = fig.gca()

	if year_labels == False:
	xfmt = mpl.dates.DateFormatter('%b')
	ax.xaxis.set_major_formatter(xfmt)

	# x-axis title
	plt.xlabel(xtitle)

	# y-axis title
	plt.ylabel(ytitle)

	# Main title
	fig.suptitle(mytitle, fontsize=12)

	# Set y-range to sensible values
	# NB. if plotting two lines, then make sure range appropriate for both datasets
	ymin = data.min()
	ymax = data.max()

	# If data2 has been passed in, then set plot range to fit both lines.
	# NB. if data2 has been passed in, then it is an array, otherwise it defaults to an empty string
	if type(data2) != str:
	ymin = min(data.min(), data2.min())
	ymax = max(data.max(), data2.max())

	# add a bit of padding so lines don't touch bottom and top of the plot
	ymin = ymin - ((ymax - ymin) * 0.1)
	ymax = ymax + ((ymax - ymin) * 0.1)

	# Set y-axis range
	plt.ylim((ymin, ymax))

	# Make plot, specifying marker style ('x'), linestyle ('-'), linewidth and line color
	line1 = ax.plot_date(times, data, 'bo-', markersize=6, linewidth=2, color='#AAAAFF')
	# Make second line, if data2 has been passed in.
	# TODO: Handle the optional second dataset better. Maybe set the Default to None instead
	# of an empty string
	if type(data2) != str:
	line2 = ax.plot_date(times, data2, 'rx-', markersize=6, linewidth=2, color='#FFAAAA')
	lines = []
	lines.extend(line1)
	lines.extend(line2)
	fig.legend((lines), ('model', 'obs'), loc='upper right')

	fig.savefig(myworkdir + '/' + myfilename + '.png')