NumPyで多次元配列をイテレーション？ char **NpyIter_GetDataPtrArray() 関数で効率的に処理しよう

NpyIter_GetDataPtrArray() は、NpyIter 構造体を受け取り、その構造体内に含まれる全ての配列のデータポインタを `char** 型の配列として返します。各要素は、対応する配列の最初の要素へのポインタとなります。

詳細説明

- iter: 操作対象の NpyIter 構造体へのポインタ
戻り値
- 成功した場合: 各配列のデータポインタを要素とする `char** 型の配列
- 失敗した場合: NULL
注意事項
- 返される char** 型の配列は、NpyIter_DECREF()` を呼び出すまで有効です。
- 各配列のデータポインタは、その配列のデータ型に対応する型にキャストする必要があります。
- 複数のスレッドから NpyIter 構造体へ同時にアクセスする場合は、適切な同期メカニズムを実装する必要があります。

例

#include <numpy/npyiter.h>

int main() {
  // NumPy 配列を 2 つ作成
  npy_intp dims[] = {2, 3};
  npy_intp strides[] = {sizeof(int), sizeof(int) * 3};
  int arr1[2][3] = {{1, 2, 3}, {4, 5, 6}};
  int arr2[2][3] = {{7, 8, 9}, {10, 11, 12}};
  PyArrayObject *array1 = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr1, NPY_ORDER_C, NULL, NULL);
  PyArrayObject *array2 = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr2, NPY_ORDER_C, NULL, NULL);

  // NumPy イテレータを作成
  NpyIter *iter = NpyIter_New(2, &array1, &array2, NPY_ITER_MULTI_LOOP);

  // 各配列のデータポインタを取得
  char **data_ptrs = NpyIter_GetDataPtrArray(iter);

  // イテレータ内の各要素を処理
  for (npy_intp i = 0; i < NpyIter_GetIterSize(iter); ++i) {
    int *val1 = (int *)data_ptrs[0];
    int *val2 = (int *)data_ptrs[1];

    printf("arr1[%d][%d] = %d, arr2[%d][%d] = %d\n",
           NpyIter_GetIterIndex(iter, 0), NpyIter_GetIterIndex(iter, 1), *val1,
           NpyIter_GetIterIndex(iter, 0), NpyIter_GetIterIndex(iter, 1), *val2);

    // 次の要素へ移動
    NpyIter_Next(iter);
    data_ptrs += NpyIter_GetInnerStrideArray(iter);
  }

  // NumPy イテレータと配列を解放
  NpyIter_DECREF(iter);
  Py_DECREF(array1);
  Py_DECREF(array2);

  return 0;
}

例：2つのNumPy配列の要素を足し合わせる

#include <numpy/npyiter.h>

int main() {
  // NumPy 配列を 2 つ作成
  npy_intp dims[] = {2, 3};
  npy_intp strides[] = {sizeof(int), sizeof(int) * 3};
  int arr1[2][3] = {{1, 2, 3}, {4, 5, 6}};
  int arr2[2][3] = {{7, 8, 9}, {10, 11, 12}};
  PyArrayObject *array1 = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr1, NPY_ORDER_C, NULL, NULL);
  PyArrayObject *array2 = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr2, NPY_ORDER_C, NULL, NULL);

  // NumPy イテレータを作成
  NpyIter *iter = NpyIter_New(2, &array1, &array2, NPY_ITER_MULTI_LOOP | NPY_ITER_READONLY);

  // 結果を格納する配列を作成
  int result[2][3];
  memset(result, 0, sizeof(result));

  // 各要素を足し合わせる
  for (npy_intp i = 0; i < NpyIter_GetIterSize(iter); ++i) {
    int *val1 = (int *)NpyIter_GetDataPtrArray(iter)[0];
    int *val2 = (int *)NpyIter_GetDataPtrArray(iter)[1];
    int *res = &result[NpyIter_GetIterIndex(iter, 0)][NpyIter_GetIterIndex(iter, 1)];

    *res = *val1 + *val2;

    // 次の要素へ移動
    NpyIter_Next(iter);
  }

  // 結果を出力
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      printf("result[%d][%d] = %d\n", i, j, result[i][j]);
    }
  }

  // NumPy イテレータと配列を解放
  NpyIter_DECREF(iter);
  Py_DECREF(array1);
  Py_DECREF(array2);

  return 0;
}

2つのNumPy配列 arr1 と arr2 を用意します。
NpyIter_New() 関数を使用して、arr1 と arr2 をイテレーションする NpyIter 構造体を作成します。
NpyIter_GetDataPtrArray() 関数を使用して、各配列のデータポインタを取得します。
ループ内で、各要素の値を足し合わせて result 配列に格納します。
結果を printf() 関数で出力します。
最後に、NpyIter_DECREF() 関数と Py_DECREF() 関数を使用して、NpyIter 構造体とNumPy配列を解放します。

高速化のために、NpyIter_CopyLoop() 関数や NpyIter_FastCopyLoop() 関数を使用することもできます。
異なるデータ型を扱う場合は、適切な型変換を行う必要があります。
このコードは、NumPy バージョン 1.20 以上で使用できます。

NpyIter_GetFlat()

NpyIter_GetFlat() は、NpyIter 構造体内の全ての要素を 1 次元配列として取得する関数です。この関数は、char **NpyIter_GetDataPtrArray() と異なり、個々の配列への直接アクセスはできませんが、メモリ割り当てやデータ処理を簡略化できます。

例：NumPy イテレータ内の全ての要素を合計する

#include <numpy/npyiter.h>

int main() {
  // NumPy 配列を作成
  npy_intp dims[] = {2, 3};
  npy_intp strides[] = {sizeof(int), sizeof(int) * 3};
  int arr[2][3] = {{1, 2, 3}, {4, 5, 6}};
  PyArrayObject *array = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr, NPY_ORDER_C, NULL, NULL);

  // NumPy イテレータを作成
  NpyIter *iter = NpyIter_New(1, &array, NPY_ITER_MULTI_LOOP | NPY_ITER_READONLY);

  // 全ての要素を合計
  int sum = 0;
  char *ptr = NpyIter_GetFlat(iter);
  for (npy_intp i = 0; i < NpyIter_GetIterSize(iter); ++i) {
    sum += *(int *)ptr;
    ptr += NpyIter_GetItemSize(iter);
  }

  // 結果を出力
  printf("合計: %d\n", sum);

  // NumPy イテレータと配列を解放
  NpyIter_DECREF(iter);
  Py_DECREF(array);

  return 0;
}

ループによるデータアクセス

シンプルな処理の場合は、NpyIter 構造体の各属性に直接アクセスして、ループ内でデータ処理を行うこともできます。この方法は、メモリ割り当てや関数呼び出しを削減できますが、コードが冗長になり、可読性が低下する可能性があります。

例：NumPy イテレータ内の全ての要素を 2 倍する

#include <numpy/npyiter.h>

int main() {
  // NumPy 配列を作成
  npy_intp dims[] = {2, 3};
  npy_intp strides[] = {sizeof(int), sizeof(int) * 3};
  int arr[2][3] = {{1, 2, 3}, {4, 5, 6}};
  PyArrayObject *array = PyArray_NewFromDescr(&PyArray_DescrInt, 2, dims, strides, (void *)arr, NPY_ORDER_C, NULL, NULL);

  // NumPy イテレータを作成
  NpyIter *iter = NpyIter_New(1, &array, NPY_ITER_MULTI_LOOP);

  // 全ての要素を 2 倍
  for (npy_intp i = 0; i < NpyIter_GetIterSize(iter); ++i) {
    int *ptr = (int *)NpyIter_GetRawPtr(iter, 0);
    *ptr *= 2;

    // 次の要素へ移動
    NpyIter_Next(iter);
  }

  // NumPy イテレータと配列を解放
  NpyIter_DECREF(iter);
  Py_DECREF(array);

  return 0;
}

NumPy 以外にも、データ処理や配列操作に特化したライブラリが存在します。例えば、Cython や PyData/Sparse などは、NumPy よりも高速で効率的な処理を提供する場合があります。

状況に応じた選択

どの方法を選択するかは、処理内容、パフォーマンス要件、コードの可読性などを考慮する必要があります。

メモリ効率を重視する場合
NpyIter_GetDataPtrArray() が適しています。
高速化が必要な場合
NpyIter_GetFlat() またはループによるデータアクセスが適しています。