ENH: Return nan values in tgt cells where a src nan cells contribute

Author: cpelley

stratify/_bounded_vinterp.py

@@ -49,8 +49,10 @@ def interpolate_conservative(z_target, z_src, fz_src, axis=-1):
 
     Note
     ----
-    Support for 1D z_target and corresponding ND z_src will be provided in
+    - Support for 1D z_target and corresponding ND z_src will be provided in
     future as driven by user requirement.
+    - Those cells, where 'nan' values in the source data contribute, a 'nan'
+      value is returned.
 
     """
     if z_src.ndim != z_target.ndim:

stratify/_conservative.pyx

@@ -72,8 +72,9 @@ cdef apply_weights(np.ndarray[np.float64_t, ndim=3] src_point,
 
     """
     cdef Py_ssize_t ind
-    cdef np.ndarray[np.float64_t, ndim=3] results
+    cdef np.ndarray[np.float64_t, ndim=3] results, weighted_contrib
     cdef np.ndarray[np.float64_t, ndim=2] weights
+    cdef np.ndarray[np.int64_t, ndim=2] nan_src_contrib
     results = np.zeros(
         [src_data.shape[0], tgt_point.shape[0], src_data.shape[2]],
         dtype='float64')
@@ -84,10 +85,16 @@ cdef apply_weights(np.ndarray[np.float64_t, ndim=3] src_point,
             msg = ('Weights calculation yields a less than conservative '
                    'result.  Aborting.')
             raise ValueError(msg)
+        weighted_contrib = weights * src_data[..., ind][..., None]
         results[..., ind] = (
-            weights * src_data[..., ind][..., None]).sum(axis=1)
-    # Return np.nan for those target cells where no source contributes.
-    results[:, weights.sum(axis=0) == 0, :] = np.nan
+            np.nansum(weighted_contrib, axis=1))
+        # Return nan values for those target cells, where there is any
+        # contribution of nan data from the source data.
+        nan_src_contrib = ((weights > 0) * np.isnan(weighted_contrib)).sum(axis=1)
+        results[..., ind][nan_src_contrib>0] = np.nan
+
+        # Return np.nan for those target cells where no source contributes.
+        results[:, weights.sum(axis=0) == 0, ind] = np.nan
     return results
 
 

stratify/tests/test_bounded_vinterp.py

@@ -53,6 +53,21 @@ def test_no_broadcasting(self):
         target_data = np.ones((3))*2
         assert_array_equal(res, target_data)
 
+    def test_source_with_nans(self):
+        # In the case of no broadcasting, transposed and reshaped array of form
+        # [broadcasting_dims, axis_interpolation, z_varying] should end up
+        # [1, axis_interpolation, 1].
+        data = self.data[0]
+        data[0] = data[-2:] = np.nan
+        target_bounds = self.gen_bounds(0, 6.5, 0.5)
+        res = bounded_vinterp.interpolate_conservative(target_bounds,
+                                                       self.bounds, data,
+                                                       axis=0)
+        target_data = np.ones((12))*.5
+        target_data[:2] = np.nan
+        target_data[-4:] = np.nan
+        assert_array_equal(res, target_data)
+
     def test_target_extends_above_source(self):
         # |-|-|-|-|-|-|-|   - Source
         # |-|-|-|-|-|-|-|-| - Target
@@ -62,7 +77,7 @@ def test_target_extends_above_source(self):
                                                        source_bounds,
                                                        self.data, axis=1)
         target_data = np.ones((4, 7))
-        target_data[:, -1] = 0
+        target_data[:, -1] = np.nan
         assert_array_equal(res, target_data)
 
     def test_target_extends_above_source_non_equally_spaced_coords(self):
@@ -88,7 +103,7 @@ def test_target_extends_below_source(self):
                                                        source_bounds,
                                                        self.data, axis=1)
         target_data = np.ones((4, 7))
-        target_data[:, 0] = 0
+        target_data[:, 0] = np.nan
         assert_array_equal(res, target_data)