Bence Béky: Using to the numpy C API

Problem


In [ ]:
from IPython.display import Image;
Image(filename="ipython.png")

In [ ]:
r = numpy.linspace(0.0, 1.0, 1000);
p = 0.1;
z = 0.5;

Python implementations


In [ ]:
def circleangleloop(r, p, z):
  # If the circle arc of radius r is disjoint from the circular disk 
  # of radius p, then the angle is zero.
  answer = numpy.zeros_like(r);
  for i in xrange(r.shape[0]):
    # If the planet entirely covers the circle, the half central angle is pi.
    if (r[i] < p-z):
      answer[i] = numpy.pi;
    # If the triangle inequalities hold between z, r, and p, 
    # then we have partial overlap.
    # If alpha is the half central angle in the triangle with sides r, p, and z,
    # with p opposite the angle, then p^2 = r^2 + z^2 - 2 rz cos(alpha)
    elif (r[i] < p+z) & (z < p+r[i]):
      answer[i] = numpy.arccos((r[i]*r[i]+z*z-p*p)/(2*z*r[i]));
  return answer;

pyplot.plot(r, circleangleloop(r, p, z), "r-");

In [ ]:
def circleanglemask(r, p, z):
  inside = (r < p-z);
  intersect = (r < p+z) & (z < r+p) & numpy.logical_not(inside);
  answer = numpy.zeros_like(r);
  answer[inside] = numpy.pi;
  answer[intersect] = numpy.arccos((numpy.power(r[intersect],2)+z*z-p*p)/(2*z*r[intersect]));
  return answer;

pyplot.plot(r, circleanglemask(r, p, z), "r-");

In [ ]:
def circleanglesorted(r, p, z):
  answer = numpy.empty_like(r);
  n = len(r);
  if (p > z):
    # Planet covers center of star.
    a, b = numpy.searchsorted(r, [p-z, p+z], side="right");
    answer[:a] = numpy.pi;
    answer[a:b] = numpy.arccos((r[a:b]*r[a:b]+z*z-p*p)/(2*z*r[a:b]));
    answer[b:] = 0.0;
  else:
    # Planet does not cover center of star.
    a, b = numpy.searchsorted(r, [z-p, z+p], side="right");
    answer[:a] = 0.0;
    answer[a:b] = numpy.arccos((r[a:b]*r[a:b]+z*z-p*p)/(2*z*r[a:b]));
    answer[b:] = 0.0;
  return answer;

pyplot.plot(r, circleanglesorted(r, p, z), "r-");

In [ ]:
from timeit import timeit;

n = 500;

r = ", numpy; r = numpy.linspace(0.0, 1.0, 1000)";
arg = "(r, 0.1, 0.5)";

time1 = timeit(stmt="toypython.circleangleloop" + arg, setup="import toypython" + r, number=n);
time2 = timeit(stmt="toypython.circleanglemask" + arg, setup="import toypython" + r, number=n);
time3 = timeit(stmt="toypython.circleanglesorted" + arg, setup="import toypython" + r, number=n);

print("Python loop:   {0:5.3f} ms.".format(1000*time1/n));
print("Python mask:   {0:5.3f} ms.".format(1000*time2/n));
print("Python sorted: {0:5.3f} ms.".format(1000*time3/n));

C implementations with the numpy API

toyc.c

#include <math.h>

void circleangleloop(double *r, double p, double z, int n, double *answer) {
  /* If the circle arc of radius r is disjoint from the circular disk 
     of radius p, then the angle is zero. */
  int i;
  double ri;
  for(i=0; i<n; i++) {
    ri = *(r+i);
    // If the planet entirely covers the circle, the half central angle is pi.
    if (ri <= p-z)
      *(answer+i) = M_PI;
    // If the triangle inequalities hold between z, r, and p, use law of cosines.
    else if ((ri < p+z) && (ri > z-p))
      *(answer+i) = acos((ri*ri+z*z-p*p)/(2*z*ri));
    else
      *(answer+i) = 0;
  }
  return;
}

toyc-wrapper.c

#include <Python.h>
#include <numpy/arrayobject.h>
#include "toyc.h"

/* Docstrings */
static char module_docstring[] = "  This module is a fast C implementation of a toy problem.";
static char circleangleloop_docstring[] = 
"  circleangleloop(r, p, z)\n"
"  Calculate half central angle of the arc of circle of radius r\n"
"  that is inside a circle of radius p with separation of centers z.";

/* Function wrappers for external use */
static PyObject *circleangleloop_wrapper(PyObject*, PyObject*, PyObject*);

/* Module specification */
static PyMethodDef module_methods[] = {
  {"circleangleloop", (PyCFunction)circleangleloop_wrapper, METH_VARARGS | METH_KEYWORDS, circleangleloop_docstring},
  {NULL, NULL, 0, NULL}
};

/* Initialize the module */
PyMODINIT_FUNC inittoyc(void) {
  PyObject *m = Py_InitModule3("toyc", module_methods, module_docstring);
  if (m == NULL)
    return;
  /* Load numpy functionality. */
  import_array();
}

/* Wrapper function for circleangleloop. */
static PyObject *circleangleloop_wrapper(PyObject *self, PyObject *args, PyObject *kwds) {
  /* Input arguments. */
  double p, z;
  PyObject *r_obj;

  // Keywords.
  static char *kwlist[] = {"r", "p", "z", NULL};

  /* Parse the input tuple */
  if (!PyArg_ParseTupleAndKeywords(args, kwds, "Odd", kwlist, &r_obj, &p, &z))
    return NULL;

  /* Interpret the input object as a numpy array. */
  PyObject *r_array = PyArray_FROM_OTF(r_obj, NPY_DOUBLE, NPY_IN_ARRAY);

  /* If that didn't work, or the resulting array does not have the correct
   * number of dimensions or type, then abort. */
  if (r_array == NULL || PyArray_NDIM(r_array) != 1 || PyArray_TYPE(r_array) != PyArray_DOUBLE) {
    PyErr_SetString(PyExc_ValueError, "r cannot be converted to a suitable array.");
    return NULL; 
  }

  /* Read out dimensions and data pointers. */
  int n = (int)PyArray_DIM(r_array, 0);
  double *r_data = (double*)PyArray_DATA(r_array);

  /* Create answer numpy array, let Python allocate memory.
     Do not allocate memory manually and then use PyArray_FromDimsAndData! */
  PyArrayObject *answer_array = (PyArrayObject*)PyArray_FromDims(1, &n, NPY_DOUBLE);

  // Evaluate the model
  circleangleloop(r_data, p, z, n, (double*)PyArray_DATA(answer_array));

  /* Clean up. */
  Py_DECREF(r_array);

  // Return.
  return PyArray_Return(answer_array);
}

toyc-setup.py

from distutils.core import setup, Extension;
import numpy.distutils.misc_util;

c_ext = Extension("toyc", ["toyc-wrapper.c", "toyc.c"], extra_compile_args=['-Ofast']);

setup(ext_modules=[c_ext], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs());

In [ ]:
from timeit import timeit;

n = 500;

r = ", numpy; r = numpy.linspace(0.0, 1.0, 1000)";
arg = "(r, 0.1, 0.5)";

time1 = timeit(stmt="toypython.circleangleloop" + arg, setup="import toypython" + r, number=n);
time2 = timeit(stmt="toypython.circleanglemask" + arg, setup="import toypython" + r, number=n);
time3 = timeit(stmt="toypython.circleanglesorted" + arg, setup="import toypython" + r, number=n);
time4 = timeit(stmt="toycython.circleangleloop" + arg, setup="import toycython" + r, number=n);
time5 = timeit(stmt="toycython.circleanglemask" + arg, setup="import toycython" + r, number=n);
time6 = timeit(stmt="toycython.circleanglesorted" + arg, setup="import toycython" + r, number=n);
time7 = timeit(stmt="toyc.circleangleloop" + arg, setup="import toyc" + r, number=n);
time8 = timeit(stmt="toyc.circleanglesorted" + arg, setup="import toyc" + r, number=n);

print("Python loop:   {0:5.3f} ms.".format(1000*time1/n));
print("Python mask:   {0:5.3f} ms.".format(1000*time2/n));
print("Python sorted: {0:5.3f} ms.".format(1000*time3/n));
print("Cython loop:   {0:5.3f} ms.".format(1000*time4/n));
print("Cython mask:   {0:5.3f} ms.".format(1000*time5/n));
print("Cython sorted: {0:5.3f} ms.".format(1000*time6/n));
print("C      loop:   {0:5.3f} ms.".format(1000*time7/n));
print("C      sorted: {0:5.3f} ms.".format(1000*time8/n));

In [ ]: