Line data Source code
1 : /*
2 : * Copyright (c) 2017 Sebastian Weber, Henri Menke. All rights reserved.
3 : *
4 : * This file is part of the pairinteraction library.
5 : *
6 : * The pairinteraction library is free software: you can redistribute it and/or modify
7 : * it under the terms of the GNU Lesser General Public License as published by
8 : * the Free Software Foundation, either version 3 of the License, or
9 : * (at your option) any later version.
10 : *
11 : * The pairinteraction library is distributed in the hope that it will be useful,
12 : * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 : * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 : * GNU Lesser General Public License for more details.
15 : *
16 : * You should have received a copy of the GNU Lesser General Public License
17 : * along with the pairinteraction library. If not, see <http://www.gnu.org/licenses/>.
18 : */
19 :
20 : #ifndef NUMPYUTILS_H
21 : #define NUMPYUTILS_H
22 :
23 : #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
24 : #include <Python.h>
25 : #include <numpy/ndarrayobject.h>
26 :
27 : #include <cstring>
28 : #include <iterator>
29 : #include <memory>
30 : #include <numeric>
31 : #include <stdexcept>
32 : #include <tuple>
33 :
34 : #include "EigenCompat.hpp"
35 : #include <Eigen/Core>
36 : #include <Eigen/SparseCore>
37 :
38 : #include "Traits.hpp"
39 :
40 : namespace numpy {
41 :
42 : /** \brief Selector for view or copy
43 : *
44 : * Deduction for view and copy should be a template parameter for
45 : * maximum flexibility. That is why we make a new type for that.
46 : */
47 : enum view_or_copy { view, copy };
48 :
49 : /** \brief Numpy array type
50 : *
51 : * Numpy arrays are just usual PyObjects. However, it doesn't
52 : * look as nice in return statements as numpy::array.
53 : */
54 : typedef PyObject *array;
55 :
56 : namespace internal {
57 :
58 : /** \brief Map C++ types to Numpy types
59 : *
60 : * This struct has specializations for common C++ types and
61 : * their Numpy counterparts.
62 : */
63 : template <typename T>
64 : struct py_type;
65 :
66 : /** \brief Specialization of py_type for int */
67 : template <>
68 : struct py_type<int> {
69 : /** \brief Numpy type identifier */
70 : static constexpr int type = NPY_INT;
71 : };
72 :
73 : /** \brief Specialization of py_type for float */
74 : template <>
75 : struct py_type<float> {
76 : /** \brief Numpy type identifier */
77 : static constexpr int type = NPY_FLOAT;
78 : };
79 :
80 : /** \brief Specialization of py_type for double */
81 : template <>
82 : struct py_type<double> {
83 : /** \brief Numpy type identifier */
84 : static constexpr int type = NPY_DOUBLE;
85 : };
86 :
87 : /** \brief Specialization of py_type for double */
88 : template <>
89 : struct py_type<std::complex<double>> {
90 : /** \brief Numpy type identifier */
91 : static constexpr int type = NPY_CDOUBLE;
92 : };
93 :
94 : /** \brief Perform sanity checks for convert
95 : *
96 : * To ensure that dimensions and types are consistent we
97 : * introduce this check function.
98 : *
99 : * \param[in] len Length of the memory
100 : * \param[in] nd Number of dimensions
101 : * \param[in] dims Initializer list with lengths of dimensions
102 : *
103 : * \return The total length of the array
104 : */
105 : template <typename T>
106 7002 : void check_array_sanity(int len, int nd, std::initializer_list<T> dims) {
107 7002 : if (len < 1)
108 0 : throw std::out_of_range("Trying to create a numpy array with zero or "
109 : "negative element count!");
110 :
111 7002 : if (nd != static_cast<int>(dims.size()))
112 0 : throw std::out_of_range("Dimension mismatch!");
113 :
114 14015 : if (len > std::accumulate(dims.begin(), dims.end(), int{1}, [](int a, int b) { return a * b; }))
115 0 : throw std::out_of_range("Requested dimension is larger than data!");
116 7002 : }
117 :
118 : /** \brief Create a Numpy view of an iterator (implementation)
119 : *
120 : * This creates a view of the data, i.e. the data still
121 : * belongs to the C++ end. When the data is deallocated on
122 : * the C++ end trying to access it in Python will probably
123 : * cause a segfault. In any case it is undefined behavior.
124 : *
125 : * \param[in] nd Number of dimensions
126 : * \param[in] dim Dimensions of the array
127 : * \param[in] dtype Value type of array elements
128 : * \param[in] data Pointer to array data
129 : *
130 : * \return PyObject* containing a Numpy array
131 : */
132 : template <view_or_copy v, bool const_tag, typename value_type>
133 : typename std::enable_if<v == numpy::view, PyObject *>::type
134 969 : convert_impl(int nd, npy_intp *dim, int dtype, void *data, int) {
135 : PyObject *ndarray =
136 969 : PyArray_New(&PyArray_Type, nd, dim, dtype, nullptr, data, 0, NPY_ARRAY_FARRAY, nullptr);
137 :
138 : if (const_tag)
139 0 : PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject *>(ndarray), NPY_ARRAY_WRITEABLE);
140 969 : return ndarray;
141 : }
142 :
143 : /** \brief Create a Numpy copy of an iterator (implementation)
144 : *
145 : * This creates a copy of the data, i.e. the data belongs to
146 : * the Python end. The array is marked with NPY_OWNDATA which
147 : * should tell the garbage collector to release this memory.
148 : *
149 : * \param[in] nd Number of dimensions
150 : * \param[in] dim Dimensions of the array
151 : * \param[in] dtype Value type of array elements
152 : * \param[in] data Pointer to array data
153 : * \param[in] len Length of data
154 : *
155 : * \return PyObject* containing a Numpy array
156 : */
157 : template <view_or_copy v, bool const_tag, typename value_type>
158 : typename std::enable_if<v == numpy::copy, PyObject *>::type
159 6033 : convert_impl(int nd, npy_intp *dim, int dtype, void *data, int len) {
160 : PyObject *ndarray =
161 6033 : PyArray_New(&PyArray_Type, nd, dim, dtype, nullptr, nullptr, 0, NPY_ARRAY_FARRAY, nullptr);
162 :
163 12066 : std::memcpy(PyArray_DATA(reinterpret_cast<PyArrayObject *>(ndarray)), data,
164 6033 : len * sizeof(value_type));
165 6033 : PyArray_ENABLEFLAGS(reinterpret_cast<PyArrayObject *>(ndarray), NPY_ARRAY_OWNDATA);
166 6033 : return ndarray;
167 : }
168 :
169 : /** \brief Create a Numpy array from an iterator (implementation)
170 : *
171 : * This is the implementation for the conversion of an
172 : * arbitrary pointer between \p begin and \p end .
173 : *
174 : * \param[in] begin Random access iterator pointing to the start
175 : * \param[in] end Random access iterator pointing to the end
176 : * \param[in] nd Number of dimensions
177 : * \param[in] dims Initializer list with lengths of dimensions
178 : *
179 : * \return PyObject* containing a Numpy array
180 : */
181 : template <view_or_copy v, typename RAIter, typename T>
182 7002 : PyObject *convert(RAIter begin, RAIter end, int nd, std::initializer_list<T> dims,
183 : std::random_access_iterator_tag) {
184 : using value_type = typename std::iterator_traits<RAIter>::value_type;
185 7002 : constexpr auto const_tag = traits::is_pointer_to_const<decltype(&(*begin))>::value;
186 :
187 7002 : int const len = std::distance(begin, end);
188 :
189 7002 : check_array_sanity(len, nd, dims);
190 :
191 : npy_intp *dim =
192 : const_cast<typename traits::pointer_remove_const<decltype(&(*dims.begin()))>::type>(
193 7002 : &(*dims.begin()));
194 7002 : auto dtype = py_type<value_type>::type;
195 7002 : void *data =
196 : const_cast<typename traits::pointer_remove_const<decltype(&(*begin))>::type>(&(*begin));
197 :
198 7002 : return convert_impl<v, const_tag, value_type>(nd, dim, dtype, data, len);
199 : }
200 : } // namespace internal
201 :
202 : // View and copy of bare pointes
203 :
204 : /** \brief Create a Numpy array from an iterator (frontend)
205 : *
206 : * This is a proxy for internal::view_impl to make sure that the
207 : * iterator passed is random access.
208 : *
209 : * \param[in] begin Iterator pointing to the start
210 : * \param[in] end Iterator pointing to the end
211 : * \param[in] nd Number of dimensions
212 : * \param[in] dims Initializer list with lengths of dimensions
213 : *
214 : * \return PyObject* containing a Numpy array
215 : */
216 : template <view_or_copy v, typename Iter, typename T>
217 7002 : PyObject *convert(Iter begin, Iter end, int nd, std::initializer_list<T> dims) {
218 7002 : return internal::convert<v>(begin, end, nd, dims,
219 7002 : typename std::iterator_traits<Iter>::iterator_category());
220 : }
221 :
222 : // Overloads for often used types
223 :
224 : /** \brief Create a Numpy array from a raw pointer
225 : *
226 : * Specialization of numpy::convert for a raw pointer.
227 : *
228 : * \param[in] begin pointer to beginning
229 : * \param[in] end pointer to end
230 : *
231 : * \return PyObject* containing a Numpy array
232 : */
233 : template <view_or_copy v, typename T>
234 969 : PyObject *convert(T *begin, T *end) {
235 969 : return numpy::convert<v>(begin, end, 1, {std::distance(begin, end)});
236 : }
237 : template <view_or_copy v, typename T>
238 : PyObject *convert(T const *begin, T const *end) {
239 : return numpy::convert<v>(begin, end, 1, {std::distance(begin, end)});
240 : }
241 :
242 : /** \brief Create a Numpy array from a Eigen::Matrix
243 : *
244 : * Specialization of numpy::convert for Eigen::Matrix
245 : *
246 : * \param[in] m Eigen::Matrix&
247 : *
248 : * \return PyObject* containing a Numpy array
249 : */
250 : template <view_or_copy v, typename T>
251 11 : PyObject *convert(Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> &m) {
252 11 : return numpy::convert<v>(m.data(), m.data() + m.size(), 2, {m.rows(), m.cols()});
253 : }
254 : template <view_or_copy v, typename T>
255 : PyObject *convert(Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> const &m) {
256 : return numpy::convert<v>(m.data(), m.data() + m.size(), 2, {m.rows(), m.cols()});
257 : }
258 :
259 : /** \brief Create a Numpy array from an Eigen::Vector
260 : *
261 : * Specialization of numpy::convert for Eigen::Vector
262 : *
263 : * \param[in] vec Eigen::Vector&
264 : *
265 : * \return PyObject* containing a Numpy array
266 : */
267 : template <view_or_copy v, typename T>
268 250 : PyObject *convert(Eigen::Matrix<T, Eigen::Dynamic, 1> &vec) {
269 250 : return numpy::convert<v>(vec.data(), vec.data() + vec.size(), 1, {vec.size()});
270 : }
271 : template <view_or_copy v, typename T>
272 0 : PyObject *convert(Eigen::Matrix<T, Eigen::Dynamic, 1> const &vec) {
273 0 : return numpy::convert<v>(vec.data(), vec.data() + vec.size(), 1, {vec.size()});
274 : }
275 :
276 : /** \brief Create a Numpy array from a std::vector
277 : *
278 : * Specialization of numpy::convert for std::vector
279 : *
280 : * \param[in] vec std::vector&
281 : *
282 : * \return PyObject* containing a Numpy array
283 : */
284 : template <view_or_copy v, typename T>
285 : PyObject *convert(std::vector<T> &vec) {
286 : return numpy::convert<v>(vec.data(), vec.data() + vec.size(), 1, {vec.size()});
287 : }
288 : template <view_or_copy v, typename T>
289 : PyObject *convert(std::vector<T> const &vec) {
290 : return numpy::convert<v>(vec.data(), vec.data() + vec.size(), 1, {vec.size()});
291 : }
292 :
293 : /** \brief Create a Numpy array from a std::array
294 : *
295 : * Specialization of numpy::convert for std::array
296 : *
297 : * \param[in] vec std::array&
298 : *
299 : * \return PyObject* containing a Numpy array
300 : */
301 : template <view_or_copy v, typename T, size_t N>
302 5772 : PyObject *convert(std::array<T, N> &vec) {
303 : using ssize_t = std::make_signed<size_t>::type;
304 5772 : return numpy::convert<v>(vec.data(), vec.data() + N, 1, {ssize_t{N}});
305 : }
306 : template <view_or_copy v, typename T, size_t N>
307 : PyObject *convert(std::array<T, N> const &vec) {
308 : using ssize_t = std::make_signed<size_t>::type;
309 : return numpy::convert<v>(vec.data(), vec.data() + N, 1, {ssize_t{N}});
310 : }
311 :
312 : // Experimental support for Sparse matrix return types
313 :
314 : namespace internal {
315 :
316 : /** \brief Create a Numpy array from an Eigen::SparseMatrix
317 : *
318 : * This is the implementation for the conversion of an
319 : * Eigen::SparseMatrix to a scipy.sparse.csc_matrix. We load
320 : * the scipy.sparse module and call the constructor for the
321 : * csc_matrix which we populate with numpy arrays.
322 : *
323 : * \param[in] sm T&&
324 : *
325 : * \return PyObject* containing a Scipy csc_matrix
326 : */
327 : template <view_or_copy v, typename T>
328 323 : PyObject *sparse_impl(T &&sm) {
329 323 : if (!sm.isCompressed())
330 : // we cannot call makeCompressed here because sm might be const
331 0 : throw std::runtime_error("Sparse matrix is not compressed!");
332 :
333 : numpy::array indptr =
334 323 : numpy::convert<v>(sm.outerIndexPtr(), sm.outerIndexPtr() + sm.outerSize() + 1);
335 :
336 : numpy::array indices =
337 323 : numpy::convert<v>(sm.innerIndexPtr(), sm.innerIndexPtr() + sm.nonZeros());
338 :
339 323 : numpy::array data = numpy::convert<v>(sm.valuePtr(), sm.valuePtr() + sm.nonZeros());
340 :
341 323 : int num_rows = sm.rows();
342 323 : int num_cols = sm.cols();
343 :
344 323 : char object[] = "csc_matrix";
345 323 : char arglist[] = "(OOO)(ii)";
346 646 : std::unique_ptr<PyObject, decltype(&Py_DecRef)> scipy{PyImport_ImportModule("scipy.sparse"),
347 323 : &Py_DecRef};
348 323 : if (scipy == nullptr) {
349 0 : PyErr_Print();
350 0 : return nullptr;
351 : }
352 323 : PyObject *mat = PyObject_CallMethod(scipy.get(), object, arglist, data, indices, indptr,
353 : num_rows, num_cols);
354 323 : if (mat == nullptr) {
355 0 : PyErr_Print();
356 0 : return nullptr;
357 : }
358 323 : return mat;
359 : }
360 :
361 : } // namespace internal
362 :
363 : /** \brief Create a Numpy array from an Eigen::SparseMatrix
364 : *
365 : * Specialization of numpy::convert for Eigen::SparseMatrix
366 : *
367 : * \param[in] sm Eigen::SparseMatrix&
368 : *
369 : * \return PyObject* containing a Numpy array
370 : */
371 : template <view_or_copy v, typename T>
372 323 : PyObject *convert(Eigen::SparseMatrix<T> &sm) {
373 323 : return internal::sparse_impl<v>(sm);
374 : }
375 : template <view_or_copy v, typename T>
376 : PyObject *convert(Eigen::SparseMatrix<T> const &sm) {
377 : return internal::sparse_impl<v>(sm);
378 : }
379 :
380 : // Experimental support for Numpy arguments
381 :
382 : namespace internal {
383 :
384 : /** \brief Check if a Numpy has given storage order and alignment
385 : *
386 : * To be convertible to an Eigen matrix type the memory of a
387 : * Numpy array has to be FORTRAN style contiguous and has to
388 : * be aligned.
389 : *
390 : * \param[in] ndarray Numpy array
391 : */
392 0 : void check_order_and_alignment(PyArrayObject *ndarray) {
393 0 : int const flags = PyArray_FLAGS(ndarray);
394 0 : if ((flags & NPY_ARRAY_F_CONTIGUOUS) == 0)
395 0 : throw std::invalid_argument("The argument is not contiguous or has wrong storage order!");
396 0 : if ((flags & NPY_ARRAY_ALIGNED) == 0)
397 0 : throw std::invalid_argument("The argument is not not aligned!");
398 0 : }
399 :
400 : /** \brief Deconstruct a numpy array for conversion to dense Eigen type
401 : *
402 : * \param[in] ndarray_ Numpy array
403 : *
404 : * \return A tuple of the number of dimensions, an array of
405 : * dimensions, and a pointer to the data.
406 : */
407 : template <typename T, typename Scalar = typename T::Scalar>
408 : std::tuple<int, npy_intp *, Scalar *> as_dense_impl(numpy::array ndarray_) {
409 : if (!ndarray_ || !PyArray_Check(ndarray_))
410 : throw std::invalid_argument("The argument is not a valid Numpy array!");
411 :
412 : PyArrayObject *ndarray = reinterpret_cast<PyArrayObject *>(ndarray_);
413 : internal::check_order_and_alignment(ndarray);
414 : int const nd = PyArray_NDIM(ndarray);
415 : npy_intp *dims = PyArray_SHAPE(ndarray);
416 : Scalar *data = reinterpret_cast<Scalar *>(PyArray_DATA(ndarray));
417 :
418 : if (PyArray_TYPE(ndarray) != internal::py_type<Scalar>::type)
419 : throw std::invalid_argument("Type mismatch.");
420 :
421 : return std::make_tuple(nd, dims, data);
422 : }
423 :
424 : } // namespace internal
425 :
426 : /** \brief Convert a Numpy array to an Eigen::Matrix
427 : *
428 : * \warning Experimental! For now the returned Eigen::Map is
429 : * read-only.
430 : *
431 : * \param[in] ndarray Numpy array
432 : *
433 : * \return An Eigen::Map of the Numpy array memory
434 : */
435 : #ifndef SCANNED_BY_DOXYGEN
436 : template <typename T, typename Scalar = typename T::Scalar>
437 : typename std::enable_if<
438 : std::is_same<T,
439 : Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, T::Options,
440 : T::MaxRowsAtCompileTime, T::MaxColsAtCompileTime>>::value,
441 : Eigen::Map<T const>>::type
442 : #else
443 : template <typename MatrixType>
444 : Eigen::Map<MatrixType>
445 : #endif
446 : as(numpy::array ndarray) {
447 : int nd;
448 : npy_intp *dims;
449 : Scalar *data;
450 : std::tie(nd, dims, data) = internal::as_dense_impl<T>(ndarray);
451 :
452 : if (nd > 2)
453 : throw std::range_error("Dimension mismatch.");
454 :
455 : return Eigen::Map<T const>(data, dims[0], (nd == 1 ? 1 : dims[1]));
456 : }
457 :
458 : /** \brief Convert a Numpy array to an Eigen::Vector
459 : *
460 : * \warning Experimental! For now the returned Eigen::Map is
461 : * read-only.
462 : *
463 : * \param[in] ndarray Numpy array
464 : *
465 : * \return An Eigen::Map of the Numpy array memory
466 : */
467 : #ifndef SCANNED_BY_DOXYGEN
468 : template <typename T, typename Scalar = typename T::Scalar>
469 : typename std::enable_if<
470 : std::is_same<T,
471 : Eigen::Matrix<Scalar, Eigen::Dynamic, 1, T::Options, T::MaxRowsAtCompileTime,
472 : T::MaxColsAtCompileTime>>::value,
473 : Eigen::Map<T const>>::type
474 : #else
475 : template <typename VectorType>
476 : Eigen::Map<VectorType>
477 : #endif
478 : as(numpy::array ndarray) {
479 : int nd;
480 : npy_intp *dims;
481 : Scalar *data;
482 : std::tie(nd, dims, data) = internal::as_dense_impl<T>(ndarray);
483 :
484 : if (nd > 1)
485 : throw std::range_error("Dimension mismatch.");
486 :
487 : return Eigen::Map<T const>(data, dims[0]);
488 : }
489 :
490 : /** \brief Convert a Numpy array to an Eigen::SparseMatrix
491 : *
492 : * \warning Experimental! For now the returned Eigen::Map is
493 : * read-only.
494 : *
495 : * \param[in] ndarray Numpy array
496 : *
497 : * \return An Eigen::Map of the Numpy array memory
498 : */
499 : #ifndef SCANNED_BY_DOXYGEN
500 : template <typename T, typename Scalar = typename T::Scalar,
501 : typename StorageIndex = typename T::StorageIndex>
502 : typename std::enable_if<
503 : std::is_same<T, Eigen::SparseMatrix<Scalar, T::Options, StorageIndex>>::value,
504 : Eigen::Map<T const>>::type
505 : #else
506 : template <typename SparseMatrixType>
507 : Eigen::Map<SparseMatrixType>
508 : #endif
509 : as(numpy::array ndarray) {
510 : auto make_unique_object = [](PyObject *o) {
511 : return std::unique_ptr<PyObject, decltype(&Py_DecRef)>{o, &Py_DecRef};
512 : };
513 : auto scipy = make_unique_object(PyImport_ImportModule("scipy.sparse"));
514 : if (scipy == nullptr) {
515 : PyErr_Print();
516 : return nullptr;
517 : }
518 : auto csc_matrix = make_unique_object(PyObject_GetAttrString(scipy.get(), "csc_matrix"));
519 : if (csc_matrix == nullptr) {
520 : PyErr_Print();
521 : return nullptr;
522 : }
523 : if (!ndarray || !PyObject_IsInstance(ndarray, csc_matrix.get()))
524 : throw std::invalid_argument("The argument is not a valid csc_matrix!");
525 :
526 : auto data_obj = make_unique_object(PyObject_GetAttrString(ndarray, "data"));
527 : PyArrayObject *data_ = reinterpret_cast<PyArrayObject *>(data_obj.get());
528 : internal::check_order_and_alignment(data_);
529 :
530 : auto indices_obj = make_unique_object(PyObject_GetAttrString(ndarray, "indices"));
531 : PyArrayObject *indices_ = reinterpret_cast<PyArrayObject *>(indices_obj.get());
532 : internal::check_order_and_alignment(indices_);
533 :
534 : auto indptr_obj = make_unique_object(PyObject_GetAttrString(ndarray, "indptr"));
535 : PyArrayObject *indptr_ = reinterpret_cast<PyArrayObject *>(indptr_obj.get());
536 : internal::check_order_and_alignment(indptr_);
537 :
538 : auto shape_obj = make_unique_object(PyObject_GetAttrString(ndarray, "shape"));
539 : int rows, cols;
540 : PyArg_ParseTuple(shape_obj.get(), "ii", &rows, &cols);
541 :
542 : npy_intp *nnz = PyArray_SHAPE(data_);
543 :
544 : Scalar *data = reinterpret_cast<Scalar *>(PyArray_DATA(data_));
545 : StorageIndex *indices = reinterpret_cast<StorageIndex *>(PyArray_DATA(indices_));
546 : StorageIndex *indptr = reinterpret_cast<StorageIndex *>(PyArray_DATA(indptr_));
547 :
548 : if (PyArray_TYPE(data_) != internal::py_type<Scalar>::type)
549 : throw std::invalid_argument("Type mismatch.");
550 :
551 : return Eigen::Map<T const>(rows, cols, *nnz, indptr, indices, data);
552 : }
553 :
554 : } // namespace numpy
555 :
556 : #endif // NUMPYUTILS_H
|