Merge pull request #526 from pydata/add_complex

Add complex arguments for new boolean functions

Author: FrancescAlted

ADDFUNCS.rst

@@ -171,6 +171,23 @@ Add clauses to generate the FUNC_CODES from the ``functions.hpp`` header, making
     };
     #endif
 
+Some functions (e.g. ``fmod``, ``isnan``) are not available in MKL, and so must be hard-coded here as well:
+
+.. code-block:: cpp
+
+    #ifdef USE_VML
+    /* no isnan, isfinite or isinf in VML */
+    static void vdIsfinite(MKL_INT n, const double* x1, bool* dest)
+    {
+        MKL_INT j;
+        for (j=0; j<n; j++) {
+            dest[j] = isfinited(x1[j]);
+        };
+    };
+    #endif
+
+The complex case is slightlñy different (see other examples in the same file).
+
 Add case handling to the ``check_program`` function
 
 .. code-block:: cpp
@@ -219,4 +236,6 @@ In many cases this process will not be very smooth since one relies on the inter
 
 * Depending on the new function signature (above all if the out type is different to the input types), one may have to edit the ``__init__`` function in the ``FuncNode`` class in ``expressions.py``.
 
+* Functions which accept and/or return complex arguments must be added to the ``complex_functions.hpp`` file (take care when adding them in ``interpreter.cpp`` and ``interp_body.cpp``, since their signatures are usually a bit different).
+
 * Depending on MSVC support, namespace clashes, casting problems, it may be necessary to make various changes to ``numexpr/numexpr_config.hpp`` and ``numexpr/msvc_function_stubs.hpp``. For example, in PR #523, non-clashing wrappers were introduced for ``isnan`` and ``isfinite`` since the float versions ``isnanf, isfinitef`` were inconsistently defined (and output ints) - depending on how strict the platform interpreter is, the implicit cast from int to bool was acceptable or not for example. In addition, the base functions were in different namespaces or had different names across platforms.

doc/user_guide.rst

@@ -1,4 +1,4 @@
-NumExpr 2.8 User Guide
+NumExpr 2.12 User Guide
 ======================
 
 The NumExpr package supplies routines for the fast evaluation of
@@ -201,6 +201,8 @@ The next are the current supported set:
 
     * :code:`where(bool, number1, number2): number` -- number1 if the bool condition
       is true, number2 otherwise.
+    * :code:`{isinf, isnan, isfinite}(float|complex): bool` -- returns element-wise True
+      for ``inf`` or ``NaN``, ``NaN``, not ``inf`` respectively.
     * :code:`{sin,cos,tan}(float|complex): float|complex` -- trigonometric sine,
       cosine or tangent.
     * :code:`{arcsin,arccos,arctan}(float|complex): float|complex` -- trigonometric

numexpr/complex_functions.hpp

@@ -424,4 +424,33 @@ nc_abs(std::complex<double> *x, std::complex<double> *r)
     r->imag(0);
 }
 
+static bool
+nc_isinf(std::complex<double> *x)
+{
+    double xr=x->real(), xi=x->imag();
+    bool bi,br;
+    bi = isinfd(xi);
+    br = isinfd(xr);
+    return bi || br;
+}
+
+static bool
+nc_isnan(std::complex<double> *x)
+{
+    double xr=x->real(), xi=x->imag();
+    bool bi,br;
+    bi = isnand(xi);
+    br = isnand(xr);
+    return bi || br;
+}
+
+static bool
+nc_isfinite(std::complex<double> *x)
+{
+    double xr=x->real(), xi=x->imag();
+    bool bi,br;
+    bi = isfinited(xi);
+    br = isfinited(xr);
+    return bi && br;
+}
 #endif // NUMEXPR_COMPLEX_FUNCTIONS_HPP

numexpr/functions.hpp

@@ -166,3 +166,17 @@ FUNC_CCC(FUNC_CCC_LAST,  NULL,      NULL)
 #undef ELIDE_FUNC_CCC
 #undef FUNC_CCC
 #endif
+
+// complex -> boolean functions
+#ifndef FUNC_BC
+#define ELIDE_FUNC_BC
+#define FUNC_BC(...)
+#endif // use wrappers as there is name collision with isnanf in std
+FUNC_BC(FUNC_ISNAN_BC,   "isnan_bc",    nc_isnan, vzIsnan)
+FUNC_BC(FUNC_ISFINITE_BC, "isfinite_bc", nc_isfinite, vzIsfinite)
+FUNC_BC(FUNC_ISINF_BC, "isinf_bc", nc_isinf, vzIsinf)
+FUNC_BC(FUNC_BC_LAST,    NULL,            NULL,        NULL)
+#ifdef ELIDE_FUNC_BC
+#undef ELIDE_FUNC_BC
+#undef FUNC_BC
+#endif

numexpr/interp_body.cpp

@@ -469,6 +469,17 @@
             VEC_ARG1(b_dest = functions_bd[arg2](d1));
 #endif
 
+        case OP_FUNC_BCN:
+#ifdef USE_VML
+            VEC_ARG1_VML(functions_bc_vml[arg2](BLOCK_SIZE,
+                        (const MKL_Complex16*)x1, (bool*)dest));
+#else
+            VEC_ARG1(ca.real(c1r);
+                     ca.imag(c1i);
+                     b_dest = functions_bc[arg2](&ca));
+#endif
+
+
         /* Reductions */
         case OP_SUM_IIN: VEC_ARG1(i_reduce += i1);
         case OP_SUM_LLN: VEC_ARG1(l_reduce += l1);

numexpr/interpreter.cpp

@@ -220,6 +220,31 @@ FuncBFPtr functions_bf[] = {
 };
 #endif
 
+#ifdef USE_VML
+/* no isnan, isfinite or isinf in VML */
+static void vfIsfinite(MKL_INT n, const float* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isfinitef_(x1[j]);
+    };
+};
+static void vfIsinf(MKL_INT n, const float* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isinff_(x1[j]);
+    };
+};
+static void vfIsnan(MKL_INT n, const float* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isnanf_(x1[j]);
+    };
+};
+#endif
+
 #ifdef USE_VML
 typedef void (*FuncBFPtr_vml)(MKL_INT, const float*, bool*);
 FuncBFPtr_vml functions_bf_vml[] = {
@@ -236,6 +261,31 @@ FuncBDPtr functions_bd[] = {
 #undef FUNC_BD
 };
 
+#ifdef USE_VML
+/* no isnan, isfinite or isinf in VML */
+static void vdIsfinite(MKL_INT n, const double* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isfinited(x1[j]);
+    };
+};
+static void vdIsinf(MKL_INT n, const double* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isinfd(x1[j]);
+    };
+};
+static void vdIsnan(MKL_INT n, const double* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isnand(x1[j]);
+    };
+};
+#endif
+
 #ifdef USE_VML
 typedef void (*FuncBDPtr_vml)(MKL_INT, const double*, bool*);
 FuncBDPtr_vml functions_bd_vml[] = {
@@ -245,6 +295,47 @@ FuncBDPtr_vml functions_bd_vml[] = {
 };
 #endif
 
+typedef bool (*FuncBCPtr)(std::complex<double>*);
+FuncBCPtr functions_bc[] = {
+#define FUNC_BC(fop, s, f, ...) f,
+#include "functions.hpp"
+#undef FUNC_BC
+};
+
+#ifdef USE_VML
+/* no isnan, isfinite or isinf in VML */
+static void vzIsfinite(MKL_INT n, const MKL_Complex16* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isfinited(x1[j].real) && isfinited(x1[j].imag);
+    };
+};
+static void vzIsinf(MKL_INT n, const MKL_Complex16* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isinfd(x1[j].real) || isinfd(x1[j].imag);
+    };
+};
+static void vzIsnan(MKL_INT n, const MKL_Complex16* x1, bool* dest)
+{
+    MKL_INT j;
+    for (j=0; j<n; j++) {
+        dest[j] = isnand(x1[j].real) || isnand(x1[j].imag);
+    };
+};
+#endif
+
+#ifdef USE_VML
+typedef void (*FuncBCPtr_vml)(MKL_INT, const MKL_Complex16[], bool*);
+FuncBCPtr_vml functions_bc_vml[] = {
+#define FUNC_BC(fop, s, f, f_vml) f_vml,
+#include "functions.hpp"
+#undef FUNC_BC
+};
+#endif
+
 #ifdef USE_VML
 /* Fake vdConj function just for casting purposes inside numexpr */
 static void vdConj(MKL_INT n, const double* x1, double* dest)
@@ -517,7 +608,14 @@ check_program(NumExprObject *self)
                         PyErr_Format(PyExc_RuntimeError, "invalid program: funccode out of range (%i) at %i", arg, argloc);
                         return -1;
                     }
-                } else if (op >= OP_REDUCTION) {
+                }
+                  else if (op == OP_FUNC_BCN) {
+                    if (arg < 0 || arg >= FUNC_BC_LAST) {
+                        PyErr_Format(PyExc_RuntimeError, "invalid program: funccode out of range (%i) at %i", arg, argloc);
+                        return -1;
+                    }
+                }
+                  else if (op >= OP_REDUCTION) {
                     ;
                 } else {
                     PyErr_Format(PyExc_RuntimeError, "invalid program: internal checker error processing %i", argloc);

numexpr/interpreter.hpp

@@ -42,6 +42,12 @@ enum FuncBDCodes {
 #undef FUNC_BD
 };
 
+enum FuncBCCodes {
+#define FUNC_BC(fop, ...) fop,
+#include "functions.hpp"
+#undef FUNC_BC
+};
+
 enum FuncDDDCodes {
 #define FUNC_DDD(fop, ...) fop,
 #include "functions.hpp"

numexpr/module.cpp

@@ -508,13 +508,15 @@ PyInit_interpreter(void) {
 #define FUNC_DD(name, sname, ...)  add_func(name, sname);
 #define FUNC_BF(name, sname, ...)  add_func(name, sname);
 #define FUNC_BD(name, sname, ...)  add_func(name, sname);
+#define FUNC_BC(name, sname, ...)  add_func(name, sname);
 #define FUNC_DDD(name, sname, ...) add_func(name, sname);
 #define FUNC_CC(name, sname, ...)  add_func(name, sname);
 #define FUNC_CCC(name, sname, ...) add_func(name, sname);
 #include "functions.hpp"
 #undef FUNC_CCC
 #undef FUNC_CC
 #undef FUNC_DDD
+#undef FUNC_BC
 #undef FUNC_BD
 #undef FUNC_BF
 #undef FUNC_DD

numexpr/opcodes.hpp

@@ -150,43 +150,44 @@ OPCODE(105, OP_CONTAINS_BSS, "contains_bss", Tb, Ts, Ts, T0)
 //Boolean outputs
 OPCODE(106, OP_FUNC_BDN, "func_bdn", Tb, Td, Tn, T0)
 OPCODE(107, OP_FUNC_BFN, "func_bfn", Tb, Tf, Tn, T0)
+OPCODE(108, OP_FUNC_BCN, "func_bcn", Tb, Tc, Tn, T0)
 
 // Reductions always have to be at the end - parts of the code
 // use > OP_REDUCTION to decide whether operation is a reduction
-OPCODE(108, OP_REDUCTION, NULL, T0, T0, T0, T0)
+OPCODE(109, OP_REDUCTION, NULL, T0, T0, T0, T0)
 
 /* Last argument in a reduction is the axis of the array the
    reduction should be applied along. */
 
-OPCODE(109, OP_SUM_IIN, "sum_iin", Ti, Ti, Tn, T0)
-OPCODE(110, OP_SUM_LLN, "sum_lln", Tl, Tl, Tn, T0)
-OPCODE(111, OP_SUM_FFN, "sum_ffn", Tf, Tf, Tn, T0)
-OPCODE(112, OP_SUM_DDN, "sum_ddn", Td, Td, Tn, T0)
-OPCODE(113, OP_SUM_CCN, "sum_ccn", Tc, Tc, Tn, T0)
-
-OPCODE(114, OP_PROD, NULL, T0, T0, T0, T0)
-OPCODE(115, OP_PROD_IIN, "prod_iin", Ti, Ti, Tn, T0)
-OPCODE(116, OP_PROD_LLN, "prod_lln", Tl, Tl, Tn, T0)
-OPCODE(117, OP_PROD_FFN, "prod_ffn", Tf, Tf, Tn, T0)
-OPCODE(118, OP_PROD_DDN, "prod_ddn", Td, Td, Tn, T0)
-OPCODE(119, OP_PROD_CCN, "prod_ccn", Tc, Tc, Tn, T0)
-
-OPCODE(120, OP_MIN, NULL, T0, T0, T0, T0)
-OPCODE(121, OP_MIN_IIN, "min_iin", Ti, Ti, Tn, T0)
-OPCODE(122, OP_MIN_LLN, "min_lln", Tl, Tl, Tn, T0)
-OPCODE(123, OP_MIN_FFN, "min_ffn", Tf, Tf, Tn, T0)
-OPCODE(124, OP_MIN_DDN, "min_ddn", Td, Td, Tn, T0)
-
-OPCODE(125, OP_MAX, NULL, T0, T0, T0, T0)
-OPCODE(126, OP_MAX_IIN, "max_iin", Ti, Ti, Tn, T0)
-OPCODE(127, OP_MAX_LLN, "max_lln", Tl, Tl, Tn, T0)
-OPCODE(128, OP_MAX_FFN, "max_ffn", Tf, Tf, Tn, T0)
-OPCODE(129, OP_MAX_DDN, "max_ddn", Td, Td, Tn, T0)
+OPCODE(110, OP_SUM_IIN, "sum_iin", Ti, Ti, Tn, T0)
+OPCODE(111, OP_SUM_LLN, "sum_lln", Tl, Tl, Tn, T0)
+OPCODE(112, OP_SUM_FFN, "sum_ffn", Tf, Tf, Tn, T0)
+OPCODE(113, OP_SUM_DDN, "sum_ddn", Td, Td, Tn, T0)
+OPCODE(114, OP_SUM_CCN, "sum_ccn", Tc, Tc, Tn, T0)
+
+OPCODE(115, OP_PROD, NULL, T0, T0, T0, T0)
+OPCODE(116, OP_PROD_IIN, "prod_iin", Ti, Ti, Tn, T0)
+OPCODE(117, OP_PROD_LLN, "prod_lln", Tl, Tl, Tn, T0)
+OPCODE(118, OP_PROD_FFN, "prod_ffn", Tf, Tf, Tn, T0)
+OPCODE(119, OP_PROD_DDN, "prod_ddn", Td, Td, Tn, T0)
+OPCODE(120, OP_PROD_CCN, "prod_ccn", Tc, Tc, Tn, T0)
+
+OPCODE(121, OP_MIN, NULL, T0, T0, T0, T0)
+OPCODE(122, OP_MIN_IIN, "min_iin", Ti, Ti, Tn, T0)
+OPCODE(123, OP_MIN_LLN, "min_lln", Tl, Tl, Tn, T0)
+OPCODE(124, OP_MIN_FFN, "min_ffn", Tf, Tf, Tn, T0)
+OPCODE(125, OP_MIN_DDN, "min_ddn", Td, Td, Tn, T0)
+
+OPCODE(126, OP_MAX, NULL, T0, T0, T0, T0)
+OPCODE(127, OP_MAX_IIN, "max_iin", Ti, Ti, Tn, T0)
+OPCODE(128, OP_MAX_LLN, "max_lln", Tl, Tl, Tn, T0)
+OPCODE(129, OP_MAX_FFN, "max_ffn", Tf, Tf, Tn, T0)
+OPCODE(130, OP_MAX_DDN, "max_ddn", Td, Td, Tn, T0)
 
 /*
 When we get to 255, will maybe have to change code again
 (change latin_1 encoding in necompiler.py, use something
 other than unsigned char for OPCODE table)
 */
 /* Should be the last opcode */
-OPCODE(130, OP_END, NULL, T0, T0, T0, T0)
+OPCODE(131, OP_END, NULL, T0, T0, T0, T0)

numexpr/tests/test_numexpr.py

@@ -709,6 +709,11 @@ def test_bool_funcs(self):
         assert np.all(evaluate("isnan(a)") == np.isnan(a))
         assert np.all(evaluate("isfinite(a)") == np.isfinite(a))
         assert np.all(evaluate("isinf(a)") == np.isinf(a))
+        a = a.astype(np.complex128)
+        assert a.dtype == np.complex128
+        assert np.all(evaluate("isnan(a)") == np.isnan(a))
+        assert np.all(evaluate("isfinite(a)") == np.isfinite(a))
+        assert np.all(evaluate("isinf(a)") == np.isinf(a))
 
     if 'sparc' not in platform.machine():
         # Execution order set here so as to not use too many threads