ufunc.hpp 5.58 KB
Newer Older
1
2
3
4
5
6
7
//
// Copyright (c) 2020 INRIA
//

#ifndef __eigenpy_ufunc_hpp__
#define __eigenpy_ufunc_hpp__

Justin Carpentier's avatar
Justin Carpentier committed
8
#include "eigenpy/register.hpp"
9
10
11
12
13
14

namespace eigenpy
{
  namespace internal
  {
  
15
16
17
18
19
20
#ifdef NPY_1_19_API_VERSION
  #define EIGENPY_NPY_CONST_UFUNC_ARG const
#else
  #define EIGENPY_NPY_CONST_UFUNC_ARG
#endif
  
21
22
#define EIGENPY_REGISTER_BINARY_OPERATOR(name,op) \
    template<typename T1, typename T2, typename R> \
23
    void binary_op_##name(char** args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * dimensions, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * steps, void * /*data*/) \
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
    { \
      npy_intp is0 = steps[0], is1 = steps[1], \
      os = steps[2], n = *dimensions; \
      char * i0 = args[0], *i1 = args[1], *o = args[2]; \
      int k; \
      for (k = 0; k < n; k++) \
      { \
        T1 & x = *static_cast<T1*>(static_cast<void*>(i0)); \
        T2 & y = *static_cast<T2*>(static_cast<void*>(i1)); \
        R & res = *static_cast<R*>(static_cast<void*>(o)); \
        res = x op y; \
        i0 += is0; i1 += is1; o += os; \
      } \
    } \
    \
    template<typename T> \
40
    void binary_op_##name(char** args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * dimensions, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * steps, void * data) \
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    { \
      binary_op_##name<T,T,T>(args,dimensions,steps,data); \
    }
  
    EIGENPY_REGISTER_BINARY_OPERATOR(add,+)
    EIGENPY_REGISTER_BINARY_OPERATOR(subtract,-)
    EIGENPY_REGISTER_BINARY_OPERATOR(multiply,*)
    EIGENPY_REGISTER_BINARY_OPERATOR(divide,/)
    EIGENPY_REGISTER_BINARY_OPERATOR(equal,==)
    EIGENPY_REGISTER_BINARY_OPERATOR(not_equal,!=)
    EIGENPY_REGISTER_BINARY_OPERATOR(less,<)
    EIGENPY_REGISTER_BINARY_OPERATOR(greater,>)
    EIGENPY_REGISTER_BINARY_OPERATOR(less_equal,<=)
    EIGENPY_REGISTER_BINARY_OPERATOR(greater_equal,>=)
  
56
57
  #define EIGENPY_REGISTER_UNARY_OPERATOR(name,op) \
    template<typename T, typename R> \
58
    void unary_op_##name(char** args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * dimensions, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * steps, void * /*data*/) \
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    { \
      npy_intp is = steps[0], \
      os = steps[1], n = *dimensions; \
      char * i = args[0], *o = args[1]; \
      int k; \
      for (k = 0; k < n; k++) \
      { \
        T & x = *static_cast<T*>(static_cast<void*>(i)); \
        R & res = *static_cast<R*>(static_cast<void*>(o)); \
        res = op x; \
        i += is; o += os; \
      } \
    } \
    \
    template<typename T> \
74
    void unary_op_##name(char** args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * dimensions, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp * steps, void * data) \
75
76
77
78
79
80
    { \
      unary_op_##name<T,T>(args,dimensions,steps,data); \
    }
  
    EIGENPY_REGISTER_UNARY_OPERATOR(negative,-)
  
81
82
  } // namespace internal
  
Justin Carpentier's avatar
Justin Carpentier committed
83
#define EIGENPY_REGISTER_BINARY_UFUNC(name,code,T1,T2,R) { \
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
   PyUFuncObject* ufunc = \
       (PyUFuncObject*)PyObject_GetAttrString(numpy, #name); \
   int _types[3] = { Register::getTypeCode<T1>(), Register::getTypeCode<T2>(), Register::getTypeCode<R>()}; \
   if (!ufunc) { \
       /*goto fail; \*/ \
   } \
   if (sizeof(_types)/sizeof(int)!=ufunc->nargs) { \
       PyErr_Format(PyExc_AssertionError, \
                    "ufunc %s takes %d arguments, our loop takes %lu", \
                    #name, ufunc->nargs, (unsigned long) \
                    (sizeof(_types)/sizeof(int))); \
       Py_DECREF(ufunc); \
   } \
   if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc, code, \
        internal::binary_op_##name<T1,T2,R>, _types, 0) < 0) { \
       /*Py_DECREF(ufunc);*/ \
       /*goto fail; \*/ \
   } \
   Py_DECREF(ufunc); \
}
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
  
#define EIGENPY_REGISTER_UNARY_UFUNC(name,code,T,R) { \
   PyUFuncObject* ufunc = \
       (PyUFuncObject*)PyObject_GetAttrString(numpy, #name); \
   int _types[2] = { Register::getTypeCode<T>(), Register::getTypeCode<R>()}; \
   if (!ufunc) { \
       /*goto fail; \*/ \
   } \
   if (sizeof(_types)/sizeof(int)!=ufunc->nargs) { \
       PyErr_Format(PyExc_AssertionError, \
                    "ufunc %s takes %d arguments, our loop takes %lu", \
                    #name, ufunc->nargs, (unsigned long) \
                    (sizeof(_types)/sizeof(int))); \
       Py_DECREF(ufunc); \
   } \
   if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc, code, \
        internal::unary_op_##name<T,R>, _types, 0) < 0) { \
       /*Py_DECREF(ufunc);*/ \
       /*goto fail; \*/ \
   } \
   Py_DECREF(ufunc); \
}
126
127
128
129

  template<typename Scalar>
  void registerCommonUfunc()
  {
130
131
    const int code = Register::getTypeCode<Scalar>();
  
132
133
134
135
136
137
138
139
140
141
142
143
    PyObject* numpy_str;
#if PY_MAJOR_VERSION >= 3
    numpy_str = PyUnicode_FromString("numpy");
#else
    numpy_str = PyString_FromString("numpy");
#endif
    PyObject* numpy;
    numpy = PyImport_Import(numpy_str);
    Py_DECREF(numpy_str);
    
    import_ufunc();

144
    // Binary operators
Justin Carpentier's avatar
Justin Carpentier committed
145
146
147
148
    EIGENPY_REGISTER_BINARY_UFUNC(add,code,Scalar,Scalar,Scalar);
    EIGENPY_REGISTER_BINARY_UFUNC(subtract,code,Scalar,Scalar,Scalar);
    EIGENPY_REGISTER_BINARY_UFUNC(multiply,code,Scalar,Scalar,Scalar);
    EIGENPY_REGISTER_BINARY_UFUNC(divide,code,Scalar,Scalar,Scalar);
149
150
  
    // Comparison operators
Justin Carpentier's avatar
Justin Carpentier committed
151
152
153
154
155
156
    EIGENPY_REGISTER_BINARY_UFUNC(equal,code,Scalar,Scalar,bool);
    EIGENPY_REGISTER_BINARY_UFUNC(not_equal,code,Scalar,Scalar,bool);
    EIGENPY_REGISTER_BINARY_UFUNC(greater,code,Scalar,Scalar,bool);
    EIGENPY_REGISTER_BINARY_UFUNC(less,code,Scalar,Scalar,bool);
    EIGENPY_REGISTER_BINARY_UFUNC(greater_equal,code,Scalar,Scalar,bool);
    EIGENPY_REGISTER_BINARY_UFUNC(less_equal,code,Scalar,Scalar,bool);
157
158
159
  
    // Unary operators
    EIGENPY_REGISTER_UNARY_UFUNC(negative,code,Scalar,Scalar);
160
161
162
163
164
165
166

    Py_DECREF(numpy);
  }
  
} // namespace eigenpy

#endif // __eigenpy_ufunc_hpp__