#=======================================================================
# Examples for Testing Function transfrom_AQobj.
# (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_trans.py:  Ignore Johnny Lin's path settings."

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


#- Set system variables for lon/lat axes tick labels and c_levels
#  vectors for u graphs and T graph:

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_levels_umid = N.arange(13)*10. - 40.
c_levels_T500 = N.arange(13)*5. + 220.


#- Read data for AtmQty object x from file, global domain with no 
#  missing values, in pressure coordinates.  Reverse latitude order
#  so IaGraph can correctly make the plot and transpose order of
#  the data axes to fit the AtmQty format:

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

T = MV.filled(f('t', squeeze=1))
u = MV.filled(f('u', squeeze=1))
v = MV.filled(f('v', squeeze=1))
q = MV.filled(f('q', squeeze=1))

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

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,:,:]

x = atmqty.AtmQty( lat, lon, lev, lev_type="pressure", missing=1e+29 )

x.add_qty({'T':T, 'u':u, 'v':v, 'q':q})
x.theta()
x.alt()

window(0)
base_fn = 'test_trans_x_u500'
contour( MV.masked_values(x.get_qty('u')[:,:,3], x.missing) \
       , x.lon, x.lat, continents=1, c_levels=c_levels_umid )
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_trans_x_T500'
contour( MV.masked_values(x.get_qty('T')[:,:,3], x.missing) \
       , x.lon, x.lat, continents=1, c_levels=c_levels_T500 )
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')


#- Set up AtmQty object x2 with different pressure levels and lat/lon
#  than x:

lev2 = N.array([900., 800., 500., 350.])
lat2 = N.arange(20)*4.0 - 30.0
lon2 = N.arange(50)*4.0 + 20.0
x2   = atmqty.AtmQty( lat2, lon2, lev2, lev_type="pressure", missing=1e+29 )
x2   = atmqty.transform_AQobj.transform_AQobj(x, x2)

window(2)
base_fn = 'test_trans_x2_u500'
contour( MV.masked_values(x2.get_qty('u')[:,:,2], x2.missing) \
       , x2.lon, x2.lat, continents=1, c_levels=c_levels_umid )
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')


#- Set up AtmQty object x3 with isentropic levels, same lat/lon as x, 
#  and different missing values:

lev3 = N.array([300., 315., 330., 350.])
x3   = atmqty.AtmQty( lat, lon, lev3, lev_type="isentropic", missing=1e+14 )
x3   = atmqty.transform_AQobj.transform_AQobj(x, x3)
x3.press()

window(3)
base_fn = 'test_trans_x3_u315K'
contour( MV.masked_values(x3.get_qty('u')[:,:,1], x3.missing) \
       , x3.lon, x3.lat, continents=1, c_levels=c_levels_umid )
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(4)
base_fn = 'test_trans_x3_p315K'
contour( MV.masked_values(x3.get_qty('p')[:,:,1], x3.missing) \
       , x3.lon, x3.lat, continents=1 )
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(5)
sysvar.__class__.x_tickvalues = sysvar.__class__.y_tickvalues = None
ilat = 18
ilon = 35
base_fn = 'test_trans_x3_x_Tpt'
print 'lat', lat[ilat], 'lon', lon[ilon]
xalt  = MV.masked_values(x.get_qty('alt')[ilat,ilon,:], x3.missing)
xu    = MV.masked_values(x.get_qty('T')[ilat,ilon,:], x3.missing)
x3alt = MV.masked_values(x3.get_qty('alt')[ilat,ilon,:], x3.missing)
x3u   = MV.masked_values(x3.get_qty('T')[ilat,ilon,:], x3.missing)
plot( xu, xalt, psym=-2, linestyle=0 \
    , xtitle='* = pressure; diamond = isentropic', ytitle='Altitude [m]' )
oplot( x3u, x3alt, psym=-4, linestyle=2 )
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 data file:

f.close()


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