SF.net SVN: matplotlib:[7707] trunk/toolkits/basemap/lib/mpl_toolkits/ basemap/netcdftime.py

[<js...@us...>]

Revision: 7707
          http://matplotlib.svn.sourceforge.net/matplotlib/?rev=7707&view=rev
Author:   jswhit
Date:     2009-09-08 12:06:35 +0000 (Tue, 08 Sep 2009)

Log Message:
-----------
update to latest version from netcdf4-python

Modified Paths:
--------------
    trunk/toolkits/basemap/lib/mpl_toolkits/basemap/netcdftime.py

Modified: trunk/toolkits/basemap/lib/mpl_toolkits/basemap/netcdftime.py
===================================================================
--- trunk/toolkits/basemap/lib/mpl_toolkits/basemap/netcdftime.py	2009-09-08 01:53:12 UTC (rev 7706)
+++ trunk/toolkits/basemap/lib/mpl_toolkits/basemap/netcdftime.py	2009-09-08 12:06:35 UTC (rev 7707)
@@ -934,18 +934,38 @@
     cdftime = utime(units,calendar=calendar)
     return cdftime.num2date(times)
 
-
 def _check_index(indices, dates, nctime, calendar):
-    """Assert that the time indices given correspond to the given dates."""
-    t = numpy.empty(len(indices), nctime.dtype)
-    for n,i in enumerate(indices):
-        t[n] = nctime[i]
-    assert numpy.all( num2date(t, nctime.units, calendar) == dates)
+    """Return True if the time indices given correspond to the given dates, 
+    False otherwise.
+    
+    Parameters: 
 
+    indices : sequence of integers
+    Positive integers indexing the time variable.
 
-def date2index(dates, nctime, calendar=None):
+    dates : sequence of datetime objects
+    Reference dates.
+
+    nctime : netCDF Variable object
+    NetCDF time object.
+
+    calendar : string
+    Calendar of nctime.
+     """
+    if  (indices <0).any():
+       return False
+       
+    if (indices >= nctime.shape[0]).any():
+        return False
+        
+    t = nctime[indices] 
+    return numpy.all( num2date(t, nctime.units, calendar) == dates)
+
+
+
+def date2index(dates, nctime, calendar=None, select='exact'):
     """
-    date2index(dates, nctime, calendar=None)
+    date2index(dates, nctime, calendar=None, select='exact')
     
     Return indices of a netCDF time variable corresponding to the given dates.
     
@@ -953,7 +973,8 @@
     The datetime objects should not include a time-zone offset.
     
     @param nctime: A netCDF time variable object. The nctime object must have a
-    C{units} attribute.
+    C{units} attribute. The entries are assumed to be stored in increasing 
+    order.
     
     @param calendar: Describes the calendar used in the time calculation.
     Valid calendars C{'standard', 'gregorian', 'proleptic_gregorian'
@@ -961,32 +982,54 @@
     Default is C{'standard'}, which is a mixed Julian/Gregorian calendar
     If C{calendar} is None, its value is given by C{nctime.calendar} or
     C{standard} if no such attribute exists.
+    
+    @param select: C{'exact', 'before', 'after', 'nearest'}
+      The index selection method. C{exact} will return the indices perfectly 
+      matching the dates given. C{before} and C{after} will return the indices 
+      corresponding to the dates just before or just after the given dates if 
+      an exact match cannot be found. C{nearest} will return the indices that 
+      correpond to the closest dates. 
     """
     # Setting the calendar.
     if calendar is None:
         calendar = getattr(nctime, 'calendar', 'standard')
 
     num = numpy.atleast_1d(date2num(dates, nctime.units, calendar))
-
-    index = numpy.empty(numpy.alen(dates), int)
-
+    
     # Trying to infer the correct index from the starting time and the stride.
-    try:
-        t0, t1 = nctime[:2]
-        dt = t1 - t0
-        index[:] = (num-t0)/dt
+    # This assumes that the times are increasing uniformly. 
+    t0, t1 = nctime[:2]
+    dt = t1 - t0
+    index = numpy.array((num-t0)/dt, int)
 
-        # Checking that the index really corresponds to the given date.
-        _check_index(index, dates, nctime, calendar)
-
-    except AssertionError:
-
-        # If check fails, use brute force method.
-        index[:] = numpy.digitize(num, nctime[:]) - 1
-
-        # Perform check again.
-        _check_index(index, dates, nctime, calendar)
-
+    # Checking that the index really corresponds to the given date.
+    # If the times do not correspond, then it means that the times
+    # are not increasing uniformly and we try the bisection method.
+    if not _check_index(index, dates, nctime, calendar):
+    
+        # Use the bisection method. Assumes the dates are ordered.
+        import bisect
+        
+        index = numpy.array([bisect.bisect_left(nctime, n) for n in num], int)
+        
+        nomatch = num2date(nctime[index], nctime.units) != dates
+        
+        if select == 'exact':
+            if not (num2date(nctime[index], nctime.units) == dates).all():
+                raise ValueError, 'Dates not found.'
+    
+        elif select == 'before':
+            index[nomatch] -= 1
+        
+        elif select == 'after':
+            pass
+        
+        elif select == 'nearest':
+            index[nomatch] = index[nomatch] - 1 * ( num[nomatch] < (nctime[index[nomatch]-1] + nctime[index[nomatch]]) / 2. )
+        
+        else:
+            raise ValueError, select
+    
     # convert numpy scalars or single element arrays to python ints.
     index = _toscalar(index)
 


This was sent by the SourceForge.net collaborative development platform, the world's largest Open Source development site.