#!/usr/bin/python -tt
#=======================================================================
#                        General Documentation

"""Module has a single public function.

   See function docstring for description.
"""

#-----------------------------------------------------------------------
#                       Additional Documentation
#
# RCS Revision Code:
#   $Id: transform_AQobj.py,v 1.13 2004/05/20 22:34:31 jlin Exp $
#
# Modification History:
# - 28 Apr 2004:  Original by Johnny Lin, Computation Institute,
#   University of Chicago.  Passed reasonable tests.
#
# Notes:
# - Written for Python 2.2.
# - Module docstrings can be tested using the doctest module.  To
#   test, execute "python transform_AQobj.py".
# - The term "log" refers to the natural logarithm.
# - See import statements throughout for other packages required.
#
# 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/.
#=======================================================================




#---------------- Module General Import and Declarations ---------------

#- Overall imports:

import copy
import MA
import Numeric as N


#- Set module version number etc. to package version number etc.:

import package_version
__version__ = package_version.version
__author__  = package_version.author
__date__    = package_version.date
__credits__ = package_version.credits
del package_version




#---- Very Private Fctns.:  A Few Ops On Arrays With Missing Values ----

def __exp(a, missing):
    """Exponential function on array a, missing values equal missing.
    """
    return MA.filled( MA.exp( MA.masked_values(a, missing) ) \
                    , missing )

def __log(a, missing):
    """Natural logarithm of array a, missing values equal missing.
    """
    return MA.filled( MA.log( MA.masked_values(a, missing) ) \
                    , missing )




#------------- Private Function:  Horizontal Interpolation -------------

def _horizontal_interpolation(x, y):
    """Deep copy of x interpolated to y lat/lon grid.

    Output is deep copy of x, but with all atmospheric quantities 
    interpolated to the horizontal grid described by y (the vertical 
    grid is not changed).  All attributes related to describing the
    new grid are changed/added to output as needed.  See docstring 
    for transform_AQobj for algorithm details.  Attributes of y that
    are unrelated to describing the new horizontal grid are ignored.

    Input Arguments:
    * x:  AtmQty object.  Not changed by function.
    * y:  AtmQty object.  Not changed by function.
    """
    import cdms
    import MV

    xcopy = copy.deepcopy(x)


    #- Create old and new grids as CDMS axis and grid objects:

    xlatAxis = cdms.createAxis(x.lat, id='lat'); xlatAxis.designateLatitude()
    xlonAxis = cdms.createAxis(x.lon);           xlonAxis.designateLongitude()
    ylatAxis = cdms.createAxis(y.lat, id='lat'); ylatAxis.designateLatitude()
    ylonAxis = cdms.createAxis(y.lon);           ylonAxis.designateLongitude()

    xgrid = cdms.createRectGrid( xlatAxis, xlonAxis, order="yx" )
    ygrid = cdms.createRectGrid( ylatAxis, ylonAxis, order="yx" )


    #- Set new quantity shape and do horizontal interpolation of all
    #  quantities:

    old_qty_shape   = xcopy.qty_shape
    new_qty_shape   = (y.qty_shape[0], y.qty_shape[1], xcopy.qty_shape[2])
    xcopy.qty_shape = new_qty_shape

    for aqty in xcopy.list_all_qty():
        old_quantity = xcopy.get_qty(aqty)
        new_quantity = N.zeros(new_qty_shape, typecode=N.Float)

        for ilev in xrange(old_qty_shape[2]):
            old_quantity_alev = \
                cdms.createVariable( MV.masked_values( old_quantity[:,:,ilev] \
                                                     , xcopy.missing ) \
                                   , axes = (xlatAxis, xlonAxis) \
                                   , grid = xgrid )
            new_quantity_alev = old_quantity_alev.regrid(ygrid)
            new_quantity[:,:,ilev] = \
                MV.filled(new_quantity_alev, xcopy.missing)

        del old_quantity
        xcopy.add_qty({aqty:new_quantity})


    #- Do horizontal interpolation of P_surf and T_surf boundary condi-
    #  tion attributes, if they exist (abc means "a boundary condition")
    #  and are not either scalars or None:

    for abc in ['P_surf', 'T_surf']:
        if N.size(xcopy.__dict__[abc]) != 1:
            tmp_old = MV.masked_values(xcopy.__dict__[abc], xcopy.missing)
            old_bc = cdms.createVariable( tmp_old \
                                        , axes = (xlatAxis, xlonAxis) \
                                        , grid = xgrid )
            xcopy.__dict__[abc] = \
                MV.filled( old_bc.regrid(ygrid), xcopy.missing )
            del tmp_old, old_bc


    #- Update lat and lon attributes and return:

    xcopy.lat = y.lat
    xcopy.lon = y.lon

    return xcopy




#-------------- Private Function:  Vertical Interpolation --------------

def _vertical_interpolation(x, y):
    """Deep copy of x interpolated from x levels to y levels.

    Output is a deep copy of x, but with all atmospheric quantities 
    interpolated to the vertical grid described by y (the horizontal 
    grid is not changed).  All attributes related to describing the
    new levels are changed/added to output as needed.  See docstring 
    for transform_AQobj for algorithm details.  Attributes of y that
    are unrelated to describing the new levels are ignored.

    Input Arguments:
    * x:  AtmQty object.  Not changed by function.
    * y:  AtmQty object.  Not changed by function.
    """
    import gemath
    from gemath.interp import interp


    #- Copy x; depending on what x and y lev_types are, calculate 
    #  values needed for the interpolations.  Variable key (not all
    #  variables described here actually exist):
    #  * old_ylev:  Value of the y.lev_type quantity at x.lev.
    #  * log_old_ylev:  Logarithm of old_ylev.
    #  * new_ylev:  Value of the y.lev_type quantity at y.lev.
    #  * log_new_ylev:  Logarithm of new_ylev.
    #  Sort log_old_ylev ascending to make sure the interpolation 
    #  routine will work correctly:

    xcopy = copy.deepcopy(x)

    if   (x.lev_type == 'pressure') and (y.lev_type == 'isentropic'):
        if xcopy.has_qty('theta'):
            if not N.allclose(xcopy.p0, y.p0):  xcopy.theta(p0=y.p0)
        else:
            xcopy.theta(p0=y.p0)
        log_old_ylev = __log(xcopy.get_qty('theta'), xcopy.missing)

    elif (x.lev_type == 'pressure') and (y.lev_type == 'pressure'):
        log_old_ylev = __log(xcopy.get_qty('p'), xcopy.missing)

    elif (x.lev_type == 'isentropic') and (y.lev_type == 'pressure'):
        if not xcopy.has_qty('p'):  xcopy.press()
        log_old_ylev = __log(xcopy.get_qty('p'), xcopy.missing)

    elif (x.lev_type == 'isentropic') and (y.lev_type == 'isentropic'):
        log_old_ylev = __log(xcopy.get_qty('theta'), xcopy.missing)

    else:
        raise ValueError, "transform_AQobj:  Bad lev_type specification"

    log_new_ylev  = N.log(y.lev)
    sort_data_idx = N.argsort(log_old_ylev)


    #- Quantities in list_qty_to_interp_log_log are linearly interpol-
    #  ated based on log(y.lev_type quantity) and log(quantity); rest 
    #  are linearly interpolated based on log(y.lev_type quantity) and 
    #  quantity.  xcopy.qty_shape is changed to new_qty_shape before
    #  interpolation so the newly interpolated values can be added to 
    #  qty using add_qty:

    list_qty_to_interp_log_log = ['p','T','theta']

    old_qty_shape   = xcopy.qty_shape
    new_qty_shape   = (xcopy.qty_shape[0], xcopy.qty_shape[1], y.qty_shape[2])
    xcopy.qty_shape = new_qty_shape

    for aqty in xcopy.list_all_qty():
        old_quantity = xcopy.get_qty(aqty)
        new_quantity = N.zeros(new_qty_shape, typecode=N.Float)

        if aqty in list_qty_to_interp_log_log:
            interp_quantity = __log(old_quantity, xcopy.missing)
        else:
            interp_quantity = old_quantity

        for ilat in xrange(old_qty_shape[0]):
            for ilon in xrange(old_qty_shape[1]):
                xdata = N.take( log_old_ylev[ilat,ilon,:] \
                              , sort_data_idx[ilat,ilon,:] )
                ydata = N.take( interp_quantity[ilat,ilon,:] \
                              , sort_data_idx[ilat,ilon,:] )
                new_quantity[ilat,ilon,:] = \
                    interp( ydata, xdata, log_new_ylev \
                          , missing=xcopy.missing )

        if aqty in list_qty_to_interp_log_log:
            new_quantity = __exp(new_quantity, xcopy.missing)

        del interp_quantity, old_quantity
        xcopy.add_qty({aqty:new_quantity})


    #- Update level related attributes and return:

    xcopy.lev_type = y.lev_type
    xcopy.lev      = y.lev

    return xcopy




#--------------------------- Public Function ---------------------------

def transform_AQobj(x, y):
    """Return copy of x transformed onto the "domain" of y.

    Calling Sequence:  Result = transform_AQobj(x, y)

    Both x and y are AtmQty objects.  Returns a deep copy of x with all
    atmospheric quantities in x.qty interpolated to the "domain" de-
    scribed by y.  Quantity-type attributes P_surf and T_surf are also 
    horizontally interpolated, if applicable.  The returned object also 
    may have changed/added attributes as part of the transformation 
    from the x "domain" to the y "domain" (e.g. lon).  Returned object 
    Result is of class AtmQty.

    This function is used primarily for regridding x atmospheric quan-
    tities, if x and y are different spatial domains.  It can also be
    used to create an object where the missing values in arrays in x 
    are set to y.missing.

    The "domain" of an AtmQty object is described by 5 (or 6, if p0 is
    applicable) attributes: lat, lon, lev, lev_type, and missing (and
    p0, if applicable).  All other attributes in y are ignored by this
    function.

    NB:  Since the Result quantities dictionary (and boundary condition
    attributes) are filled by interpolation and not recomputation, 
    though all the derived variables in x are consistent with one 
    another because they were all computed from the same given state 
    variables, all the derived variables in Result are only consistent 
    in the sense they were all interpolated from the same x counter-
    parts.


    Positional Input Arguments:
    * x:  AtmQty object.  Not changed by function.
    * y:  AtmQty object.  Not changed by function.


    References:
    * Drach, R., et al. (2002):  Climate Data Management System, Version
      3.3.  Livermore, CA:  Lawrence Livermore National Laboratory, UCRL-
      JC-134897.  See http://esg.llnl.gov/cdat/cdmsV3.3.pdf.


    Example to create xnew, which has the fields in x interpolated to the 
    domain in y, and with the missing value value the same as in y (xnew, 
    x, y are all AtmQty objects):
      from transform_AQobj import transform_AQobj
      xnew = transform_AQobj(x, y)


    Notes:
    * Horizontal regridding:  Latitude-longitude grid regridding is done
      using the regrid method of CDAT transient variables.  Information
      on CDAT is available at http://esg.llnl.gov/cdat/.

    * Vertical interpolation:  For pressure, temperature, and potential 
      temperature, interpolation is done on a log-log basis (i.e. linear 
      interpolation between log(p) or log(theta) and log(quantity)).  
      For all other fields, interpolation is done on a log-linear basis 
      (i.e. linear interpolation between log(p) or log(theta) and the 
      quantity).  Note that if there are missing values in quantities
      p or theta (such as from data gaps in temperature), the interpol-
      ation routine will ignore those points completely and assume that 
      interpolation using only the remaining points is legitimate.

    * Output values at domain locations in y that are outside of the 
      data domain specified in x are set to y.missing.  However, the
      atmospheric quantity corresponding to lev_type is set from the 
      lev attribute, and so there are no missing values, Also, if the
      quantity fc is in x, the transformed value of fc is recalculated
      directly from the y domain information (so there are also no
      missing values in fc).
    """
    import gemath
    from gemath.where_close import where_close


    #- Check input is of the correct type:

    if (x.__class__.__name__ != "AtmQty") or \
       (y.__class__.__name__ != "AtmQty"):
        raise TypeError, "transform_AQobj:  Arg(s) not class AtmQty"


    #- Initialize the working AtmQty object (xnew, which begins as a
    #  deep copy of x) and do vertical coordinate regridding of 
    #  quantities (if needed):

    if x.lev_type == y.lev_type:
        try:
            if N.allclose(x.lev, y.lev):
                xnew = copy.deepcopy(x)
            else:
                xnew = _vertical_interpolation(x, y)
        except:
            xnew = _vertical_interpolation(x, y)
    else:
        xnew = _vertical_interpolation(x, y)


    #- Horizontal coordinate regridding of quantities.  If x and y
    #  lat/lon domains are same, just pass through.  If they are
    #  different, or if the comparison cannot be made since the
    #  domain vectors are of different lengths, do a horizontal
    #  interpolation:

    try:
        if N.allclose(x.lat, y.lat) and N.allclose(x.lon, y.lon):
            xnew = xnew
        else:
            xnew = _horizontal_interpolation(xnew, y)
    except:
        xnew = _horizontal_interpolation(xnew, y)


    #- Change missing quantities elements in xnew to the value in y:

    if not N.allclose(x.missing, y.missing):
        for aqty in xnew.list_all_qty():
            miss_mask = where_close(xnew.get_qty(aqty), xnew.missing)
            N.putmask(xnew.get_qty(aqty), miss_mask, y.missing)
        xnew.missing = y.missing


    #- Ensure quantity corresponding to lev_type has no missing 
    #  values by overwriting with lev values:

    tmp = N.zeros( xnew.qty_shape, typecode=N.Float )
    for ilev in xrange(xnew.qty_shape[2]):
        tmp[:,:,ilev] = xnew.lev[ilev]

    if xnew.lev_type == 'pressure':
        xnew.add_qty({'p':tmp})
    elif xnew.lev_type == 'isentropic':
        xnew.add_qty({'theta':tmp})
    else:
        raise ValueError, "tranform_AQobj:  Bad lev_type"
    del tmp


    #- Ensure quantity fc (Coriolis parameter) is not interpolated 
    #  but computed from new domain information:

    if xnew.has_qty('fc'):  xnew.fc()


    #- Return transformed copy of x onto domain of y:

    return xnew




#-------------------------- Main:  Test Module -------------------------

#- Define additional examples for doctest to use:

__test__ = {'Additional Examples':
"""
#- Import modules:

>>> import atmqty
>>> import MV
>>> import Numeric as N


#- 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)

>>> ['%.7g' % x.get_qty('T')[i,65,2] for i in range(15,20)]
['249.3242', '249.4546', '249.5855', '249.7167', '249.8479']
>>> ['%.7g' % x.get_qty('u')[i,65,2] for i in range(15,20)]
['12.64021', '12.56337', '12.35342', '12.06663', '11.77985']
>>> ['%.7g' % x.get_qty('v')[i,65,2] for i in range(15,20)]
['3.939625', '3.818202', '1e+29', '1e+29', '1e+29']
>>> ['%.7g' % x.get_qty('theta')[i,65,2] for i in range(15,20)]
['303.9701', '304.1291', '304.2887', '304.4486', '304.6085']
>>> ['%.7g' % x.get_qty('alt')[i,65,2] for i in range(15,20)]
['4969.983', '4972.506', '4975.038', '4977.576', '4980.113']


#- 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)

>>> ['%.7g' % x2.get_qty('T')[i,35,2] for i in range(15,20)]
['250.5066', '250.6046', '250.8997', '250.8997', '251.1924']
>>> ['%.7g' % x2.get_qty('u')[i,35,2] for i in range(15,20)]
['11.40322', '11.37764', '11.30059', '11.30059', '11.37745']


#- 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)

>>> ['%.7g' % x3.get_qty('T')[i,35,2] for i in range(15,20)]
['252.0216', '252.1196', '252.4147', '252.4147', '252.7074']
>>> ['%.7g' % x3.get_qty('u')[i,35,2] for i in range(15,20)]
['8.562556', '8.536971', '8.459926', '8.459926', '8.536791']

>>> ['%.7g' % x3.get_qty('T')[i,35,3] for i in range(15,20)]
['250.5066', '250.6046', '250.8997', '250.8997', '251.1924']
>>> ['%.7g' % x3.get_qty('u')[i,35,3] for i in range(15,20)]
['11.40322', '11.37764', '11.30059', '11.30059', '11.37745']


#- Test fc recalculation scheme on an x3 domain:

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

>>> x.has_qty('fc')
0
>>> x3a.has_qty('fc')
0

>>> xtmp = copy.deepcopy(x)
>>> xtmp.fc()
>>> x3a  = atmqty.transform_AQobj.transform_AQobj(xtmp, x3)
>>> xtmp.has_qty('fc')
1
>>> x3a.has_qty('fc')
1
>>> ['%.7g' % xtmp.lat[i] for i in range(15,20)]
['-15', '-10', '-5', '0', '5']
>>> ['%.7g' % xtmp.get_qty('fc')[i,35,3] for i in range(15,20)]
['-3.774669e-05', '-2.53252e-05', '-1.271097e-05', '0', '1.271097e-05']
>>> ['%.7g' % x3a.lat[i] for i in range(15,20)]
['10', '12', '14', '16', '18']
>>> ['%.7g' % x3a.get_qty('fc')[i,35,3] for i in range(15,20)]
['2.53252e-05', '3.032226e-05', '3.528237e-05', '4.01995e-05', '4.506766e-05']


#- 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)

>>> ['%.7g' % x4.get_qty('T')[i,30,0] for i in range(15,20)]
['248.7729', '249.2471', '249.723', '250.1997', '250.6763']
>>> ['%.7g' % x4.get_qty('u')[i,30,0] for i in range(15,20)]
['1e+14', '1e+14', '1e+14', '11.15918', '11.30548']

>>> ['%.7g' % x4.get_qty('T')[i,30,2] for i in range(15,20)]
['250.1451', '250.5309', '250.9182', '251.3064', '251.6946']
>>> ['%.7g' % x4.get_qty('u')[i,30,2] for i in range(15,20)]
['11.83163', '11.89344', '12.07126', '12.31606', '12.56088']
"""}


#- Execute doctest if module is run from command line:

if __name__ == "__main__":
    """Test the module.

    Tests the examples in all the module documentation strings, plus
    __test__.

    Note:  To help ensure that module testing of this file works, the
    parent directory to the current directory is added to sys.path.
    """
    import doctest, sys, os
    sys.path.append(os.pardir)
    doctest.testmod(sys.modules[__name__])




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