Parallelization in Python using Weave and OpenMP

If you're using python a lot, you probably stumbled across the performance python howto at some point. Since inner loops are so
expensive in python, it makes sense to move them to C++, and
PerformancePython shows how to do that. When standard libraries like
numpy and weave.blitz are not expressive enough, your best bet becomes
SciPy Weave. In weave you write your C++ code in a string and access
your numpy arrays via the blitz converters. The code is compiled when
it has changed and can make use of arbitrary libraries as we shall see.

Parallelizing Weave-Code: Per-Pixel eigenvalues

Lets suppose we have a 3D image $D$, with the second order derivatives
in $x$, $y$ and $z$ direction given by $Dxx, Dxy, Dyy, \ldots$. We
want to calculate the eigenvalues at each pixel. For efficiency, we
use LAPACK and for clarity, we use Boost uBLAS. Lets first write down
the C++ code:

namespace ublas = boost::numeric::ublas;
namespace lapack = boost::numeric::bindings::lapack;
using namespace std;
typedef ublas::bounded_matrix<double,3,3,ublas::column_major> mat;
ublas::bounded_vector<double,3> lambda;
ublas::bounded_vector<double,34*3> work; // the 34*width is from syev.hpp

mat A;
int i,j,k,idx;
#pragma omp parallel for private(i,j,k,A,idx,lambda,work)
for(i=0;i<nx;i++){
  for(j=0;j<ny;j++){
    for(k=0;k<nz;k++){
      A(0,0) = Dxx(i,j,k);  // fill upper triangle of A
      A(1,1) = Dyy(i,j,k);
      A(2,2) = Dzz(i,j,k);
      A(0,1) = Dxy(i,j,k);
      A(0,2) = Dxz(i,j,k);
      A(1,2) = Dyz(i,j,k);

      // determine eigenvalues (N) of upper (U) triangle of A in lambda
      lapack::syev('N','U',A,lambda,lapack::workspace(work));
      lambda1(i,j,k) = lambda(0);  // eigenvalues in ascending order
      lambda2(i,j,k) = lambda(1);
      lambda3(i,j,k) = lambda(2);
    }
  }
 }  

Note that we use OpenMP here to parallelize the outmost for-loop using
a #pragma directive. I took me a while to figure out that one needs
to explicitly name all variables which must be private to a thread in
this for-loop, otherwise you will not notice any speedups and/or weird
program behavior. I suppose further restrictions are that the size of
the variables needs to be known at compile time (here I use
bounded_vector for that purpose). Finally, the loop index must be of
type int for OpenMP.

That is only the body of a function. We save this in a python string
called codestring and execute it as follows:

from scipy.weave import inline, converters
variables = "nx ny nz lambda1 lambda2 lambda3 Dxx Dyy Dzz Dxy Dxz Dyz".split()
inline(codestring,
       variables,
       extra_compile_args =['-O3 -fopenmp'],
       extra_link_args=['-lgomp'],
       headers=[
           '<boost/numeric/bindings/lapack/syev.hpp>',
           '<boost/numeric/bindings/traits/ublas_matrix.hpp> ',
           '<boost/numeric/bindings/traits/ublas_vector.hpp> ',
           '<boost/numeric/ublas/matrix.hpp>',
           '<boost/numeric/ublas/banded.hpp>',
           '<boost/numeric/ublas/vector.hpp>',
           '<cmath>'],
       type_converters=converters.blitz,
       libraries=["lapack"])

A couple of things to point out here: since codestring is only the
body of a function, we need some other place to declare include
files and declarations. Include files are listed in headers,
additional code with declarations can be supplied using
support_code. To use OpenMP, we need the compiler argument
-fopenmp, and the linker needs to link the gomp library. The
type_converters argument tells Weave to copy and use conversion
routines into the final program which convert NumPy arrays in Blitz
arrays, which have nice accessor operators.

Second Order Parallelization

We can go even one step further and parallelize the resulting code
over multiple processors or even computers. In Python, a comparably
new and exciting way to achieve this is Parallel IPython. I'll leave
the details on this to a later posting, instead I'll point to a
problem I found while running Weave code in parallel on multiple machines.


When Weave compiles your code, it writes the generated source code to
~/.python26_compiled. The file name is created from a hash code over
the contained source code. This way, Weave knows when to recompile and
when to reuse an old binary.


Now to the downside:

Multiple Cores

You will get problems if you run multiple
instances at the same time: They try to write the same source code
file at the same time.

Multiple Computers

You will get the same problem if your home
directory is on the same file system (NFS or the like)

My solution to problem number one is to compile once with a single
thread, then reuse using all threads. Problem two can be solved by
symlinking ~/.python26_compiled to /tmp.

Directory Color Embedding

Do you also save all your experiment data in folder names which
represent the settings? I do. If you have many experiments, however,
this strategy will defy comparison of the experiments, as it becomes
hard to put everything in one plot. The default settings in plotting
programs such as gnuplot or matplotlib do not have a large enough
range of colors, and if they do, it becomes hard to assign the colors
in a meaningful way for exploratory data analysis. The script here
might come to rescue.

Reading pre-processing directory names

We first transform directory names such that test-paramA_0.1 becomes
test-paramA_0000.1, allowing easier string comparison.

import Levenshtein as L
import numpy as np
from colormath.color_objects import LabColor
import os, re

def fillfunc(mo):
    """ returns the matched number, with padded zeros to the front """
    s = mo.group(1)
    while(len(s)<6): s = '0'+s
    return s
def get_unified_names(fns):
    n = len(fns)
    strs = [x for x in fns]
    for i in xrange(len(strs)):
        strs[i] = re.sub(r'(?<!\d)([\d.]+)(?!\d)',fillfunc,strs[i])
    return strs

Creating the distance matrix based on Levenstein string distance

The Levenshtein string distance is a measure of how many
edits/insertions/deletions one needs to transform one string into the
other. With this we automatically derive a dissimilarity measure for
the unified directory strings:

mat = np.zeros((n,n))
for i in xrange(n):
    for j in xrange(n):
        if j<=i: continue
        mat[i,j] = L.distance(strs[i],strs[j])
        mat[j,i] = mat[i,j]
return np.exp(-mat)

Embedding the dissimilarity matrix using Multidimensional Scaling (MDS) in R

The distance matrix does not necessarily directly map to coordinates
in 3D, so we need an embedding algorithm which deals with distance
matrices only.


For simplicity, we call R from python with a saved dissimilarity matrix.

def emb(mat):
    np.savetxt("file.dat", mat)
    os.system("./analyse/emb.R")
    res = np.loadtxt("mds.dat")
    return res

Now to the part written in R, which only does the embedding and saves
the result in a text file:

#!/usr/bin/Rscript
library(MASS)
library(vegan)
tab <- read.table("file.dat", header = FALSE, sep=" ")
data.m    <- as.matrix(tab)
data.mds <- vegan::metaMDS(data.m, k=3, trymax=50)
write.table(data.mds$points, file="mds.dat", quote=F, row.names=F, col.names=F)

Transforming embedded coordinates into colors

The coordinates are in some – it seems arbitrary – range and
therefore need to be mapped to colors. We want similar colors to have
similar edit distance, we therefore map our coordinates directly to
Lab Color Space and then transform the Lab color coordinates to RGB:

def getcolor(embedded,i):
    #lab = LabColor(0.8,*embedded[i,:])
    lab = LabColor(*embedded[i,:])
    rgb = lab.convert_to('RGB', debug=False).get_rgb_hex()
    return rgb

The resulting string can be used directly in the color specification
of a matplotlib plot and is then best combined with picking to find
out more about a particularly interesting plotline.

Example Embedding of directory names