#=======================================================================
# Examples for Testing Relative Vorticity Method
# (Using ECMWF ERA-40 reanalysis)
#
# Program uses the ImageMagick convert program to create JPEG files.
#
# Copyright (c) 2004 by Johnny Lin.  For licensing, distribution
# conditions, contact information, and additional documentation see
# the URL http://www.johnny-lin.com/py_pkgs/atmqty/doc/.
#=======================================================================


#- Import modules:

try:  execfile('/home/jlin/.pythonrc.py')
except:  print "test_vort.py:  Ignore Johnny Lin's path settings."

import atmqty
import cdms
import IaGraph
from IaGraph import *
import MV
import Numeric as N
import os


#- Format lat/lon axes and set contour levels for the 2 fields:

sysvar = IaGraph.Sysvar()
sysvar.__class__.x_tickvalues = \
    {0:'0', 90:'90E', 180:'180E', 270:'270E', 360:'0'}
sysvar.__class__.y_tickvalues = \
    {-90:'90S', -45:'45S', 0:'EQ', 45:'45N', 90:'90N'}

tmp = N.array([0., 2e-5, 4e-5, 6e-5, 8e-5, 1e-4, 1.5e-4, 2e-4])
c_lev_vort = N.concatenate([-tmp[::-1], tmp[1:]])
N.putmask( c_lev_vort, N.where(N.absolute(c_lev_vort) < 1e-12, 1, 0), 0. )


#- Read ERA-40 reanalysis data and flip latitudes so the latitudes
#  go from S to N:

f = cdms.open('era40_2001-01-15.nc')

T_obj = f['t']
lon = T_obj.getLongitude()[:]
lat = T_obj.getLatitude()[::-1]
lev = T_obj.getLevel()[:].astype(N.Float)

T = T_obj(level=(1000.,100.), raw=1)
u = f['u'](level=(1000.,100.), raw=1)
v = f['v'](level=(1000.,100.), raw=1)
q = f['q'](level=(1000.,100.), raw=1)
vort = f['vo'](level=(1000.,100.), raw=1)

T = N.transpose(T, axes=(1,2,0))[::-1,:,:]
u = N.transpose(u, axes=(1,2,0))[::-1,:,:]
v = N.transpose(v, axes=(1,2,0))[::-1,:,:]
q = N.transpose(q, axes=(1,2,0))[::-1,:,:]
vort = N.transpose(vort, axes=(1,2,0))[::-1,:,:]


#- Compare relative vorticity from ECMWF and calculated from ECMWF
#  fields using the AtmQty class, on pressure levels.  Plot the 500 
#  hPa level:

x = atmqty.AtmQty( lat, lon, lev, lev_type="pressure" )
x.add_qty({'T':T, 'u':u, 'v':v, 'q':q})
x.vort_rel()

window(0)
base_fn = 'test_vort_AQvort'
contour( x.get_qty_asMA('vort_rel')[:,:,3], x.lon, x.lat \
       , title='AtmQty Calc. Rel. Vorticity' \
       , continents=1, c_levels=c_lev_vort )
active2gif('AQd_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQd_'+base_fn+'.gif AQd_'+base_fn+'.jpg')

window(1)
base_fn = 'test_vort_ECMvort'
contour( vort[:,:,3], lon, lat \
       , title='ERA-40 Rel. Vorticity' \
       , continents=1, c_levels=c_lev_vort )
active2gif('AQd_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQd_'+base_fn+'.gif AQd_'+base_fn+'.jpg')


#- Close file:

f.close()


# ===== end file =====