Large point-matrix array multiplication in numpy

Question

Given two large numpy arrays, one for a list of 3D points, and another for a list of transformation matrices. Assuming there is a 1 to 1 correspondence between the two lists, i'm looking for the best way to calculate the result array of each point transformed by it's corresponding matrix.

My solution to do this was to use slicing (see "test4" in the example code below) which worked fine with small arrays, but fails with large arrays because of how memory-wasteful my method is :)

import numpy as np
COUNT    = 100
matrix   = np.random.random_sample((3,3,))   # A single matrix
matrices = np.random.random_sample((COUNT,3,3,)) # Many matrices
point    = np.random.random_sample((3,))     # A single point
points   = np.random.random_sample((COUNT,3,)) # Many points

# Test 1, result of a single point multiplied by a single matrix
# This is as easy as it gets
test1 = np.dot(point,matrix)
print 'done'

# Test 2, result of a single point multiplied by many matrices
# This works well and returns a transformed point for each matrix
test2 = np.dot(point,matrices)
print 'done'

# Test 3, result of many points multiplied by a single matrix
# This works also just fine
test3 = np.dot(points,matrix)
print 'done'

# Test 4, this is the case i'm trying to solve. Assuming there's a 1-1
# correspondence between the point and matrix arrays, the result i want
# is an array of points, where each point has been transformed by it's
# corresponding matrix
test4 = np.zeros((COUNT,3))
for i in xrange(COUNT):
    test4[i] = np.dot(points[i],matrices[i])
print 'done' 

With a small array, this works fine. With large arrays, (COUNT=1000000) Test #4 works but gets rather slow.

Is there a way to make Test #4 faster? Presuming without using a loop?

Answer 1

You can use numpy.einsum. Here's an example with 5 matrices and 5 points:

In [49]: matrices.shape
Out[49]: (5, 3, 3)

In [50]: points.shape
Out[50]: (5, 3)

In [51]: p = np.einsum('ijk,ik->ij', matrices, points)

In [52]: p[0]
Out[52]: array([ 1.16532051,  0.95155227,  1.5130032 ])

In [53]: matrices[0].dot(points[0])
Out[53]: array([ 1.16532051,  0.95155227,  1.5130032 ])

In [54]: p[1]
Out[54]: array([ 0.79929572,  0.32048587,  0.81462493])

In [55]: matrices[1].dot(points[1])
Out[55]: array([ 0.79929572,  0.32048587,  0.81462493])

The above is doing matrix[i] * points[i] (i.e. multiplying on the right), but I just reread the question and noticed that your code uses points[i] * matrix[i]. You can do that by switching the indices and arguments of einsum:

In [76]: lp = np.einsum('ij,ijk->ik', points, matrices)

In [77]: lp[0]
Out[77]: array([ 1.39510822,  1.12011057,  1.05704609])

In [78]: points[0].dot(matrices[0])
Out[78]: array([ 1.39510822,  1.12011057,  1.05704609])

In [79]: lp[1]
Out[79]: array([ 0.49750324,  0.70664634,  0.7142573 ])

In [80]: points[1].dot(matrices[1])
Out[80]: array([ 0.49750324,  0.70664634,  0.7142573 ])

Answer 2

It doesn't make much sense to have multiple transform matrices. You can combine transform matrices as in this question:

If I want to apply matrix A, then B, then C, I will multiply the matrices in reverse order np.dot(C,np.dot(B,A))

So you can save some memory space by precomputing that matrix. Then applying a bunch of vectors to one transform matrix should be easily handled (within reason).

I don't know why you would need one million transformation on one million vectors, but I would suggest buying a larger RAM.

Edit: There isn't a way to reduce the operations, no. Unless your transformation matrices hold a specific property such as sparsity, diagonality, etc. you're going to have to run all multiplications and summations. However, the way you process these operations can be optimized across cores and/or using vector operations on GPUs.

Also, python is notably slow. You can try splitting numpy across your cores using NumExpr. Or maybe use a BLAS framework on C++ (notably quick ;))