numpy-allocator.hpp 3.88 KB
Newer Older
Justin Carpentier's avatar
Justin Carpentier committed
1
2
3
4
5
6
7
8
9
10
11
/*
 * Copyright 2020 INRIA
 */

#ifndef __eigenpy_numpy_allocator_hpp__
#define __eigenpy_numpy_allocator_hpp__

#include "eigenpy/fwd.hpp"
#include "eigenpy/numpy-type.hpp"
#include "eigenpy/eigen-allocator.hpp"

12
13
#include "eigenpy/user-type.hpp"

Justin Carpentier's avatar
Justin Carpentier committed
14
15
namespace eigenpy
{
16

Justin Carpentier's avatar
Justin Carpentier committed
17
18
19
20
21
22
23
24
25
  template<typename MatType>
  struct NumpyAllocator
  {
    template<typename SimilarMatrixType>
    static PyArrayObject * allocate(const Eigen::MatrixBase<SimilarMatrixType> & mat,
                                    npy_intp nd, npy_intp * shape)
    {
      typedef typename SimilarMatrixType::Scalar Scalar;
      
26
27
      const int code = Register::getTypeCode<Scalar>();
      PyArrayObject * pyArray = (PyArrayObject*) call_PyArray_SimpleNew(static_cast<int>(nd), shape, code);
Justin Carpentier's avatar
Justin Carpentier committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
      
      // Copy data
      EigenAllocator<SimilarMatrixType>::copy(mat,pyArray);
      
      return pyArray;
    }
  };

  template<typename MatType>
  struct NumpyAllocator<MatType &>
  {
    template<typename SimilarMatrixType>
    static PyArrayObject * allocate(Eigen::PlainObjectBase<SimilarMatrixType> & mat,
                                    npy_intp nd, npy_intp * shape)
    {
      typedef typename SimilarMatrixType::Scalar Scalar;
      enum { NPY_ARRAY_MEMORY_CONTIGUOUS = SimilarMatrixType::IsRowMajor ? NPY_ARRAY_CARRAY : NPY_ARRAY_FARRAY };
      
46
47
      if(NumpyType::sharedMemory())
      {
48
49
        const int Scalar_type_code = Register::getTypeCode<Scalar>();
        PyArrayObject * pyArray = (PyArrayObject*) call_PyArray_New(getPyArrayType(),
50
51
                                                                    static_cast<int>(nd),
                                                                    shape,
52
                                                                    Scalar_type_code,
53
54
55
56
57
58
59
60
61
                                                                    mat.data(),
                                                                    NPY_ARRAY_MEMORY_CONTIGUOUS | NPY_ARRAY_ALIGNED);
        
        return pyArray;
      }
      else
      {
        return NumpyAllocator<MatType>::allocate(mat.derived(),nd,shape);
      }
Justin Carpentier's avatar
Justin Carpentier committed
62
63
64
    }
  };

65
66
#if EIGEN_VERSION_AT_LEAST(3,2,0)

67
68
  template<typename MatType, int Options, typename Stride>
  struct NumpyAllocator<Eigen::Ref<MatType,Options,Stride> > : NumpyAllocator<MatType &>
69
70
71
72
73
  {
  };

#endif

Justin Carpentier's avatar
Justin Carpentier committed
74
75
76
77
78
79
80
81
82
83
  template<typename MatType>
  struct NumpyAllocator<const MatType &>
  {
    template<typename SimilarMatrixType>
    static PyArrayObject * allocate(const Eigen::PlainObjectBase<SimilarMatrixType> & mat,
                                    npy_intp nd, npy_intp * shape)
    {
      typedef typename SimilarMatrixType::Scalar Scalar;
      enum { NPY_ARRAY_MEMORY_CONTIGUOUS_RO = SimilarMatrixType::IsRowMajor ? NPY_ARRAY_CARRAY_RO : NPY_ARRAY_FARRAY_RO };
      
84
85
      if(NumpyType::sharedMemory())
      {
86
87
        const int Scalar_type_code = Register::getTypeCode<Scalar>();
        PyArrayObject * pyArray = (PyArrayObject*) call_PyArray_New(getPyArrayType(),
88
89
                                                                    static_cast<int>(nd),
                                                                    shape,
90
                                                                    Scalar_type_code,
91
92
93
94
95
96
97
98
99
                                                                    const_cast<SimilarMatrixType &>(mat.derived()).data(),
                                                                    NPY_ARRAY_MEMORY_CONTIGUOUS_RO | NPY_ARRAY_ALIGNED);
                                                                    
        return pyArray;
      }
      else
      {
        return NumpyAllocator<MatType>::allocate(mat.derived(),nd,shape);
      }
Justin Carpentier's avatar
Justin Carpentier committed
100
101
    }
  };
102
103
104

#if EIGEN_VERSION_AT_LEAST(3,2,0)

105
106
  template<typename MatType, int Options, typename Stride>
  struct NumpyAllocator<const Eigen::Ref<const MatType,Options,Stride> > : NumpyAllocator<const MatType &>
107
108
109
110
  {
  };

#endif
Justin Carpentier's avatar
Justin Carpentier committed
111
112
113
}

#endif // ifndef __eigenpy_numpy_allocator_hpp__