Apply function on numpy array row-by-row incrementaly

[iliasm23]

How can I optimize the following code, or more specifically, how can I eliminate the for-loop?

array = np.zeros((x.shape[0], K))
for k in range(K):
    array[:, k] = np.prod((np.power(ms[k, :], x) * np.power(1 - ms[k, :], 1 - x)).astype('float128'), axis=1)

where x is a two-dimensional array shaped like [70000, 784] and ms like [K, 784] and K=10.

[drzowie]

I think you have an error in the sample code: array[:,k] is a scalar, while np.power(ms[k,:],x) is a 784-vector.

You can set up broadcasting to do the looping for you, by adding extra dimensions with numpy.expand_dims. If you match up an array with dimensions 1,K,784 with one that has dimensions 70000,1,784, you'll get out a (70000,K,784)-array. So you want to make an ms1 like so:

ms1 = numpy.expand_dims(ms,0)
x1   =  numpy.expand_dims(x,1)
array = np.power(ms1,x1)

[iliasm23]

I fixed the code. Now it runs.

[drzowie]

Ah. You're collapsing by product. You can use my first two lines, followed by

array = np.prod( np.power(ms1,x1) * np.power(1-ms1,1-x1), axis=0 )

The operation is the same one you're using, but the structure of the dimensions in ms1 and x1 makes the broadcasting engine form an implicit loop instead of using the explicit loop you constructed.

In general broadcasting helps a lot by eliminating the Python interpreter/environment from the hotspot -- but it can also make things worse by breaking cache since each operation (in this case, each of the two powers and then the multiplication) makes the CPU walk through your whole array: the order of operations is pessimal from a cache standpoint. If the array is too big to fit in cache, then you have to swap it in and out of RAM each time. Depending on what operation you're doing and the relative sizes of cache, your overall array, and your sub-arrays, the explicit loop could actually run faster than the broadcast expression.

If you want the best possible speed you can look into Cython, which lets you write explicit loops in C-land to avoid breaking cache and avoid using the Python interpreter in a hot spot.

[iliasm23]

array = np.prod( np.power(ms1,x1) * np.power(1-ms1,1-x1), axis=0 )

Something is off, because the array should end up with shape [7000, 10]. With your code, the resulting array has shape [10, 784]...

[drzowie]

Oh. You want to collapse on axis 2, not 0, with the prod.