PythonのライブラリをC++で作成する（3）

いよいよ本題、C++でMatrixの掛け算をするPythonのライブラリを作ります。

まず、リストをstd::vectorに変換します。

PyObject *mul(PyObject *self, PyObject *args) {
    PyObject    *obj1, *obj2;
    long    D;
    if(!PyArg_ParseTuple(args, "OOl", &obj1, &obj2, &D))
        ...

まず、このように引数のタプルからPyObjectを取り出します。このobj1、obj2がリストになっています。そして、このリストから要素を取り出します。以下は、リストの要素が整数の場合です。

    const int   M = PyList_Size(obj_list);
    for(int j = 0; j < M; ++j) {
        PyObject    *obj = PyList_GetItem(obj_list, (Py_ssize_t)j);
        const long  d = PyLong_AsLong(obj);
        v.push_back(d);     // std::vector<long>
    }

タプルの取り出し方とだいたい同じですね。
Matrixはこれを2層で行えばいいだけです。簡単ですね。

次に、std::vectorをPythonのリストに変換します。

    const size_t    L = v.size();   // v : std::vector<long>
    PyObject    *obj_list = PyList_New(L);
    for(size_t i = 0U; i < L; ++i)
        PyList_SET_ITEM(obj_list, i, PyLong_FromLong(v[i]));
    return obj_list;

長さLのリストオブジェクトを作り、longをPyObjectに変換して、リストにセットします。

リストオブジェクト - Python 3.9.1 ドキュメント

コードは、

# setup.py
from distutils.core import setup, Extension

setup(name = 'mymatrix', version = '1.0.0', \
    ext_modules = [Extension('mymatrix', ['matrix.cpp'])])

// matrix.h
#include <Python.h>
#include <vector>

typedef std::vector<long>   Vec;
typedef std::vector<Vec>    Matrix;

static Vec make_vector(PyObject *obj_list);
static Matrix make_matrix(PyObject *args);
static PyObject *mul(PyObject *self, PyObject *args);
static Matrix mul_core(const Matrix& A, const Matrix& B, long D);
static Matrix transpose(const Matrix& A);
static PyObject *make_python_matrix(const Matrix& A);
static PyObject *make_python_vector(const Vec& v);

// matrix.cpp
#include <iostream>
#include "matrix.h"

using namespace std;

Vec make_vector(PyObject *obj_list) {
    if(!PyList_Check(obj_list))
        return Vec(0);
    
    Vec v;
    const int   M = PyList_Size(obj_list);
    for(int j = 0; j < M; ++j) {
        PyObject    *obj = PyList_GetItem(obj_list, (Py_ssize_t)j);
        if(!PyLong_Check(obj)) {
            printf("vector[%d] is not long.\n", j);
            return Vec(0);
        }
        
        const long  d = PyLong_AsLong(obj);
        v.push_back(d);
    }
    
    return v;
}

Matrix make_matrix(PyObject *obj_table) {
    if(!PyList_Check(obj_table))
        return Matrix(0);
    
    const int   L = PyList_Size(obj_table);
    Matrix  X(L);
    for(int i = 0; i < L; ++i) {
        PyObject    *obj = PyList_GetItem(obj_table, (Py_ssize_t)i);
        Vec v = make_vector(obj);
        if(v.empty())
            return Matrix(0);
        X[i] = v;
    }
    
    return X;
}

PyObject *mul(PyObject *self, PyObject *args) {
    PyObject    *obj1, *obj2;
    long    D;
    if(!PyArg_ParseTuple(args, "OOl", &obj1, &obj2, &D))
        return make_python_matrix(Matrix(0));
    
    const Matrix    A = make_matrix(obj1);
    const Matrix    B = make_matrix(obj2);
    Matrix  C = mul_core(A, B, D);
    return make_python_matrix(C);
}

Matrix mul_core(const Matrix& A, const Matrix& B, long D) {
    const size_t    H = A.size();
    const size_t    M = B.size();
    const size_t    W = B.front().size();
    
    Matrix  C(H, Vec(W));
    for(size_t i = 0U; i < H; ++i) {
        for(size_t j = 0U; j < W; ++j) {
            for(size_t k = 0U; k < M; ++k)
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % D;
        }
    }
    return C;
}

PyObject *make_python_matrix(const Matrix& A) {
    const size_t    H = A.size();
    PyObject    *obj_mat = PyList_New(H);
    for(size_t i = 0U; i < H; ++i) {
        PyObject    *obj_list = make_python_vector(A[i]);
        PyList_SET_ITEM(obj_mat, i, obj_list);
    }
    return obj_mat;
}

PyObject *make_python_vector(const Vec& v) {
    const size_t    L = v.size();
    PyObject    *obj_list = PyList_New(L);
    for(size_t i = 0U; i < L; ++i)
        PyList_SET_ITEM(obj_list, i, PyLong_FromLong(v[i]));
    return obj_list;
}

static PyObject *print_matrix(PyObject *self, PyObject *args) {
    PyObject    *obj_table;
    if(!PyArg_ParseTuple(args, "O", &obj_table))
        return make_python_matrix(Matrix(0));
    
    Matrix  X = make_matrix(obj_table);
    for(auto p = X.begin(); p != X.end(); ++p) {
        for(auto q = p->begin(); q != p->end(); ++q)
            cout << *q << ' ';
        cout << endl;
    }
    
    return PyLong_FromLong(0);
}

static PyMethodDef mymatrixMethods[] = {
    { "print_matrix", print_matrix, METH_VARARGS },
    { "mul", mul, METH_VARARGS },
    { NULL }
};

static struct PyModuleDef mymatrix = {
    PyModuleDef_HEAD_INIT, "mymatrix", "Python3 C API Module",
    -1, mymatrixMethods
};

PyMODINIT_FUNC PyInit_mymatrix() {
    return PyModule_Create(&mymatrix);
}

Python側は、

import mymatrix

def mat_pow(M, e, D):
    if e == 1:
        return M
    elif e % 2 == 1:
        return mymatrix.mul(M, mat_pow(M, e-1, D), D)
    else:
        A = mat_pow(M, e/2, D)
        return mymatrix.mul(A, A, D)

ここが最初のコードと違うだけです。

実行してみましょう。

1134
256330968
206.85379028320312

そんなに変わりませんね。なぜでしょう。工夫する余地が2つあります。