1+1=10

扬长避短 vs 取长补短

"My first C Extension To Numpy"

A Simple Example

Let's create an extension module called opee and let's say we want to create a Python interface to one C function example. This function takes a numpy array as argument and return an integer. We want this function to be callable from Python as follows:

>>> import opee
>>> import numpy
>>> a = numpy.zeros( (2,3,4) )
>>> opee.example()

To do this we need two files. The first is the C file which contains the actual code, and the second is the setup.py file used to create the module.

Souce Code

Begin by creating a file opeemodule.c.

#include "Python.h"
#include "numpy/arrayobject.h"

static PyObject*
example (PyObject *dummy, PyObject *args)
{
    PyObject *arg1=NULL;
    PyObject *arr1=NULL;
    int nd;

    if (!PyArg_ParseTuple(args, "O", &arg1))
        return NULL;

    arr1 = PyArray_FROM_OTF(arg1, NPY_DOUBLE, NPY_IN_ARRAY);
    if (arr1 == NULL)
        return NULL;

    nd = PyArray_NDIM(arr1);   //number of dimensions

    Py_DECREF(arr1);

    return PyInt_FromLong(nd);
}

static struct PyMethodDef methods[] = {
    {"example", example, METH_VARARGS, "descript of example"},
    {NULL, NULL, 0, NULL}
};

PyMODINIT_FUNC
initopee (void)
{
    (void)Py_InitModule("opee", methods);
    import_array();
}

Building

First, we create a setup script.

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from distutils.core import setup, Extension
import numpy as np

ext_modules = [ Extension('opee', sources = ['opeemodule.c']) ]

setup(
        name = 'Opee',
        version = '1.0',
        include_dirs = [np.get_include()], #Add Include path of numpy
        ext_modules = ext_modules
      )

Then run

python setup.py instll

Second Example

test.py

import opee
import numpy as np

a = np.zeros((8, 7, 3, 2))
b = np.array([[1.1, 2.2, 3.3], [1.2, 1.3, 1.4]])
c = np.array([[1.1, 2.2, 3.3], [1.2, 1.3, 1.4]])
opee.example(a, b, c)
print c

opeemodule.cpp

#include <iostream>

#include "Python.h"
#include "numpy/arrayobject.h"

static PyObject*
example (PyObject *dummy, PyObject *args)
{
    PyObject *arg1=NULL, *arg2=NULL, *out=NULL;
    PyArrayObject *arr1=NULL, *arr2=NULL, *oarr=NULL;

    if (!PyArg_ParseTuple(args, "OOO!", &arg1, &arg2,
        &PyArray_Type, &out)) return NULL;

    arr1 = (PyArrayObject*)PyArray_FROM_OTF(arg1, NPY_DOUBLE, NPY_IN_ARRAY);
    if (arr1 == NULL) return NULL;
    arr2 = (PyArrayObject*)PyArray_FROM_OTF(arg2, NPY_DOUBLE, NPY_IN_ARRAY);
    if (arr2 == NULL) goto fail;
    oarr = (PyArrayObject*)PyArray_FROM_OTF(out, NPY_DOUBLE, NPY_INOUT_ARRAY);
    if (oarr == NULL) goto fail;

    /*vv* code that makes use of arguments *vv*/

    int nd = PyArray_NDIM(arr1);   //number of dimensions
    npy_intp *shape = PyArray_DIMS(arr1);  // npy_intp array of length nd showing length in each dim.
    for (int i=0; i<nd; ++i)
        std::cout<<" "<<shape[i];
    std::cout<<std::endl;

    for (int i=0; i<arr2->nd; ++i)
        std::cout<<" "<<arr2->dimensions[i];
    std::cout<<std::endl;

    for (int i=0; i<oarr->dimensions[0]; ++i) {
        for (int j=0; j<oarr->dimensions[1]; ++j) {
            double *v = (double*)PyArray_GETPTR2(oarr, i, j);
            *v = *v * 2;
        }
    }
    /*^^* code that makes use of arguments *^^*/

    Py_DECREF(arr1);
    Py_DECREF(arr2);
    Py_DECREF(oarr);
    Py_INCREF(Py_None);
    return Py_None;

 fail:
    Py_XDECREF(arr1);
    Py_XDECREF(arr2);
    PyArray_XDECREF_ERR(oarr);
    return NULL;
}

//....

Run

python setup.py build_ext --inplace

Reference

Comments