Merge pull request #216 from kenkehoe/qc_indexing

QC indexing

act/qc/qcfilter.py

@@ -220,7 +220,7 @@ def add_test(self, var_name, index=None, test_number=None,
         ----------
         var_name : str
             data variable name
-        index : int or list of int or numpy array
+        index : int, bool, list of int or bool, numpy array, tuple of numpy arrays
             Indexes into quality control array to set the test bit.
             If not set or set to None will not set the test on any
             element of the quality control variable but will still
@@ -258,8 +258,12 @@ def add_test(self, var_name, index=None, test_number=None,
                              'keyword when calling the add_test method')
 
         # This ensures the indexing will work even if given float values.
-        if index is not None:
-            index = np.array(index, dtype=np.int)
+        # Preserves tuples from np.where() or boolean arrays for standard
+        # python indexing.
+        if index is not None and not isinstance(index, (np.ndarray, tuple)):
+            index = np.array(index)
+            if index.dtype.kind == 'f':
+                index = index.astype(int)
 
         # Ensure assessment is lowercase and capitalized to be consistent
         test_assessment = test_assessment.lower().capitalize()

act/qc/qctests.py

@@ -95,15 +95,13 @@ def add_missing_value_test(self, var_name, missing_value=None,
         # Ensure missing_value attribute is matching data type
         missing_value = np.array(missing_value, dtype=self._obj[var_name].values.dtype.type)
 
-        if np.isnan(missing_value) is False:
-            data = np.ma.masked_equal(self._obj[var_name].values,
-                                      missing_value)
-        else:
-            data = np.ma.masked_invalid(self._obj[var_name].values)
-
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)
+        # New method using straight numpy instead of masked array
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            if np.isnan(missing_value) is False:
+                index = np.equal(self._obj[var_name].values, missing_value)
+            else:
+                index = np.isnan(self._obj[var_name].values)
 
         test_dict = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -179,10 +177,10 @@ def add_less_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_less(self._obj[var_name].values, limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        # New method with straight numpy
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.less(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -259,11 +257,9 @@ def add_greater_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_greater(self._obj[var_name].values,
-                                    limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.greater(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -342,10 +338,9 @@ def add_less_equal_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_less_equal(self._obj[var_name].values, limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.less_equal(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -424,11 +419,9 @@ def add_greater_equal_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_greater_equal(self._obj[var_name].values,
-                                          limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.greater_equal(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -505,10 +498,9 @@ def add_equal_to_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_equal(self._obj[var_name].values, limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.equal(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -585,10 +577,9 @@ def add_not_equal_to_test(self, var_name, limit_value, test_meaning=None,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        data = np.ma.masked_not_equal(self._obj[var_name].values, limit_value)
-        if data.mask.size == 1:
-            data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
+            index = np.not_equal(self._obj[var_name].values, limit_value)
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -674,12 +665,14 @@ def add_outside_test(self, var_name, limit_value_lower, limit_value_upper,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        with np.errstate(invalid='ignore'):
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
             data = np.ma.masked_outside(self._obj[var_name].values,
                                         limit_value_lower, limit_value_upper)
         if data.mask.size == 1:
             data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+
+        index = data.mask
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -767,12 +760,14 @@ def add_inside_test(self, var_name, limit_value_lower, limit_value_upper,
         if prepend_text is not None:
             test_meaning = ': '.join((prepend_text, test_meaning))
 
-        with np.errstate(invalid='ignore'):
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=RuntimeWarning)
             data = np.ma.masked_inside(self._obj[var_name].values,
                                        limit_value_lower, limit_value_upper)
         if data.mask.size == 1:
             data.mask = np.full(data.data.shape, data.mask, dtype=bool)
-        index = np.where(data.mask)[0]
+
+        index = data.mask
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -852,7 +847,7 @@ def add_persistence_test(self, var_name, window=10, test_limit=0.0001,
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=RuntimeWarning)
             stddev = data.rolling(time=window, min_periods=min_periods, center=True).std()
-            index = np.where(stddev < test_limit)
+            index = stddev < test_limit
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,
@@ -1001,7 +996,7 @@ def add_difference_test(self, var_name, dataset2_dict=None, ds2_var_name=None,
                 else:
                     diff = np.absolute(pd_c[var_name] - pd_c[ds2_var_name])
 
-                index = np.where(diff > diff_limit)
+                index = diff > diff_limit
 
         result = self._obj.qcfilter.add_test(
             var_name, index=index,