Animate surface relative vorticity on blue marble earth

demo_animate_relative_vortcitiy.py

import netCDF4
import matplotlib.pyplot as plt
import numpy
from mpl_toolkits.basemap import Basemap
import math

lon = netCDF4.Dataset('/archive/gold/datasets/MOM6z_SIS_025/siena/mosaic.unpacked/ocean_hgrid.nc').variables['x'][1::2,1::2]
lat = netCDF4.Dataset('/archive/gold/datasets/MOM6z_SIS_025/siena/mosaic.unpacked/ocean_hgrid.nc').variables['y'][1::2,1::2]
dx = netCDF4.Dataset('/archive/gold/datasets/MOM6z_SIS_025/siena/mosaic.unpacked/ocean_hgrid.nc').variables['dx'][:]
dx = dx[1::2,1::2]+numpy.roll(dx,-1,axis=-1)[1::2,1::2]
dy = netCDF4.Dataset('/archive/gold/datasets/MOM6z_SIS_025/siena/mosaic.unpacked/ocean_hgrid.nc').variables['dy'][:]
topRow = dy[-1,1::2]
dy = dy[1::2,1::2]+numpy.roll(dy,-1,axis=0)[1::2,1::2]; dy[-1,:] = topRow + topRow[::-1]
area = netCDF4.Dataset('/archive/gold/datasets/MOM6z_SIS_025/siena/mosaic.unpacked/ocean_hgrid.nc').variables['area'][:]
area = area[:,1::2] + numpy.roll(area,-1,axis=-1)[:,1::2]
topRow = area[-1,:]
area = area[1::2,:] + numpy.roll(area,-1,axis=0)[1:2,:]; area[-1,:] = topRow + topRow[::-1]


dailyFiles = '/archive/aja/awg/tikal_201403/CM4_c192L48_am4a1_2000/gfdl.ncrc2-intel-prod-openmp/history/unpack/000[3456]0101.ocean_daily.nc'
ssu = netCDF4.MFDataset(dailyFiles).variables['ssu']
ssv = netCDF4.MFDataset(dailyFiles).variables['ssv']
dayNum = netCDF4.MFDataset(dailyFiles).variables['time']

def JulianDate(n):
  """Returns a (year/month/day) date for the given Julian day number."""
  a = int(n) + 32044
  b = (4*a + 3)//146097
  c = a - (146097*b)//4
  d = (4*c + 3)//1461
  e = c - (1461*d)//4
  m = (5*e + 2)//153
  day = e + 1 - (153*m + 2)//5
  month = m + 3 - 12*(m//10)
  year = 100*b + d - 4800 + m/10
  return year, month, day

def vort(u,v,dx,dy,area):
  """This function calculate relative vorticity"""
  vort = u*dy - numpy.roll(u*dy, -1, axis=0) # -d_j U
  vort[-1,:] = u[-1,:]*dy[-1,:] - u[-1,::-1]*dy[-1,::-1] # Tri-polar fold
  vort = vort + (numpy.roll(v*dx, -1, axis=1) - v*dx) # +d_i V
  return vort / area

def makePlot(fig, ssu, ssv, dx, dy, area, n, year, month, day):
  """This function plots ssh[n] to fig."""
  plt.clf()
  ax = fig.add_axes([0.03,0.03,0.94,0.94])
  m = Basemap(projection='kav7',lon_0=-120,resolution=None)
  m.drawmapboundary(fill_color='0.3')
  im0 = m.bluemarble(scale=0.25)
  im1 = m.pcolormesh(numpy.minimum(lon,60.),lat, 86400./4/math.pi*vort(ssu[n],ssv[n],dx,dy,area) ,shading='flat',cmap=plt.cm.RdYlBu,latlon=True)
  plt.clim(-.2,0.2)
  cbaxes = fig.add_axes([0.01, 0.91, 0.16, 0.03]) 
  cb = plt.colorbar(im1, cax = cbaxes, orientation='horizontal', extend='both', ticks=[-.2,0.,.2])
  cb.set_label('[non-dim]',fontsize=10)
  cb.ax.tick_params(labelsize=10)
  plt.figtext(0.09,.975,'Surface relative vorticity',verticalalignment='top',size='medium',horizontalalignment='center')
  plt.figtext(0.81,.975,'Model date = %4.4i.%2.2i.%2.2i'%(year,month,day),verticalalignment='top')
  plt.draw()

fig = plt.figure(figsize=(12.8,7.2),dpi=100) # 720p
#plt.show(block=False)
for n in range(dayNum.shape[0]):
  year, month, day = JulianDate(1721426 + dayNum[n])
  print 'record %i (y/m/d %i/%i/%i)'%(n,year,month,day)
  makePlot(fig, ssu, ssv, dx, dy, area, n, year, month, day)
  plt.savefig('figs/q.%4.4i.%2.2i.%2.2i.png'%(year,month,day))