averaging intervals work with masked arrays

metobs/data/__init__.py

@@ -8,7 +8,8 @@ import math
 from datetime import timedelta
 
 import fpconst
-from numpy import array, average, zeros, radians, degrees, arctan2, sin, cos
+from numpy import array, average, zeros, radians, degrees, arctan2, sin, cos, ndarray
+from numpy.ma import masked_array, average as masked_average, MaskedArray, masked_where, column_stack
 from pydap.client import open_url as dapopen
 
 from metobs import mytime
@@ -32,6 +33,12 @@ def average_for_interval(basetime, arr, interval):
     :type arr: ``numpy.array``
     :param interval: the averaging interval in seconds
     """
+    if type(arr) == ndarray:
+        avg_func = average
+        array_func = array
+    elif type(arr) == MaskedArray:
+        avg_func = masked_average
+        array_func = masked_array
     arr_out = []
     stop_dt = basetime + timedelta(seconds=interval)
     stop_idx = 1
@@ -55,12 +62,13 @@ def average_for_interval(basetime, arr, interval):
             # matrix of data to be averaged
             avg_input = arr[start_idx:stop_idx,1:]
             avg_out = array(zeros(arr[0].shape[0]))
-            avg_out.fill(NaN)
+            if type(arr) == ndarray: avg_out.fill(NaN)
+            elif type(arr) == MaskedArray: masked_where(avg_out == 0, avg_out)
             
             # timestamp for center of averaging interval
             int_dt = stop_dt - timedelta(seconds=(interval/2))
             avg_out[0] = mytime.datetime_to_epoch(int_dt)
-            avg_out[1:] = average(avg_input, axis=0)
+            avg_out[1:] = avg_func(avg_input, axis=0)
             arr_out.append(avg_out)
             
             # adjust conditions for next average
@@ -69,8 +77,12 @@ def average_for_interval(basetime, arr, interval):
             stop_dt += timedelta(seconds=interval)
 
     avg_arr = arr[start_idx:,:]
-    arr_out.append(average(avg_arr, axis=0))
-    return array(arr_out)
+    avg_arr = avg_func(avg_arr, axis=0)
+    arr_out.append(avg_arr)
+    if type(arr) == ndarray: arr_out = array_func(arr_out)
+    elif type(arr) == MaskedArray:
+        arr_out = column_stack(arr_out).transpose()
+    return arr_out
 
 def average_for_interval_degree(basetime, arr, interval):
     """Generate averages for timestamped data at a particular interval.