Merge pull request #12988 from qwhelan/bool_ufunc

ENH: Create boolean and integer ufuncs for isnan, isinf, and isfinite.

5 files changed, 88 insertions, 3 deletions

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index 687a8467b..de0bb81fe 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -827,7 +827,7 @@ defdict = {
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isnan'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'isnat':
     Ufunc(1, 1, None,
@@ -839,13 +839,13 @@ defdict = {
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isinf'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'isfinite':
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isfinite'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'signbit':
     Ufunc(1, 1, None,
diff --git a/numpy/core/src/umath/fast_loop_macros.h b/numpy/core/src/umath/fast_loop_macros.h
index 37656dcf5..e3cfa1f72 100644
--- a/numpy/core/src/umath/fast_loop_macros.h
+++ b/numpy/core/src/umath/fast_loop_macros.h
@@ -64,6 +64,8 @@
 #define IS_UNARY_CONT(tin, tout) (steps[0] == sizeof(tin) && \
                                   steps[1] == sizeof(tout))
 
+#define IS_OUTPUT_CONT(tout) (steps[1] == sizeof(tout))
+
 #define IS_BINARY_REDUCE ((args[0] == args[2])\
         && (steps[0] == steps[2])\
         && (steps[0] == 0))
@@ -84,6 +86,29 @@
 
 /*
  * loop with contiguous specialization
+ * op should be the code storing the result in `tout * out`
+ * combine with NPY_GCC_OPT_3 to allow autovectorization
+ * should only be used where its worthwhile to avoid code bloat
+ */
+#define BASE_OUTPUT_LOOP(tout, op) \
+    OUTPUT_LOOP { \
+        tout * out = (tout *)op1; \
+        op; \
+    }
+#define OUTPUT_LOOP_FAST(tout, op) \
+    do { \
+    /* condition allows compiler to optimize the generic macro */ \
+    if (IS_OUTPUT_CONT(tout)) { \
+        BASE_OUTPUT_LOOP(tout, op) \
+    } \
+    else { \
+        BASE_OUTPUT_LOOP(tout, op) \
+    } \
+    } \
+    while (0)
+
+/*
+ * loop with contiguous specialization
  * op should be the code working on `tin in` and
  * storing the result in `tout * out`
  * combine with NPY_GCC_OPT_3 to allow autovectorization
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index 04e6cbdee..1e4ab350b 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -644,6 +644,19 @@ BOOL__ones_like(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UN
 }
 
 
+/**begin repeat
+ * #kind = isnan, isinf, isfinite#
+ * #func = npy_isnan, npy_isinf, npy_isfinite#
+ * #val = NPY_FALSE, NPY_FALSE, NPY_TRUE#
+ **/
+NPY_NO_EXPORT void
+BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    OUTPUT_LOOP_FAST(npy_bool, *out = @val@);
+}
+
+/**end repeat**/
+
 /*
  *****************************************************************************
  **                           INTEGER LOOPS
@@ -875,6 +888,18 @@ NPY_NO_EXPORT void
     }
 }
 
+/**begin repeat1
+ * #kind = isnan, isinf, isfinite#
+ * #func = npy_isnan, npy_isinf, npy_isfinite#
+ * #val = NPY_FALSE, NPY_FALSE, NPY_TRUE#
+ **/
+NPY_NO_EXPORT void
+@TYPE@_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    OUTPUT_LOOP_FAST(npy_bool, *out = @val@);
+}
+/**end repeat1**/
+
 /**end repeat**/
 
 /**begin repeat
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index 5264a6533..f48319056 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -38,6 +38,13 @@ BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED
 NPY_NO_EXPORT void
 BOOL__ones_like(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data));
 
+/**begin repeat
+ * #kind = isnan, isinf, isfinite#
+ **/
+NPY_NO_EXPORT void
+BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+/**end repeat**/
+
 /*
  *****************************************************************************
  **                           INTEGER LOOPS
@@ -146,6 +153,13 @@ NPY_NO_EXPORT void
 NPY_NO_EXPORT void
 @S@@TYPE@_lcm(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
 
+/**begin repeat2
+ * #kind = isnan, isinf, isfinite#
+ **/
+NPY_NO_EXPORT void
+@S@@TYPE@_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+/**end repeat2**/
+
 /**end repeat1**/
 
 /**end repeat**/
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 478a08397..b6b68d922 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -1915,3 +1915,24 @@ class TestUfunc(object):
         exc = pytest.raises(TypeError, np.sqrt, None)
         # minimally check the exception text
         assert 'loop of ufunc does not support' in str(exc)
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_not_finite(self, nat):
+        try:
+            assert not np.isfinite(nat)
+        except TypeError:
+            pass  # ok, just not implemented
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_nan(self, nat):
+        try:
+            assert np.isnan(nat)
+        except TypeError:
+            pass  # ok, just not implemented
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_not_inf(self, nat):
+        try:
+            assert not np.isinf(nat)
+        except TypeError:
+            pass  # ok, just not implemented