#=======================================================================
# Examples for Testing Function transfrom_AQobj.
# (Additional examples using synthetic output.)
#
# 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_add.py:  Ignore Johnny Lin's path settings."

import atmqty
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:

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

c_levels_u500 = N.arange(12)*0.2 + 11.
c_levels_v500 = N.arange(12)*0.2 + 2.
c_levels_T500 = N.arange(16)*1. + 242.


#- Set up synthetic data for AtmQty object x, global domain with 
#  some missing values, in pressure coordinates.  Variables T, u, 
#  v are created:

nx = 73
ny = 37
lon = N.arange(nx) * 5.0 - 180.0
lat = N.arange(ny) * 5.0 - 90.0
lev = N.array([850., 700., 500., 300., 250.])

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

lat_2d = N.reshape( N.repeat(lat, nx), (ny, nx) )
lon_2d = N.reshape( N.repeat(N.reshape(lon,(1,nx)), ny), (ny, nx) )

u_2d = N.sin(lon_2d*N.pi/180.)*N.sin((lat_2d-90.)*N.pi/30.)
v_2d = N.sin((lon_2d-30.)*N.pi/180.)*N.sin((lat_2d-90.)*N.pi/30.)
N.putmask( u_2d, N.where(N.absolute(u_2d) < 1e-10, 1, 0), 0. )
N.putmask( v_2d, N.where(N.absolute(v_2d) < 1e-10, 1, 0), 0. )

T_2d = N.outerproduct( N.sin(lat*N.pi/360.) \
                     , N.cos(lon*N.pi/360.) ) * 10. + 257.3

T = N.zeros(x.qty_shape, typecode=N.Float)
u = N.zeros(x.qty_shape, typecode=N.Float)
v = N.zeros(x.qty_shape, typecode=N.Float)

T[:,:,0] = T_2d
u[:,:,0] = u_2d
v[:,:,0] = v_2d

for ilev in range(1,len(lev)):
    T[:,:,ilev] = T[:,:,ilev-1] - N.sin(ilev*4.0/N.pi)*5.0
    u[:,:,ilev] = u[:,:,ilev-1] + N.sin(ilev*4.0/N.pi)*4.0 + 3.0
    v[:,:,ilev] = v[:,:,ilev-1] + N.sin(ilev*4.0/N.pi)*2.0

T[23:30,17,3] = 1e+29
u[10:18,23:32,1] = 1e+29
v[17:24,62:70,2] = 1e+29

x.add_qty({'T':T, 'u':u, 'v':v})
x.theta()
x.alt(T_surf=T_2d)

window(0)
base_fn = 'test_trans_x_u500'
contour( MV.masked_values(x.get_qty('u')[:,:,2], x.missing) \
       , x.lon, x.lat, continents=1, title=base_fn \
       , c_levels = c_levels_u500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')

window(1)
base_fn = 'test_trans_x_v500'
contour( MV.masked_values(x.get_qty('v')[:,:,2], x.missing) \
       , x.lon, x.lat, continents=1, title=base_fn \
       , c_levels = c_levels_v500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')

window(2)
base_fn = 'test_trans_x_T500'
contour( MV.masked_values(x.get_qty('T')[:,:,2], x.missing) \
       , x.lon, x.lat, continents=1, title=base_fn \
       , c_levels = c_levels_T500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')


#- Set up AtmQty object x2 with same pressure levels as x but lat/lon 
#  sub-set:

lat2 = N.arange(40)*2.0 - 20.0
lon2 = N.arange(50)*4.0 - 90.0
x2   = atmqty.AtmQty( lat2, lon2, lev, lev_type="pressure", missing=1e+29 )
x2   = atmqty.transform_AQobj.transform_AQobj(x, x2)

window(3)
base_fn = 'test_trans_x2_u500'
contour( MV.masked_values(x2.get_qty('u')[:,:,2], x2.missing) \
       , x2.lon, x2.lat, continents=1, title=base_fn \
       , c_levels = c_levels_u500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')


#- Set up AtmQty object x2a with same pressure levels and lat/lon as 
#  x, but interpolate the data from x2, not x:

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

window(4)
base_fn = 'test_trans_x2a_u500'
contour( MV.masked_values(x2a.get_qty('u')[:,:,2], x2a.missing) \
       , x2a.lon, x2a.lat, continents=1, title=base_fn \
       , c_levels = c_levels_u500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')


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

lev4 = N.array([900., 800., 600., 500.])
x3   = atmqty.AtmQty( lat2, lon2, lev4, lev_type="pressure", missing=1e+17 )
x3   = atmqty.transform_AQobj.transform_AQobj(x, x3)

window(5)
base_fn = 'test_trans_x3_u500'
contour( MV.masked_values(x3.get_qty('u')[:,:,-1], x3.missing) \
       , x3.lon, x3.lat, continents=1, title=base_fn \
       , c_levels = c_levels_u500 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')


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

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

window(6)
sysvar.__class__.x_tickvalues = sysvar.__class__.y_tickvalues = None

base_fn = 'test_trans_x4_x_upt'
ilat = 22
ilon = 51
print 'lat', lat[ilat], 'lon', lon[ilon]
xalt  = MV.masked_values(x.get_qty('alt')[ilat,ilon,:], x4.missing)
xu    = MV.masked_values(x.get_qty('u')[ilat,ilon,:], x4.missing)
x4alt = MV.masked_values(x4.get_qty('alt')[ilat,ilon,:], x4.missing)
x4u   = MV.masked_values(x4.get_qty('u')[ilat,ilon,:], x4.missing)
plot( xu, xalt, psym=-2, linestyle=0 \
    , xtitle='* = pressure; diamond = isentropic', ytitle='Altitude [m]' \
    , title=base_fn )
oplot( x4u, x4alt, psym=-4, linestyle=2 )
active2gif('AQs_'+base_fn+'.gif')
os.system('convert -verbose -trim -border 8x8 -bordercolor white ' \
         +'-mattecolor black -frame 2x2 ' \
         +'AQs_'+base_fn+'.gif AQs_'+base_fn+'.jpg')


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