Python C/C++ 拡張: データの一部を np.array として Python に渡す (stride を設定)

C/C++ のデータを np.array として Python から見る。

typedef struct Particle {
  float x[3], v[3];
} Particle;

const int np=10;
Particle* const p= calloc(sizeof(Particle), np);

for(int i=0; i<np; ++i) {
  p[i].x[0]= (float) i + 1;
}

これを np.array として Python に渡したい場合は PyArray_SimpleNewFromData 関数を使えば良い。¹

しかし、ここでは、データの一部、x[3] だけを np.array としてPythonに渡す方法を紹介する。データ構造が複数の型を含んでいて、全てをひとつの np.array にできない場合などもあるでしょう。

このようにデータがメモリに連続して存在しないけれど、一定間隔で並んでいる場合は strides を設定すればよい。²

PyObject* PyArray_New(PyTypeObject* subtype, int nd, npy_intp* dims, 
			           int type_num, npy_intp* strides, void* data, 
			           int itemsize, int flags, PyObject* obj);

strides[] とは a を np.array とするとき、strides[0] が a[i,j] と a[i+1,j] のバイト差、strides[1] が a[i,j] と a[i,j+1] のバイト差。

static PyObject* py_as_nparray(PyObject *self, PyObject *args)
{
  const int nd=2;
  const int ncol= 3;
  npy_intp dims[]= {np, ncol};
  npy_intp strides[]= {sizeof(Particle), sizeof(float)};

  return PyArray_New(&PyArray_Type, nd, dims, NPY_FLOAT, strides, p->x, 0, 0, 0);
}

Python から使った結果

import numpy as np
import cext08

a = cext08.as_nparray()

[[  1.   0.   0.]
 [  2.   0.   0.]
 [  3.   0.   0.]
 [  4.   0.   0.]
 ...

＊ 実際に使うときは Particle* p のメモリは Python で使っている間は解放されず、解放する場合は C のほうでなんとかするとします。

コード全体は github にあります

1. http://qiita.com/junkoda/items/17df11d7a20dc9d50e7d ↩

2. Array API http://docs.scipy.org/doc/numpy-1.10.1/reference/c-api.array.html#array-structure-and-data-access ↩