BUG: fix replace on non-contiguous array views (#513)

Author: emmanuel-ferdman

bottleneck/src/nonreduce_template.c

@@ -28,19 +28,20 @@ replace_DTYPE0(PyArrayObject *a, double old, double new) {
     BN_BEGIN_ALLOW_THREADS
     const npy_DTYPE0 oldf = (npy_DTYPE0)old;
     const npy_DTYPE0 newf = (npy_DTYPE0)new;
+    const npy_intp stride = it.stride;
     if (old == old) {
         WHILE {
             npy_DTYPE0* array = PA(DTYPE0);
             FOR {
-                array[it.i] = array[it.i] == oldf ? newf : array[it.i];
+                array[it.i * stride] = array[it.i * stride] == oldf ? newf : array[it.i * stride];
             }
             NEXT
         }
     } else {
         WHILE {
             npy_DTYPE0* array = PA(DTYPE0);
             FOR {
-                array[it.i] = array[it.i] != array[it.i] ? newf : array[it.i];
+                array[it.i * stride] = array[it.i * stride] != array[it.i * stride] ? newf : array[it.i * stride];
             }
             NEXT
         }
@@ -69,12 +70,13 @@ replace_DTYPE0(PyArrayObject *a, double old, double new) {
             return NULL;
         }
         BN_BEGIN_ALLOW_THREADS
+        const npy_intp stride = it.stride;
         WHILE {
             npy_DTYPE0* array = (npy_DTYPE0 *)it.pa;
             npy_intp i;
             // clang has a large perf regression when using the FOR macro here
             for (i=0; i < it.length; i++) {
-                array[i] = array[i] == oldint ? newint : array[i];
+                array[i * stride] = array[i * stride] == oldint ? newint : array[i * stride];
             }
             NEXT
         }

bottleneck/tests/nonreduce_test.py

@@ -8,7 +8,7 @@
 
 import bottleneck as bn
 
-from .util import DTYPES, INT_DTYPES, array_order, arrays
+from .util import DTYPES, FLOAT_DTYPES, INT_DTYPES, array_order, arrays
 
 
 @pytest.mark.parametrize(
@@ -138,3 +138,31 @@ def test_replace_newaxis(dtype):
     array = np.ones((2, 2), dtype=dtype)[..., np.newaxis]
     result = bn.replace(array, 1, 2)
     assert (result == 2).all().all()
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_replace_view(dtype):
+    """Test replace on non-contiguous view"""
+    expected_array = np.arange(20, dtype=dtype)
+    expected_view = expected_array[::2]
+    bn.slow.replace(expected_view, 10, -1)
+    array = np.arange(20, dtype=dtype)
+    view = array[::2]
+    bn.replace(view, 10, -1)
+    assert_array_equal(view, expected_view)
+    assert_array_equal(array, expected_array)
+
+
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_replace_nan_view(dtype):
+    """Test replace NaN on non-contiguous view"""
+    expected_array = np.ones((4, 3, 2), dtype=dtype)
+    expected_array[::2, :, 0] = np.nan
+    expected_view = expected_array[:, :, 0]
+    bn.slow.replace(expected_view, np.nan, 0)
+    array = np.ones((4, 3, 2), dtype=dtype)
+    array[::2, :, 0] = np.nan
+    view = array[:, :, 0]
+    bn.replace(view, np.nan, 0)
+    assert_array_equal(view, expected_view)
+    assert_array_equal(array, expected_array)