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

"""Define class StateSet for model State objects.

   The StateSet class is a container that is a collection of State
   class objects.
"""

#-----------------------------------------------------------------------
#                       Additional Documentation
#
# RCS Revision Code:
#   $Id: stateset.py,v 1.1.1.1 2005/01/13 00:15:42 jlin Exp $
#
# Modification History:
# - 27 Aug 2004:  Original by Johnny Lin, Computation Institute,
#   University of Chicago.  Passed passably reasonable tests.
#
# Notes:
# - Written for Python 2.2.
# - See import statements throughout module for dependencies.
#
# Copyright (c) 2004 by Johnny Lin.  For licensing, distribution 
# conditions, contact information, and additional documentation see
# the URL http://geosci.uchicago.edu/csc/modelutil/doc/.
#=======================================================================




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

#- Set module version to package version:

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


#- If you're importing this module in testing mode, or you're running 
#  pydoc on this module via the command line, import user-specific 
#  settings to make sure any non-standard libraries are found:

import os, sys
if (__name__ == "__main__") or \
   ("pydoc" in os.path.basename(sys.argv[0])):
    import user
del os, sys


#- Import Numeric/numarray as appropriate (see gemath.num_settings 
#  module for details as to what gets imported).  Import copy:

from gemath.num_settings import *
import copy


#- Create module variable _typeState which is the type() of objects
#  of class State:

from field import Field
from state import State
_f1 = Field(0, id='a')
_f2 = Field(0, id='b')
_typeState = type(State(_f1, _f2))
del _f1, _f2
del Field, State




#-------- Module Functions:  Basic State id Manipulation Methods -------

def stateid2int(inlist):
    """Returns non-"tendency" State ids list inlist as list of integers.

    In the returned list, "tm*" ids become negative integers and "tp*" 
    ids become positive.  "t0" becomes 0.  The magnitude of the integer 
    is equal to the numeral(s) after the "tm" or "tp" prefix.  There 
    can be no "tend*" terms or any other named values in inlist or an 
    exception is thrown.  The returned list does not share memory with 
    inlist, but the order of the list is the same as the input list.

    Positional Input Parameter:
    * inlist:  List of State ids.  Usually this list is the keys of a 
      StateSet object.  List of strings.

    Examples:
    >>> a = ['tp1', 't0', 'tm2']
    >>> print stateid2int(a)
    [1, 0, -2]
    >>> a = ['tp1', 't0', 'tm2', 'tend1']
    >>> print stateid2int(a)
    Traceback (most recent call last):
        ...
    ValueError: Bad keys
    """
    id_list_num = []
    for anid in inlist:
        if   anid.startswith('tm'):
            id_list_num.append( -int(anid[2:]) )
        elif anid.startswith('tp'):
            id_list_num.append(  int(anid[2:]) )
        elif anid == 't0':
            id_list_num.append( 0 )
        else:
            raise ValueError, "Bad keys"
    return id_list_num


def int2stateid(inlist):
    """Returns list of integers as list of non-"tendency" State ids.

    In the returned list, negative integers become "tm" State ids and 
    positive integers become "tp" ids.  0 becomes "t0".  The numeral(s)
    after the "tm" or "tp" prefix is converted from the magnitude of 
    the integer.  The returned list does not share memory with inlist, 
    but the order of the list is the same as the input list.

    Positional Input Parameter:
    * inlist:  List of integers.  Function does not test that this is
      the case, but use of different values will likely give an unin-
      telligible result or an exception.
    
    Examples:
    >>> a = [1, 0, -2]
    >>> print int2stateid(a)
    ['tp1', 't0', 'tm2']
    >>> a = [1, 0, -2, None]
    >>> print int2stateid(a)
    Traceback (most recent call last):
        ...
    TypeError: bad operand type for abs()
    """
    id_list = []
    for anum in inlist:
        if   anum < 0:
            id_list.append( 'tm' + str(abs(anum)) )
        elif anum > 0:
            id_list.append( 'tp' + str(abs(anum)) )
        elif anum == 0:
            id_list.append( 't0' )
        else:
            raise ValueError, "Bad keys"
    return id_list




#----------------------- StateSet Class:  Header -----------------------

class StateSet(object):
    """Collection of State objects at different ids.
    
    The StateSet class is a container that is a collection of State
    class objects.  The StateSet keys for mapping the State objects 
    are the id attributes in the State objects; the keys in a State-
    Set object are unique, and thus each StateSet object can have 
    only as many State objects as there are unique ids.

    Public Instance Attributes:
    * delt:  Timestep (i.e. the difference between the time for States
      at "tp1" and "t0"), in units time_units (of the State variables).
      Default is None.
    * id:  Name of the object.  The name contains information about the 
      set of states.  There are no standardized set of values that this
      attribute has to conform to.  String.  Default is None.

    All the instance attributes can be set directly or via keywords
    passed in the instantiation parameter list.  See the __init__ func-
    tion docstring for details.

    Example:
    >>> from field import Field
    >>> from state import State
    >>> data1 = N.array([1.1, 3.0, -4.8, 2.9, 2.1])
    >>> data2 = N.array([31.1, 0.4, 4.0, -4.9, 9.1])
    >>> data3 = N.array([3.1, 2.4, 4.9, 8.1, -9.2])
    >>> f1a = Field(data1, id='u', units='m/s', axis=['Latitude'])
    >>> f2a = Field(data2, id='v', units='m/s', axis=['Latitude'])
    >>> f3a = Field(data3, id='w', units='m/s', axis=['Latitude'])
    >>> f1b = Field(data1*0.1, id='u', units='m/s', axis=['Latitude'])
    >>> f2b = Field(data2*0.2, id='v', units='m/s', axis=['Latitude'])
    >>> s1 = State(f1a, f2a, f3a, long_name='winds', id='t0')
    >>> s2 = State(f1b, f2b, long_name='horiz_winds', id='tm1')
    >>> set = StateSet(s1, s2, id='wind_collection')

    >>> print set.id
    wind_collection
    >>> print set.delt
    None
    >>> print set.keys()
    ['tm1', 't0']

    >>> print type(set['t0'])
    <class 'state.State'>
    >>> print set['tm1'].keys()
    ['u', 'v']
    >>> print set['tm1']['v'].units
    m/s
    """




#---------------- StateSet Class:  Initialization Method ---------------

    def __init__(self, *args, **kwds):
        """Initialize class object.

        Positional Arguments:
        * All arguments of this initialization function are assumed to
          be State objects which are put into the container's data
          structure, private attribute _data.  If any of the arguments 
          have the same id attribute, only the last State object in
          the argument list with that id will be stored; the rest will 
          be ignored.

        Keyword Arguments:
        * All keywords are attributes of the instance.  Any other key-
          words in the calling line are ignored.  See the class doc-
          string for details.

        The private variable _data is the data structure that holds
        all the States.  The keys to _data (accessible via the public
        function keys) are the State.id attributes and the values of
        _data (accessible via the public function values) are State
        objects.
        """
        self._ok_init_kwds = ('delt', 'id')
        self._data = {}     #- Note:  _data is just a dictionary.

        for an_arg in args:
            self._data[an_arg.id] = an_arg

        for akey in self._ok_init_kwds:
            if kwds.has_key(akey):
                setattr(self, akey, kwds[akey])
            else:
                setattr(self, akey, None)




#----------------- StateSet Class:  Container Methods ------------------

    def __contains__(self, item):
        return self._data.has_key(item)

    def __delitem__(self, key):
        del self._data[key]

    def __getitem__(self, key):
        return self._data[key]

    def __iter__(self):
        """Returns an iterkeys equivalent iterator over StateSet data."""
        return self._data.iterkeys()

    def __len__(self):
        return len(self._data)

    def __setitem__(self, key, value):
        """Method to add/rebind items.

        Raises a KeyError exception if key is not equal to value.id.
        """
        if key != value.id:
            raise KeyError, "Key/value.id mismatch"
        self._data[key] = value




#------------------ StateSet Class:  Mapping Methods -------------------

    def clear(self):
        """Clears all State objects in the StateSet object."""
        self._data.clear()

    def has_key(self, key):
        """Returns True if State object of id key is in StateSet object."""
        return self._data.has_key(key)

    def items(self):
        """Returns copy of list of all State items in StateSet object.
        
        The copy is a shallow copy (I think) and thus references to
        other objects are not copied recursively.
        """
        return self._data.items()

    def keys(self):
        """Returns copy of list of all State object ids in StateSet object.
        
        The copy is a shallow copy (I think) and thus references to
        other objects are not copied recursively.
        """
        return self._data.keys()

    def values(self):
        """Returns copy of list of all State objects in StateSet object.
        
        The copy is a shallow copy (I think) and thus references to
        other objects are not copied recursively.
        """
        return self._data.values()




#--------------------- StateSet Class:  add Method ---------------------

    def add(self, *args):
        """Add/overwrite State objects given in args to StateSet object.

        Positional Arguments:
        * All arguments are State objects and are added to the StateSet
          object self in the order they are given in the argument list.
          If any of the arguments have id attributes equal to an exist-
          ing key in self, the key is detached from its previous value 
          and associated with the new value given in the argument.
        """
        #- Check list of arguments does not contain States with dupli-
        #  cate ids:
    
        list_ids = [anarg.id for anarg in args]
        for anid in list_ids:
            if list_ids.count(anid) != 1:
                raise ValueError, "State list has duplicate ids"


        #- Add State objects:

        for aState in args:
            if type(aState) != _typeState:
                raise TypeError, "Adding a non-State object"
            else:
                self._data[aState.id] = aState




#---------------- StateSet Class:  Various Copy Methods ----------------

    def copy(self):
        """Returns deepcopy of StateSet object."""
        return copy.deepcopy(self)


    def copymeta(self):
        """Returns deepcopy of StateSet object metadata.

        The entire StateSet object is copied, except for the States
        held in the StateSet container.  Thus, the returned StateSet
        object contains no data or State metadata, but rather only
        StateSet metadata.
        """
        tmp = StateSet()
        for akey in self.__dict__.keys():
            if akey != '_data':
                tmp.__dict__[akey] = copy.deepcopy(self.__dict__[akey])
            else:
                pass
        return tmp




#----------- StateSet Class:  Method to Remove Certain States ----------

    def del_earliest(self):
        """Removes the earliest State object from StateSet container.

        Earliest is defined in terms of relative timestep.  Only non-
        "tendency" State ids (i.e. those with prefix "tm", "tp", or 
        value of "t0") in the StateSet object are considered.  All
        other State ids, valid or not (e.g. "tendency" State ids like 
        "tend1") are ignored.
        
        Examples of the Idea of How the Method Is Used and Works:
           set = StateSet(s1, s2, s3)
           set.keys() -> ['tp1', 't0', 'tm2']
           set.del_earliest() 
           set.keys() -> ['tp1', 't0']

           set = StateSet(s1, s2, s3)
           set.keys() -> ['tend1', 't0', 'tm2']
           set.del_earliest() 
           set.keys() -> ['tend1', 't0']
        """
        sorted_stateid = self.sort_stateid()
        del self[sorted_stateid[0]]


    def del_latest(self):
        """Removes the latest State object from StateSet container.

        Latest is defined in terms of relative timestep.  Only non-
        "tendency" State ids (i.e. those with prefix "tm", "tp", or 
        value of "t0") in the StateSet object are considered.  All
        other State ids, valid or not (e.g. "tendency" State ids like 
        "tend1") are ignored.
        
        Examples of the Idea of How the Method Is Used and Works:
           set = StateSet(s1, s2, s3)
           set.keys() -> ['tp1', 't0', 'tm2']
           set.del_latest() 
           set.keys() -> ['tm2', 't0']

           set = StateSet(s1, s2, s3)
           set.keys() -> ['tend1', 't0', 'tm2']
           set.del_latest() 
           set.keys() -> ['tend1', 'tm2']
        """
        sorted_stateid = self.sort_stateid()
        del self[sorted_stateid[-1]]




#------------- StateSet Class:  meta_ok_to_interface Method ------------

    def meta_ok_to_interface(self, IMobj, check_long_name=False):
        """Returns True if metadata passes checks for interfaces.

        Tests whether the Field objects metadata in the State objects
        that make up this StateSet object pass consistency checks in 
        accordance with the description given in InterfaceMeta object 
        IMobj.  If all the Field objects pass, this method returns 
        True.  See the docstring for the Field class for details as to 
        what the consistency checks entail.

        Positional Input Argument:
        * IMobj:  InterfaceMeta object that defines the accepted meta-
          data that you wish to check this State variable's Field var-
          iable's metadata against.  For details see:

             http://geosci.uchicago.edu/csc/modelutil/doc/imeta.html

        Keyword Input Arguments:
        * check_long_name:  If set to True, the method not only checks 
          units but also long_name.  Default is False.
        """
        for aobjkey in self.keys():
            aobj = self[aobjkey]
            if not aobj.meta_ok_to_interface( IMobj \
                 , check_long_name=check_long_name ):  return False
        return True




#------------- StateSet Class:  shift_time_stateid Method --------------

    def shift_time_stateid(self, nstep):
        """Shift State objects nstep timesteps; delete "tendencies".

        The State objects in the StateSet container are all shifted
        forward or backward the number of time steps given in nstep.
        For example, if nstep=-1:  the State object with id "t0" has
        its id changed to "tm1" and the value of its time attribute 
        is decreased by delt; the State object with id "tm1" has its
        id changed to "tm2", etc.  "Tendency" State variables (e.g. 
        id with prefix "tend") are deleted from the StateSet, as are 
        State objects with None as key.

        If either self.delt is None or the time attribute of a State
        variable is None, the time attribute of that State variable
        after the shift is set to None (if the time attribute does
        not exist for that State variable pre-shift, it remains
        undefined).

        Positional Input Parameter:
        * nstep:  Number of timesteps to move all State objects for-
          ward or backward.  Integer.  Required input, with no default 
          value.

        Since shifting timesteps requires that all the timesteps in
        the StateSet are consistent with one another, this method also
        executes the time_ok() method on the post-shift States and 
        throws an exception if it returns False.

        Example:
        >>> from field import Field
        >>> from state import State
        >>> data1 = N.array([1.1, 3.0, -4.8, 2.9, 2.1])
        >>> data2 = N.array([31.1, 0.4, 4.0, -4.9, 9.1])
        >>> data3 = N.array([3.1, 2.4, 4.9, 8.1, -9.2])
        >>> f1a = Field(data1, id='u', units='m/s')
        >>> f2a = Field(data2, id='v', units='m/s')
        >>> f3a = Field(data3, id='w', units='m/s')
        >>> s1 = State(f3a, id='tend1')
        >>> s2 = State(f1a, f2a, id='t0')
        >>> s3 = State(f2a, id='tm1')
        >>> s = StateSet(s1, s2, s3, id='velocities')
        >>> s.keys()
        ['tm1', 'tend1', 't0']
        >>> s.shift_time_stateid(-1)
        >>> s.keys()
        ['tm1', 'tm2']
        """
        nstep_use = int(nstep)

        for anid in self.keys():           #- Remove "tendencies" and
            if anid == None:               #  id=None State objects
                del self[anid]
            elif anid.startswith('tend'):
                del self[anid]


        #- Calculate list of ids pre- and post-shift.  The lists are
        #  in the same order so corresponding elements refer to the
        #  same object:

        old_id_list = self.keys()
        old_id_list_num = self.stateid2int()
        new_id_list_num = N.array(old_id_list_num) + nstep_use
        new_id_list = int2stateid(new_id_list_num.tolist())


        #- Replace ids in StateSet:

        tmp_old_state_list = [ self[id] for id in old_id_list ]
        self.clear()
        for i in range(len(old_id_list)):
            tmp_state = tmp_old_state_list[i]
            tmp_state.id = new_id_list[i]
            self.add( tmp_state )
            del tmp_state
        del tmp_old_state_list


        for aStatekey in self.keys():     #- Shift State.time attribute
            aState = self[aStatekey]
            if hasattr(aState, 'time'):
                if (self.delt is None) or (aState.time is None):
                    aState.time = None
                else:
                    aState.time = aState.time + (nstep_use * self.delt)


        #- Check time metadata post-shift:

        if not self.time_ok():
            raise ValueError, "Inconsistent time metadata"




#---------------- StateSet Class:  sort_stateid Method -----------------

    def sort_stateid(self):
        """Returns a list of the State ids in time-ascending order.

        Only non-"tendency" State ids (i.e. those with prefix "tm",
        "tp", or value of "t0") in the StateSet object are sorted and 
        returned.  All other State ids, valid or not (e.g. "tendency" 
        State ids like "tend1") are ignored.

        Examples of the Idea of How the Method Is Used and Works:
           set = StateSet(s1, s2, s3)
           set.keys() -> ['tp1', 't0', 'tm2']
           set.sort_stateid() -> ['tm2', 't0', 'tp1']

           set = StateSet(s1, s2, s3, s4)
           set.keys() -> ['tp1', 't0', 'tm2', 'tend1']
           set.sort_stateid() -> ['tm2', 't0', 'tp1']
        """
        id_list_non_tend = []
        for anid in self.keys():
            if anid.startswith('tm') or anid.startswith('tp') or \
               anid.startswith('t0'):
                id_list_non_tend.append(anid)
            else:
                pass
        id_list_num = stateid2int(id_list_non_tend)
        id_list_num.sort()
        return int2stateid(id_list_num)




#----------------- StateSet Class:  stateid2int Methods ----------------

    def stateid2int(self):
        """Returns list of non-"tendency" State ids as integers.

        The source list of State ids is the keys in the present State-
        Set object.  In the returned list, "tm" ids are negative and 
        "tp" ids are positive.  "t0" is 0.  Any "tend*" terms in the 
        State id list are ignored.  The returned list does not share
        memory with the original keys list, but the order of the list
        is the same as the input keys.

        Example of the Idea of How the Method Is Used and Works:
           set = StateSet(s1, s2, s3)
           set.keys() -> ['tp1', 't0', 'tm2']
           set.stateid2int() -> [1, 0, -2]
        """
        return stateid2int(self.keys())




#------------------- StateSet Class:  time_ok Method -------------------

    def time_ok(self):
        """Returns True if passes consistency checks for time parameters.

        Examines the time parameters for the StateSet container as well
        as the constituent State objects to see that they are consis-
        tent with each other.  The StateSet container can only hold
        non-"tendency" State objects (e.g. the presence of State ob-
        jects with id prefixed by "tend" will throw an exception). 

        Method makes the following consistency checks:
        - Each of the State variables have the same time_units (or un-
          defined or equal to None).
        - Values of time, delt, and id are consistent with each other.
          If the "t0" State id does not exist, the "t0" object does 
          not have the time attribute (or it equals None), or the delt 
          attribute does not exist or equals None, this check is igno-
          red and True is returned.
        
        In the consistency checks, if an attribute is undefined or set
        to None, the value is assumed to match any value.  If no checks 
        are done, default of the method is to return True.
        """
        #- Establish list of all State ids and integer "equivalents":

        id_list = self.keys()
        id_list_num = self.stateid2int()


        #- Check each State variable has same time_units:

        first_id_to_check = True
        for aStatekey in self.keys():
            aState = self[aStatekey]
            if hasattr(aState, 'time_units') and (aState.time_units != None):
                if first_id_to_check:
                    time_units_to_compare = aState.time_units
                    first_id_to_check = False
                    continue
                if aState.time_units != time_units_to_compare:
                    return False


        #- Check time, delt, and id are consistent with each other:

        if not hasattr(self, 'delt'):        return True
        if self.delt == None:                return True
        if "t0" not in id_list:              return True
        if not hasattr(self["t0"], 'time'):  return True
        if self["t0"].time == None:          return True

        for aStatekey in self.keys():
            aState = self[aStatekey]
            idx = id_list.index(aState.id)
            if hasattr(aState, 'time') and (aState.time != None):
                test_time = ( id_list_num[idx] * self.delt ) + self["t0"].time
                if aState.time != test_time:  return False

        return True




#----------------- StateSet Class:  Additional Methods -----------------

    def refdata(self):
        """Returns reference to data of StateSet object.
        
        This returns a reference to the private attribute _data, which
        is a dictionary.
        """
        return self._data


    def setdata(self, data):
        """Reference StateSet private _data attribute to data.
        
        This makes a reference of data to the private attribute _data, 
        which is a dictionary.  Thus, argument data must be a diction-
        ary.
        """
        self._data = data


    def state_list(self):
        """Returns summary list of State objects in StateSet object.

        A list of tuples is returned, each tuple describing one of the
        State objects the StateSet object is a container for.  Each
        tuple holds the string values of the following State attributes,
        in this form:  (id, long_name).  If either the id or long_name
        attributes are empty, an empty string is returned.

        Example of the Idea of How the Method Is Used and Works:
           s1 = State(f1a, f2a, f3a, long_name='3-D_winds', id='t0')
           s2 = State(f1b, f2b, long_name='2-D_winds')
           st = StateSet(s1, s2, id='velocities')
           st.state_list() -> [(None, '2-D_winds'), ('t0', '3-D_winds')]
        """
        return [(akey, self[akey].long_name) for akey in self.keys()]


    def wrap(self):
        """Connect the MPI parallel subdomains of all State variables.

        Runs through all State variables and executes the wrap method
        for the State variable.
        """
        for aState in self.values():  aState.wrap()




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

__test__ = {
'Additional Examples and Tests':
"""
#- Initialize a StateSet object:

>>> from field import Field
>>> from state import State
>>> data1 = N.array([1.1, 3.0, -4.8, 2.9, 2.1])
>>> data2 = N.array([31.1, 0.4, 4.0, -4.9, 9.1])
>>> data3 = N.array([3.1, 2.4, 4.9, 8.1, -9.2])
>>> f1a = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2a = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3a = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> f1b = Field(data1*0.1, id='u', units='m/s', axis=['Latitude'])
>>> f2b = Field(data2*0.2, id='v', units='m/s', axis=['Latitude'])
>>> s1 = State(f1a, f2a, f3a, long_name='3-D_winds', id='t0')
>>> s2 = State(f1b, f2b, long_name='2-D_winds')
>>> s = StateSet(s1, s2, id='velocities')

>>> print s.id
velocities


#- Test container methods:

>>> print 't0' in s
True
>>> print None in s
True
>>> print 'u' in s
False

>>> print s.keys()
[None, 't0']
>>> del s['t0']
>>> print s.keys()
[None]
>>> del s['tp1']
Traceback (most recent call last):
    ...
KeyError: 'tp1'

>>> print type(s[None])
<class 'state.State'>
>>> print s[None].long_name
2-D_winds

>>> [a for a in s]
[None]

>>> len(s)
1

>>> s['new_winds'] = s2
Traceback (most recent call last):
    ...
KeyError: 'Key/value.id mismatch'
>>> s['t0'] = s1
>>> print s.keys()
[None, 't0']
>>> print s['tp1'].long_name
Traceback (most recent call last):
    ...
KeyError: 'tp1'
>>> ['%.6g' % s['t0']['w'][i] for i in range(len(s['t0']['w']))]
['3.1', '2.4', '4.9', '8.1', '-9.2']


#- Test mapping methods:

>>> print s.keys()
[None, 't0']
>>> s.clear()
>>> print s.keys()
[]
>>> print s._data
{}

>>> st = StateSet(s1, s2, id='velocities')
>>> print st.has_key('u')
False
>>> print st.has_key(None)
True

>>> items = st.items()
>>> print items[0][0]
None
>>> print type(items[0][1])
<class 'state.State'>

>>> print st.keys()
[None, 't0']

>>> print type(st.values()[1])
<class 'state.State'>


#- Test add method:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2 = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3 = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> sset = StateSet(s1, s2)
>>> sset.keys()
['tm1', 't0']
>>> ['%.6g' % sset['t0']['u'][i] for i in range(len(sset['t0']['u']))]
['1.1', '3', '-4.8', '2.9', '2.1']
>>> ['%.6g' % sset['tm1']['u'][i] for i in range(len(sset['tm1']['u']))]
Traceback (most recent call last):
    ...
KeyError: 'u'
>>> sset.add(s3, s4)
>>> ['%.6g' % sset['tm2']['u'][i] for i in range(len(sset['tm2']['u']))]
Traceback (most recent call last):
    ...
KeyError: 'u'
>>> ['%.6g' % sset['tm2']['w'][i] for i in range(len(sset['tm2']['w']))]
['3.1', '2.4', '4.9', '8.1', '-9.2']
>>> ['%.6g' % sset['tm3']['w'][i] for i in range(len(sset['tm3']['w']))]
['3.1', '2.4', '4.9', '8.1', '-9.2']
>>> sset.add(f1)
Traceback (most recent call last):
    ...
TypeError: Adding a non-State object


#- Test copy and copymeta methods:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', long_name='f1')
>>> f2 = Field(data2, id='v', units='m/s', long_name='f2')
>>> f3 = Field(data3, id='w', units='m/s', long_name='f3')
>>> s1 = State(f1, id='t0', long_name='s1')
>>> s2 = State(f2, id='tm1', long_name='s2')
>>> s3 = State(f3, id='tm2', long_name='s3')
>>> s4 = State(f2, f3, id='tp2', long_name='s4')
>>> sset = StateSet(s1, s2, s3, s4, id='modelstuff')
>>> ['%.6g' % sset['t0']['u'][i] for i in range(len(sset['t0']['u']))]
['1.1', '3', '-4.8', '2.9', '2.1']
>>> news = sset.copy()
>>> ['%.6g' % news['t0']['u'][i] for i in range(len(news['t0']['u']))]
['1.1', '3', '-4.8', '2.9', '2.1']
>>> sset['t0']['u'][1] = -999.
>>> news['t0']['u'][3] = -888.
>>> ['%.6g' % sset['t0']['u'][i] for i in range(len(sset['t0']['u']))]
['1.1', '-999', '-4.8', '2.9', '2.1']
>>> ['%.6g' % news['t0']['u'][i] for i in range(len(news['t0']['u']))]
['1.1', '3', '-4.8', '-888', '2.1']
>>> newer = sset.copymeta()
>>> print newer.id
modelstuff
>>> newer.id = 'newerstuff'
>>> print newer.id
newerstuff
>>> print news.id
modelstuff
>>> print newer._data
{}
>>> print newer.delt
None


#- Test del_earliest and del_latest methods:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', long_name='f1')
>>> f2 = Field(data2, id='v', units='m/s', long_name='f2')
>>> f3 = Field(data3, id='w', units='m/s', long_name='f3')
>>> s1 = State(f1, id='t0', long_name='s1')
>>> s2 = State(f2, id='tm1', long_name='s2')
>>> s3 = State(f3, id='tm2', long_name='s3')
>>> s4 = State(f2, f3, id='tp2', long_name='s4')
>>> sset = StateSet(s1, s2, s3, s4)
>>> sset.keys()
['tm1', 'tm2', 'tp2', 't0']
>>> print sset['tm1'].long_name
s2
>>> sset.del_earliest()
>>> sset.del_latest()
>>> sset.keys()
['tm1', 't0']
>>> print sset['tm1'].long_name
s2

>>> s1 = State(f1, id='t0', long_name='s1')
>>> s2 = State(f2, id='tend1', long_name='s2')
>>> s3 = State(f3, id='tm2', long_name='s3')
>>> s4 = State(f2, f3, id='tp2', long_name='s4')
>>> sset = StateSet(s1, s2, s3, s4)
>>> sset.keys()
['tm2', 'tend1', 't0', 'tp2']
>>> print sset['t0'].long_name
s1
>>> sset.del_earliest()
>>> sset.del_latest()
>>> sset.keys()
['tend1', 't0']
>>> print sset['t0'].long_name
s1


#- Test shift_time_stateid method:

    * Test shifting does not change the values (pos. and neg. shifts):

>>> data1 = N.array([1.1, 3.0, -4.8, 2.9, 2.1])
>>> data2 = N.array([31.1, 0.4, 4.0, -4.9, 9.1])
>>> data3 = N.array([3.1, 2.4, 4.9, 8.1, -9.2])
>>> f1a = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2a = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3a = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> f1b = Field(data1*0.1, id='u', units='m/s', axis=['Latitude'])
>>> f2b = Field(data2*0.2, id='v', units='m/s', axis=['Latitude'])
>>> s1 = State(f3a, long_name='winds1', id='tp1')
>>> s2 = State(f1a, f2a, f3a, long_name='winds2', id='t0')
>>> s3 = State(f1b, f2b, long_name='winds3', id='tm1')
>>> s4 = State(f1b, f3a, long_name='winds4', id='tm2')
>>> s5 = State(f2a, f2b, long_name='winds5', id='tend1')
>>> s6 = State(f2a, f2b, long_name='winds6', id='tm23')
>>> s7 = State(f2a, f2b, long_name='winds7', id='tp12')
>>> s = StateSet(s1, s2, s3, s4, s5, s6, s7, id='velocities')
>>> s.keys()
['tm1', 'tm2', 'tp12', 't0', 'tm23', 'tend1', 'tp1']
>>> print s['tm2'].long_name
winds4
>>> ['%.6g' % s['tm2']['u'][i] for i in range(len(s['tm2']['u']))]
['0.11', '0.3', '-0.48', '0.29', '0.21']

>>> s.shift_time_stateid(-1)
>>> s.keys()
['tm1', 'tm2', 'tm3', 'tp11', 't0', 'tm24']
>>> print s['tm3'].long_name
winds4
>>> ['%.6g' % s['tm3']['u'][i] for i in range(len(s['tm3']['u']))]
['0.11', '0.3', '-0.48', '0.29', '0.21']
>>> ['%.6g' % s['tm3']['v'][i] for i in range(len(s['tm3']['v']))]
Traceback (most recent call last):
    ...
KeyError: 'v'
>>> ['%.6g' % s['tm3']['w'][i] for i in range(len(s['tm3']['w']))]
['3.1', '2.4', '4.9', '8.1', '-9.2']

>>> print s['tm23'].long_name
Traceback (most recent call last):
    ...
KeyError: 'tm23'
>>> print s['tm24'].long_name
winds6
>>> ['%.6g' % s['tm24']['u'][i] for i in range(len(s['tm24']['u']))]
Traceback (most recent call last):
    ...
KeyError: 'u'
>>> ['%.6g' % s['tm24']['v'][i] for i in range(len(s['tm24']['v']))]
['6.22', '0.08', '0.8', '-0.98', '1.82']

>>> s.shift_time_stateid(-2)
>>> s.keys()
['tm2', 'tm3', 'tm4', 'tm5', 'tm26', 'tp9']
>>> ['%.6g' % s['tm26']['v'][i] for i in range(len(s['tm26']['v']))]
['6.22', '0.08', '0.8', '-0.98', '1.82']
>>> ['%.6g' % s['tm5']['u'][i] for i in range(len(s['tm5']['u']))]
['0.11', '0.3', '-0.48', '0.29', '0.21']

>>> s.shift_time_stateid(4)
>>> s.keys()
['tm1', 'tp13', 't0', 'tp2', 'tp1', 'tm22']
>>> ['%.6g' % s['tm1']['u'][i] for i in range(len(s['tm1']['u']))]
['0.11', '0.3', '-0.48', '0.29', '0.21']

    * Check time attribute correctly changed:

>>> s1 = State(f3a, long_name='winds1', id='tp1', time=10)
>>> s2 = State(f1a, f2a, f3a, long_name='winds2', id='t0', time=0)
>>> s3 = State(f1b, f2b, long_name='winds3', id='tm1', time=-10)
>>> s5 = State(f2a, f2b, long_name='winds5', id='tend1', time=0)
>>> s = StateSet(s1, s2, s3, s5, d='velocities', delt=10)
>>> s.keys()
['tm1', 't0', 'tp1', 'tend1']
>>> print s['t0'].time
0
>>> print s['t0'].time_units
None
>>> print s['tm1'].time
-10
>>> print s['tp1'].time
10
>>> print s['tp1'].time_units
None
>>> s.shift_time_stateid(-1)
>>> s.keys()
['tm1', 'tm2', 't0']
>>> print s['t0'].time
0
>>> print s['t0'].time_units
None
>>> print s['tm1'].time
-10
>>> print s['tp1'].time
Traceback (most recent call last):
    ...
KeyError: 'tp1'

    * Case where the time units are not ok:

>>> s1 = State(f3a, long_name='winds1', id='tp1', time=100)
>>> s2 = State(f1a, f2a, f3a, long_name='winds2', id='t0', time=0)
>>> s3 = State(f1b, f2b, long_name='winds3', id='tm1', time=-10)
>>> s5 = State(f2a, f2b, long_name='winds5', id='tend1', time=0)
>>> s = StateSet(s1, s2, s3, s5, d='velocities', delt=10)
>>> s.shift_time_stateid(-1)
Traceback (most recent call last):
    ...
ValueError: Inconsistent time metadata


#- Test sort_stateid method:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2 = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3 = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tp2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.keys()
['tm1', 'tm3', 'tp2', 't0']
>>> set.sort_stateid()
['tm3', 'tm1', 't0', 'tp2']
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tend1')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.keys()
['tm1', 'tm3', 'tend1', 't0']
>>> set.sort_stateid()
['tm3', 'tm1', 't0']
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tp2')
>>> s3 = State(f3, id='tend1')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.keys()
['tm3', 'tp2', 't0', 'tend1']
>>> set.sort_stateid()
['tm3', 't0', 'tp2']
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='hello')
>>> s3 = State(f3, id='tend1')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.keys()
['tm3', 'hello', 't0', 'tend1']
>>> set.sort_stateid()
['tm3', 't0']


#- Test stateid2int method:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2 = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3 = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.keys()
['tm1', 'tm2', 'tm3', 't0']
>>> set.stateid2int()
[-1, -2, -3, 0]


#- Test time_ok method:

    * With no time related metadata explicitly set besides the State ids:

>>> data1 = [1.1, 3.0, -4.8, 2.9, 2.1]
>>> data2 = [31.1, 0.4, 4.0, -4.9, 9.1]
>>> data3 = [3.1, 2.4, 4.9, 8.1, -9.2]
>>> f1 = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2 = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3 = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='tend1')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
Traceback (most recent call last):
    ...
ValueError: Bad keys
>>> s1 = State(f1, id='t0')
>>> s2 = State(f2, id='hello')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
Traceback (most recent call last):
    ...
ValueError: Bad keys

    * Check State time_units consistency:

>>> s1 = State(f1, id='t0', time_units='s')
>>> s2 = State(f2, id='tm1')
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time_units='s')
>>> s2 = State(f2, id='tm1', time_units='s')
>>> s3 = State(f3, id='tm2', time_units='s')
>>> s4 = State(f2, f3, id='tm3', time_units='s')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time_units='s')
>>> s2 = State(f2, id='tm1', time_units='s')
>>> s3 = State(f3, id='tm2', time_units='sec')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> print s4.time_units
None
>>> set.time_ok()
False
>>> s1 = State(f1, id='t0', time_units='s')
>>> s2 = State(f2, id='tm1', time_units='s')
>>> s3 = State(f3, id='tm2', time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> del s4.__dict__['time_units']
>>> print hasattr(s4, 'time_units')
False
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time_units='hr')
>>> s2 = State(f2, id='tm1', time_units='s')
>>> s3 = State(f3, id='tm2', time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> del s4.__dict__['time_units']
>>> print hasattr(s4, 'time_units')
False
>>> set.time_ok()
False
>>> s3 = State(f3, id='tm2')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s3, s4)
>>> del s3.__dict__['time_units']
>>> del s4.__dict__['time_units']
>>> set.time_ok()
True

    * Check delt, id, and time are consistent:

>>> s2 = State(f2, id='tm1', time=0, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time=100, time_units='s')
>>> s2 = State(f2, id='tm1', time=0, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time=100000, time_units='s')
>>> s2 = State(f2, id='tm1', time=0, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time=100, time_units='s')
>>> s2 = State(f2, id='tm1', time=0, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4, delt=100)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time=100, time_units='s')
>>> s2 = State(f2, id='tm1', time=9, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4, delt=100)
>>> set.time_ok()
False

    * Check non-contiguous timesteps:

>>> s1 = State(f1, id='tm3', time=100, time_units='hr')
>>> s2 = State(f2, id='tp4', time=0, time_units='hr')
>>> s3 = State(f3, id='tm2', time=-100, time_units='hr')
>>> set = StateSet(s1, s2, s3, delt=100)
>>> set.time_ok()
True
>>> s1 = State(f1, id='tm3', time=100, time_units='hr')
>>> s2 = State(f2, id='tp4', time=0, time_units='hr')
>>> s3 = State(f3, id='tm2', time=-100, time_units='hrs')
>>> set = StateSet(s1, s2, s3)
>>> set.time_ok()
False
>>> s1 = State(f1, id='tm3', time=100, time_units='hr')
>>> s2 = State(f2, id='tp4', time=0, time_units='hr')
>>> s3 = State(f3, id='tm2', time=-100, time_units='hr')
>>> set = StateSet(s1, s2, s3)
>>> set.time_ok()
True
>>> s1 = State(f1, id='t0', time=100, time_units='s')
>>> s2 = State(f2, id='tp4', time=500, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm33')
>>> set = StateSet(s1, s2, s3, s4, delt=100)
>>> set.time_ok()
True

    * Check some combinations for consistency:

>>> s1 = State(f1, id='t0', time=100, time_units='s')
>>> s2 = State(f2, id='tm1', time=0, time_units='s GMT')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4, delt=100)
>>> set.time_ok()
False
>>> s1 = State(f1, id='t0', time=300, time_units='s')
>>> s2 = State(f2, id='tm1', time=0, time_units='s GMT')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> s4 = State(f2, f3, id='tm3')
>>> set = StateSet(s1, s2, s3, s4, delt=100)
>>> set.time_ok()
False
>>> s1 = State(f1, id='t0', time='100', time_units='s')
>>> s2 = State(f2, id='tp4', time=500, time_units='s')
>>> s3 = State(f3, id='tm2', time=-100, time_units='s')
>>> set = StateSet(s1, s2, s3, delt=100)
>>> set.time_ok()
Traceback (most recent call last):
    ...
TypeError: unsupported operand type(s) for +: 'int' and 'str'


#- Test additional methods:

    * Test refdata and setdata methods:

>>> data1 = N.array([1.1, 3.0, -4.8, 2.9, 2.1])
>>> data2 = N.array([31.1, 0.4, 4.0, -4.9, 9.1])
>>> data3 = N.array([3.1, 2.4, 4.9, 8.1, -9.2])
>>> f1a = Field(data1, id='u', units='m/s', axis=['Latitude'])
>>> f2a = Field(data2, id='v', units='m/s', axis=['Latitude'])
>>> f3a = Field(data3, id='w', units='m/s', axis=['Latitude'])
>>> f1b = Field(data1*0.1, id='u', units='m/s', axis=['Latitude'])
>>> f2b = Field(data2*0.2, id='v', units='m/s', axis=['Latitude'])
>>> s1 = State(f1a, f2a, f3a, long_name='winds1', id='t0')
>>> s2 = State(f1b, f2b, long_name='winds2', id='tm1')
>>> s3 = State(f1a, f2b, long_name='winds3', id='tm4')
>>> st = StateSet(s1, s2, id='velocities')
>>> ['%.6g' % st['tm1']['u'][i] for i in range(len(st['tm1']['u']))]
['0.11', '0.3', '-0.48', '0.29', '0.21']
>>> st.setdata({s3.id:s3})
>>> ['%.6g' % st['tm1']['u'][i] for i in range(len(st['tm1']['u']))]
Traceback (most recent call last):
    ...
KeyError: 'tm1'
>>> ['%.6g' % st['tm4']['u'][i] for i in range(len(st['tm4']['u']))]
['1.1', '3', '-4.8', '2.9', '2.1']
>>> f1a[1] = -999.
>>> ['%.6g' % f1a[i] for i in range(len(f1a))]
['1.1', '-999', '-4.8', '2.9', '2.1']
>>> ['%.6g' % st['tm4']['u'][i] for i in range(len(st['tm4']['u']))]
['1.1', '-999', '-4.8', '2.9', '2.1']
>>> a = st.refdata()
>>> a['tm4']['u'][3] = -888.
>>> ['%.6g' % f1a[i] for i in range(len(f1a))]
['1.1', '-999', '-4.8', '-888', '2.1']
>>> ['%.6g' % st['tm4']['u'][i] for i in range(len(st['tm4']['u']))]
['1.1', '-999', '-4.8', '-888', '2.1']
>>> type(a)
<type 'dict'>
>>> ['%.6g' % a['tm4']['u'][i] for i in range(len(a['tm4']['u']))]
['1.1', '-999', '-4.8', '-888', '2.1']

    * Test state_list method:

>>> s1 = State(f1a, f2a, f3a, long_name='3-D_winds', id='t0')
>>> s2 = State(f1b, f2b, long_name='2-D_winds')
>>> st = StateSet(s1, s2, id='velocities')
>>> print st.state_list()
[(None, '2-D_winds'), ('t0', '3-D_winds')]
"""}


if __name__ == "__main__":
    """Test the module documentation strings and __test__.
    """
    import doctest, sys, os
    sys.path.append(os.pardir)
    doctest.testmod(sys.modules[__name__])




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