#=======================================================================
# Examples for Testing IPV Method
# (Using NCEP/NCAR 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_ipv.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, 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'}

c_lev_ipv = N.arange(17)*1e-6 - 7e-6
N.putmask( c_lev_ipv, N.where(N.absolute(c_lev_ipv) < 1e-12, 1, 0), 0. )


#- Read NCEP/NCAR reanalysis data on isentropic surfaces (except for
#  IPV, which is read below right before it is plotted).  Set the
#  missing value value to use for this script and the index of the
#  level to plot at (ilev=1 => 315 K):

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

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

T = MV.filled( T_obj )
u = MV.filled( f['UGRD'] )
v = MV.filled( f['VGRD'] )
rh = MV.filled( f['RH']/100.0 )
n2 = MV.filled( f['BVF2'] )

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,:,:]
rh = N.transpose(rh, axes=(1,2,0))[::-1,:,:]
n2 = N.transpose(n2, axes=(1,2,0))[::-1,:,:]

miss_val = T_obj.missing_value[0]
ilev = 1


#- Calculate and plot IPV calculated from NCEP/NCAR reanalysis fields 
#  on isentropic levels (not including buoyancy frequency):

x1 = atmqty.AtmQty( lat, lon, lev, lev_type="isentropic", missing=miss_val )
x1.add_qty({'T':T, 'u':u, 'v':v, 'RH':rh})
x1.ipv()

window(0)
base_fn = 'test_ipv_AQipv1'
contour( x1.get_qty_asMA('pv')[:,:,ilev], x1.lon, x1.lat \
       , title='AtmQty Calc. IPV (calc. N2)' \
       , continents=1, c_levels=c_lev_ipv )
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')


#- Calculate and plot IPV calculated from NCEP/NCAR reanalysis fields 
#  on isentropic levels (including buoyancy frequency):

x2 = atmqty.AtmQty( lat, lon, lev, lev_type="isentropic", missing=miss_val )
x2.add_qty({'T':T, 'u':u, 'v':v, 'RH':rh, 'N2':n2})
x2.ipv()

window(1)
base_fn = 'test_ipv_AQipv2'
contour( x2.get_qty_asMA('pv')[:,:,ilev], x2.lon, x2.lat \
       , title='AtmQty Calc. IPV (read-in N2)' \
       , continents=1, c_levels=c_lev_ipv )
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')


#- Read in and plot IPV calculated by the reanalysis:

pv = MV.transpose(f['PV___'], axes=(1,2,0))[::-1,:,:] * 1e-6

window(2)
base_fn = 'test_ipv_NNipv'
contour( pv[:,:,ilev], lon, lat \
       , title='NN Reanalysis IPV' \
       , continents=1, c_levels=c_lev_ipv )
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 =====