Python Optimization: extending with C

So, you want to code your project with your favourite language, which obviously :) is Python, and at the same time, you have some computationally heavy problem at your hand. But using Python for large calculations, for example, is not always preferred, because the interpreted nature of the language. Does this sound familiar?

Optimize, optimize...

One solution, of course, is to optimize your python code. This topic isn't covered here, but is a link to great tips for optimizing python code. Other solution, is to make your problem solving algorithm more effective. This is one major aspect in making your program work faster.

But what if you have optimized your code already? And you are using the best algorithms which have been known to be the fastest, and still the program is slow like a turtle? Then, there is one solution (besides upgrading your computer) left: extend your program with C! (or C++) Use Python to code the other parts with comfort, and then, write the routines which take up the CPU time with C!

Case study: calculating primes

Let's say that you have a program, which needs to calculate the first n primes. (Why? Dunno, that's up to you. :) So, you write your code with python, the outcome might be something like this:

def isprime1(input):
    if input < 1:
	return 0
    
    n = input-1
    
    while n > 1:
	if input%n == 0:
	    return 0
	n = n - 1

    return 1

def main():
    i = 0
    l = 0 

    max = int(sys.argv[1])

    while 1:
	
	if isprime1(i):
	    l = l +1

	i = i + 1
	if l == max:
	    break

    print max, "primes calculated"
    
if __name__ == "__main__":
    main()

The isprime1 function calculates if a given number is a prime number. I guess not the best 'algorithm' available, but that's not relevant here. Since Python is highly readable, no further comments for the code is required, right? :)

The program runs beautifully, but then, you notice that when the amount of calculated primes goes up, it starts to take enormously time. As we observe , the equivalent C program runs something like 30 times faster! Whoah, wouldn't like to wait 30 times longer with our Python version. Here's the C code:

int isprime1(int input)
{
	int n;

	if (input < 1) {
		return 0;
	}
	n = input - 1;

	while (n > 1){
		if (input%n == 0) return 0;
		n--;
	}
	return 1;
}

int main(int argc, char *argv[]){
	int i=0, l=0;
	int max;
	
	max = atoi(argv[1]);
	
	while (1) {
		if (isprime1(i)==1) {
			l++;
		}
		i++;
		if (l==max){
			break;
		}
	}
	
	printf("%i primes calculated\n",max);
}

Looks pretty much the same that the Python version.

Solution

So, what now? As I hinted before, we have to make the prime checking routine with C. First thought might be that it could help some, but not so much it would be significant. Guess again...

Let's take a look at some timings I made... Here the graph shows the time used to calculate first n primes.

As you can see, calculation with pure Python code starts to slow down much before the C version. And surprise surprise, the Python + C version of our program run nearly as fast as the pure C version! Very nice.

Here's a closer look to relation between the pure C version and the Python + C version of our program.

With large n, the difference between the C and the Python + C program becomes much smaller. Which is no other than good news for us! :)

The Extending

So, what is this extending business anyway and how do we do it? Basically, it's nothing more than writing a module to Python with C. The more into depth details and documentation of this process is described in Python documents, of course, but I'll give you a little example to whet your appitite. :)

Step 1

Decide what you want to code with C. In our case, the prime checking function is obviously the one we need to write with C, to get some speed up.

We have the code for that function ready, so now all we have to do is make it 'compatible' with Python. We'll name the new source file 'prmodule.c'. (I'll just show an example here, which is pretty much an imitation of that in Python documentation, except that we are calculating primes, so consult Python documentation for further information.)

The isprime function in prmodule.c:

static PyObject *pr_isprime(PyObject *self, PyObject *args){
	int n, input;

	if (!PyArg_ParseTuple(args, "i", &input))
	  return NULL;
	
	if (input < 1) {
		return Py_BuildValue("i", 0);
	}
	
	n = input - 1;

	while (n > 1){
		if (input%n == 0) return Py_BuildValue("i", 0);;
		n--;
	}

	return Py_BuildValue("i", 1);
}

The function must return a PyObject pointer and it takes PyObject pointers' as parameters too. That's how we get the real arguments from python call. PyArg_ParseTuple() is used to parse the arguments passed from python. We parse a integer from parameter, which tells us what number is to be checked. When returning a value, we must pass it in python object, Py_BuildValue() is used for that purpose.

Pretty simple and straightforward, when you get the idea. Though, prmodule.c isn't quite finished yet. We need to tell what methods are available in our module and last, make a initialization method for the module.

So we add:

static PyMethodDef PrMethods[] = {
	{"isprime", pr_isprime, METH_VARARGS, "Check if prime."},
	{NULL, NULL, 0, NULL}
};

void initpr(void){
	(void) Py_InitModule("pr", PrMethods);
}

to the prmodule.c source file.

Step 2

Next, we need to compile our module. There's the nice distutils package which does pretty much all we need. Perhaps you've already seen those setup.py named files from some python libraries. That's what we need to do now. (We're not really installing anything anywhere, just compiling.)

Let's write the setup.py:

from distutils.core import setup, Extension

module = Extension('pr', sources = ['prmodule.c'])
setup(name = 'Pr test', version = '1.0', ext_modules = [module])

and run 'python setup.py build'. This should compile the module. That's it! The result can be found under build/lib.linux-i686-2.1 (or similar) directory, named 'pr.so'. You can copy it to your source directory, then use it like any other module. For example:

comet:~/prog/python/prime$ python
Python 2.1.3 (#1, Apr 20 2002, 10:14:34) 
[GCC 2.95.4 20011002 (Debian prerelease)] on linux2
Type "copyright", "credits" or "license" for more information.
>>> import pr
>>> pr.isprime(17)
1
>>>