eigenpy 3.12.0
Bindings between Numpy and Eigen using Boost.Python
Loading...
Searching...
No Matches
ufunc.hpp
1//
2// Copyright (c) 2020-2025 INRIA
3// code aptapted from
4// https://github.com/numpy/numpy/blob/41977b24ae011a51f64faa75cb524c7350fdedd9/numpy/core/src/umath/_rational_tests.c.src
5//
6
7#ifndef __eigenpy_ufunc_hpp__
8#define __eigenpy_ufunc_hpp__
9
10#include "eigenpy/register.hpp"
11#include "eigenpy/user-type.hpp"
12#include "eigenpy/utils/python-compat.hpp"
13
14namespace eigenpy {
15namespace internal {
16
17#ifdef NPY_1_19_API_VERSION
18#define EIGENPY_NPY_CONST_UFUNC_ARG const
19#else
20#define EIGENPY_NPY_CONST_UFUNC_ARG
21#endif
22
23template <typename T>
24void matrix_multiply(char **args, npy_intp const *dimensions,
25 npy_intp const *steps) {
26 /* pointers to data for input and output arrays */
27 char *ip1 = args[0];
28 char *ip2 = args[1];
29 char *op = args[2];
30
31 /* lengths of core dimensions */
32 npy_intp dm = dimensions[0];
33 npy_intp dn = dimensions[1];
34 npy_intp dp = dimensions[2];
35
36 /* striding over core dimensions */
37 npy_intp is1_m = steps[0];
38 npy_intp is1_n = steps[1];
39 npy_intp is2_n = steps[2];
40 npy_intp is2_p = steps[3];
41 npy_intp os_m = steps[4];
42 npy_intp os_p = steps[5];
43
44 /* core dimensions counters */
45 npy_intp m, p;
46
47 /* calculate dot product for each row/column vector pair */
48 for (m = 0; m < dm; m++) {
49 for (p = 0; p < dp; p++) {
50 SpecialMethods<T>::dotfunc(ip1, is1_n, ip2, is2_n, op, dn, NULL);
51
52 /* advance to next column of 2nd input array and output array */
53 ip2 += is2_p;
54 op += os_p;
55 }
56
57 /* reset to first column of 2nd input array and output array */
58 ip2 -= is2_p * p;
59 op -= os_p * p;
60
61 /* advance to next row of 1st input array and output array */
62 ip1 += is1_m;
63 op += os_m;
64 }
65}
66
67template <typename T>
68void gufunc_matrix_multiply(char **args,
69 npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *dimensions,
70 npy_intp EIGENPY_NPY_CONST_UFUNC_ARG *steps,
71 void *NPY_UNUSED(func)) {
72 /* outer dimensions counter */
73 npy_intp N_;
74
75 /* length of flattened outer dimensions */
76 npy_intp dN = dimensions[0];
77
78 /* striding over flattened outer dimensions for input and output arrays */
79 npy_intp s0 = steps[0];
80 npy_intp s1 = steps[1];
81 npy_intp s2 = steps[2];
82
83 /*
84 * loop through outer dimensions, performing matrix multiply on
85 * core dimensions for each loop
86 */
87 for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {
88 matrix_multiply<T>(args, dimensions + 1, steps + 3);
89 }
90}
91
92#define EIGENPY_REGISTER_BINARY_OPERATOR(name, op) \
93 template <typename T1, typename T2, typename R> \
94 void binary_op_##name( \
95 char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
96 EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
97 npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions; \
98 char *i0 = args[0], *i1 = args[1], *o = args[2]; \
99 int k; \
100 for (k = 0; k < n; k++) { \
101 T1 &x = *static_cast<T1 *>(static_cast<void *>(i0)); \
102 T2 &y = *static_cast<T2 *>(static_cast<void *>(i1)); \
103 R &res = *static_cast<R *>(static_cast<void *>(o)); \
104 res = x op y; \
105 i0 += is0; \
106 i1 += is1; \
107 o += os; \
108 } \
109 } \
110 \
111 template <typename T> \
112 void binary_op_##name( \
113 char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
114 EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
115 binary_op_##name<T, T, T>(args, dimensions, steps, data); \
116 }
117
118EIGENPY_REGISTER_BINARY_OPERATOR(add, +)
119EIGENPY_REGISTER_BINARY_OPERATOR(subtract, -)
120EIGENPY_REGISTER_BINARY_OPERATOR(multiply, *)
121EIGENPY_REGISTER_BINARY_OPERATOR(divide, /)
122EIGENPY_REGISTER_BINARY_OPERATOR(equal, ==)
123EIGENPY_REGISTER_BINARY_OPERATOR(not_equal, !=)
124EIGENPY_REGISTER_BINARY_OPERATOR(less, <)
125EIGENPY_REGISTER_BINARY_OPERATOR(greater, >)
126EIGENPY_REGISTER_BINARY_OPERATOR(less_equal, <=)
127EIGENPY_REGISTER_BINARY_OPERATOR(greater_equal, >=)
128
129#define EIGENPY_REGISTER_UNARY_OPERATOR(name, op) \
130 template <typename T, typename R> \
131 void unary_op_##name( \
132 char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
133 EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void * /*data*/) { \
134 npy_intp is = steps[0], os = steps[1], n = *dimensions; \
135 char *i = args[0], *o = args[1]; \
136 int k; \
137 for (k = 0; k < n; k++) { \
138 T &x = *static_cast<T *>(static_cast<void *>(i)); \
139 R &res = *static_cast<R *>(static_cast<void *>(o)); \
140 res = op x; \
141 i += is; \
142 o += os; \
143 } \
144 } \
145 \
146 template <typename T> \
147 void unary_op_##name( \
148 char **args, EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *dimensions, \
149 EIGENPY_NPY_CONST_UFUNC_ARG npy_intp *steps, void *data) { \
150 unary_op_##name<T, T>(args, dimensions, steps, data); \
151 }
152
153EIGENPY_REGISTER_UNARY_OPERATOR(negative, -)
154EIGENPY_REGISTER_UNARY_OPERATOR(square, x *)
155
156} // namespace internal
157
158#define EIGENPY_REGISTER_BINARY_UFUNC(name, code, T1, T2, R) \
159 { \
160 PyUFuncObject *ufunc = \
161 (PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
162 int _types[3] = {Register::getTypeCode<T1>(), Register::getTypeCode<T2>(), \
163 Register::getTypeCode<R>()}; \
164 if (!ufunc) { \
165 /*goto fail; \*/ \
166 } \
167 if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
168 PyErr_Format(PyExc_AssertionError, \
169 "ufunc %s takes %d arguments, our loop takes %lu", #name, \
170 ufunc->nargs, \
171 (unsigned long)(sizeof(_types) / sizeof(int))); \
172 Py_DECREF(ufunc); \
173 } \
174 if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
175 internal::binary_op_##name<T1, T2, R>, \
176 _types, 0) < 0) { \
177 /*Py_DECREF(ufunc);*/ \
178 /*goto fail; \*/ \
179 } \
180 Py_DECREF(ufunc); \
181 }
182
183#define EIGENPY_REGISTER_UNARY_UFUNC(name, code, T, R) \
184 { \
185 PyUFuncObject *ufunc = \
186 (PyUFuncObject *)PyObject_GetAttrString(numpy, #name); \
187 int _types[2] = {Register::getTypeCode<T>(), Register::getTypeCode<R>()}; \
188 if (!ufunc) { \
189 /*goto fail; \*/ \
190 } \
191 if (sizeof(_types) / sizeof(int) != ufunc->nargs) { \
192 PyErr_Format(PyExc_AssertionError, \
193 "ufunc %s takes %d arguments, our loop takes %lu", #name, \
194 ufunc->nargs, \
195 (unsigned long)(sizeof(_types) / sizeof(int))); \
196 Py_DECREF(ufunc); \
197 } \
198 if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, code, \
199 internal::unary_op_##name<T, R>, _types, \
200 0) < 0) { \
201 /*Py_DECREF(ufunc);*/ \
202 /*goto fail; \*/ \
203 } \
204 Py_DECREF(ufunc); \
205 }
206
207template <typename Scalar>
208void registerCommonUfunc() {
209 const int type_code = Register::getTypeCode<Scalar>();
210
211 PyObject *numpy_str;
212 numpy_str = PyStr_FromString("numpy");
213 PyObject *numpy;
214 numpy = PyImport_Import(numpy_str);
215 Py_DECREF(numpy_str);
216
217 import_ufunc();
218
219 // Matrix multiply
220 {
221 int types[3] = {type_code, type_code, type_code};
222
223 std::stringstream ss;
224 ss << "return result of multiplying two matrices of ";
225 ss << bp::type_info(typeid(Scalar)).name();
226 PyUFuncObject *ufunc =
227 (PyUFuncObject *)PyObject_GetAttrString(numpy, "matmul");
228 if (!ufunc) {
229 std::stringstream ss;
230 ss << "Impossible to define matrix_multiply for given type "
231 << bp::type_info(typeid(Scalar)).name() << std::endl;
232 eigenpy::Exception(ss.str());
233 }
234 if (PyUFunc_RegisterLoopForType((PyUFuncObject *)ufunc, type_code,
235 &internal::gufunc_matrix_multiply<Scalar>,
236 types, 0) < 0) {
237 std::stringstream ss;
238 ss << "Impossible to register matrix_multiply for given type "
239 << bp::type_info(typeid(Scalar)).name() << std::endl;
240 eigenpy::Exception(ss.str());
241 }
242
243 Py_DECREF(ufunc);
244 }
245
246 // Binary operators
247 EIGENPY_REGISTER_BINARY_UFUNC(add, type_code, Scalar, Scalar, Scalar);
248 EIGENPY_REGISTER_BINARY_UFUNC(subtract, type_code, Scalar, Scalar, Scalar);
249 EIGENPY_REGISTER_BINARY_UFUNC(multiply, type_code, Scalar, Scalar, Scalar);
250 EIGENPY_REGISTER_BINARY_UFUNC(divide, type_code, Scalar, Scalar, Scalar);
251
252 // Comparison operators
253 EIGENPY_REGISTER_BINARY_UFUNC(equal, type_code, Scalar, Scalar, bool);
254 EIGENPY_REGISTER_BINARY_UFUNC(not_equal, type_code, Scalar, Scalar, bool);
255 EIGENPY_REGISTER_BINARY_UFUNC(greater, type_code, Scalar, Scalar, bool);
256 EIGENPY_REGISTER_BINARY_UFUNC(less, type_code, Scalar, Scalar, bool);
257 EIGENPY_REGISTER_BINARY_UFUNC(greater_equal, type_code, Scalar, Scalar, bool);
258 EIGENPY_REGISTER_BINARY_UFUNC(less_equal, type_code, Scalar, Scalar, bool);
259
260 // Unary operators
261 EIGENPY_REGISTER_UNARY_UFUNC(negative, type_code, Scalar, Scalar);
262 EIGENPY_REGISTER_UNARY_UFUNC(square, type_code, Scalar, Scalar);
263
264 Py_DECREF(numpy);
265}
266
267} // namespace eigenpy
268
269#endif // __eigenpy_ufunc_hpp__