Numpy: Efficient Selective Multidimensional Operations (partial Ops)

Is there a way in numpy to do operations partially along some dimensions? For example, say I have 2 matrices of shapes a=(MxN), b=(NxK) and a 3rd one I want to store the dot-produc

Solution 1:

I'm a beginner with numpy, the indexing syntax is probably very well known by experienced users. For beginners like me though, I've just confirmed that the above syntax (see update in question) does exactly what I needed, and efficiently. Moreover, there's support for index-selections by both boolean-mask and indexes-lists as well.

Code snippet:

import numpy as np
import timeit

a = np.ones((1000, 2000))
b = np.ones((2000, 1000))
c = np.zeros((1000, 1000))

# Prepare dimensions selection (list)
dims=np.arange(0,500,2)
# Prepare dimensions selection (conditions)
cond=np.full((1000), False)
cond[dims] = Truedeffull_dot(): c = a.dot(b)
defdot_by_dims(): c[dims,:] = a[dims,:].dot(b)
defdot_by_cond(): c[cond,:] = a[cond,:].dot(b)
defdot_by_pyfor():
    for d in dims:
        c[d,:] = a[d,:].dot(b)

print('Full dot time: {} seconds'.format(timeit.Timer(full_dot).timeit(number=100)))
print('Dot partial python for time: {} seconds'.format(timeit.Timer(dot_by_pyfor).timeit(number=100)))
print('Dot partial by dims time: {} seconds'.format(timeit.Timer(dot_by_dims).timeit(number=100)))
print('Dot partial by cond time: {} seconds'.format(timeit.Timer(dot_by_cond).timeit(number=100)))

Output:

Full dot time: 4.831596400000002 seconds
Dot partial python for time: 12.7231105 seconds
Dot partialby dims time: 1.4666931999999946 seconds
Dot partialby cond time: 1.5701424000000017 seconds

(using a classic python for is about 8.5X slower than the numpy syntax! Even 2.6X slower than the full-dot, even though it does only partial dotting...)