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

"""Single-function module.

   See function docstring for description.
"""

#-----------------------------------------------------------------------
#                       Additional Documentation
#
# RCS Revision Code:
#   $Id: theta.py,v 1.2 2004/04/06 20:07:49 jlin Exp $
#
# Modification History:
# - 08 Oct 2003:  Original by Johnny Lin, Computation Institute,
#   University of Chicago.  Email:  air_jlin@yahoo.com.  Passed
#   passably reasonable tests.
# - 17 Nov 2003:  Added ability to handle missing values.  Passed
#   passably reasonable tests.
#
# Notes:
# - Written for Python 2.2.2.
# - Module docstrings can be tested using the doctest module.  To
#   test, execute "python theta.py".
# - Non-"built-in" packages and modules required:  atmqty
#
# Copyright (c) 2003 by Johnny Lin.  For licensing, distribution 
# conditions, contact information, and additional documentation see
# the URL http://www.johnny-lin.com/py_pkgs/atmqty/.
#=======================================================================




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

#- Set module version number to package version number:

import package_version
__version__ = package_version.version
del package_version




#---------------------- General Function:  theta -----------------------

def theta(p_in, T_in, p0=1000.0, missing=1e+20):
    """Calculate potential temperature.

    The reference pressure is set to 1000 hPa by default.  To use a 
    different reference pressure, set p0 to a different value when 
    calling the method.

    Method Arguments:
    * p_in:  Pressure [hPa].  Numeric floating point array.  
      Required.
    * T_in:  Temperature [K].  Numeric floating point array.  
      Required.
    * p0:  Reference pressure [hPa].  Floating point scalar.  
      Optional.
    * missing:  If p_in and/or T_in has missing values, this is the 
      missing value value.  Floating point scalar.  Default is 1e+20.
      Optional.

    Arguments p_in and T_in can be Numeric arrays of any number of 
    dimensions, as long as they are of the same shape and size.

    Output:
    * Potential temperature [K].  Numeric floating point array of 
      same dimensions and size as arguments.  If there are any
      missing values in output, those values are set to the value
      in argument missing from the input.  If there are missing 
      values in the output due to math errors and missing is set 
      to None, output will fill those missing values with the MA 
      default value of 1e+20.

    Reference:
    * Wallace, J. M., and P. V. Hobbs (1977): Atmospheric Science:  
      An Introductory Survey.  San Diego, CA:  Academic Press, ISBN 
      0-12-732950-1, p 69.

    Example without missing values:
    >>> import theta
    >>> import Numeric as N
    >>> p = N.array([1000., 1010., 1008.2])
    >>> T = N.array([273.1, 280.2, 278.4])
    >>> a = theta.theta(p, T)
    >>> ['%.8g' % a[i] for i in range(len(a))]
    ['273.1', '279.404', '277.75073']

    Example with missing values:
    >>> import theta
    >>> import Numeric as N
    >>> p = N.array([1000., 1010., 1008.2])
    >>> T = N.array([273.1, 1e20, 278.4])
    >>> a = theta.theta(p, T, missing=1e20)
    >>> ['%.8g' % a[i] for i in range(len(a))]
    ['273.1', '1e+20', '277.75073']
    """
    import atmqty
    import MA
    import Numeric as N
    from atmconst import AtmConst
    from is_numeric_float import is_numeric_float


    #- Do some error checks on input:

    if is_numeric_float(p_in, T_in) != 1:
        raise TypeError, "theta:  Args not Numeric floating"
    if is_numeric_float(N.array(p0)) != 1:
        raise TypeError, "theta:  p0 not floating"
    if missing != None:
        if is_numeric_float(N.array(missing)) != 1:
            raise TypeError, "theta:  missing not floating"


    #- Change input to MA:  just set to input value unless there
    #  are missing values, in which case add mask:
    
    if missing == None:
        T = MA.masked_array(T_in)
        p = MA.masked_array(p_in)
    else:
        T = MA.masked_values(T_in, missing, copy=0)
        p = MA.masked_values(p_in, missing, copy=0)


    #- Calculate and return potential temperature:

    const = AtmConst()
    return MA.filled( T * (p0 / p)**(const.kappa), missing )




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

#- Define additional examples for doctest to use:

__test__ = { 'Additional Example 1':
    """
    >>> import theta
    >>> import Numeric as N
    >>> p = N.array([400., 410., 308.2])
    >>> T = N.array([173.1, 180.2, 178.4])
    >>> a = theta.theta(p, T, p0=1010.)
    >>> ['%.8g' % a[i] for i in range(len(a))]
    ['225.58289', '233.18352', '250.48191']
    >>> a = theta.theta(p, T, p0=1010)
    Traceback (most recent call last):
        ...
    TypeError: theta:  p0 not floating
    >>> p = [400., 410., 308.2]
    >>> a = theta.theta(p, T)
    Traceback (most recent call last):
        ...
    TypeError: theta:  Args not Numeric floating
    """ }


#- 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:  On at least 2 platforms the doctest module does not execute
    the commands in $HOME/.pythonrc.py when "python theta.py" is
    executed at the command line.  Since atmqty might not be included
    on the default Python path, specifying it in .pythonrc.py is one
    way to insure Python can access the package in an interactive
    session.  If doctest does not execute .pythonrc.py, any call to 
    atmqty will fail.
    
    As the theta module requires the atmqty package (in order to 
    access the AtmConst class) and theta is distributed as part of the 
    atmqty package, the parent directory to the atmqty directory is 
    added to sys.path for this module testing case to enable doctest
    to work on theta.
    """
    import doctest, sys, os
    sys.path.append(os.pardir)
    doctest.testmod(sys.modules[__name__])




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