Calculate function on array slices in a vectorized way

user_10639468 0
Viewed 127

Let's say, I have 1D numpy arrays X (features) and Y (binary classes) and a function f that takes two slices of X and Y and calculates a number.

I also have an array of indices S, by which I need to split X and Y. It is guaranteed, that each slice will be not empty.

So my code looks like this:

def f(x_left, y_left, x_right, y_right):
    n = x_left.shape[0] + x_right.shape[0]

    lcond = y_left == 1
    rcond = y_right == 1

    hleft = 1 - ((y_left[lcond].shape[0])**2
                     + (y_left[~lcond].shape[0])**2) / n**2

    hright = 1 - ((y_right[rcond].shape[0])**2
                     + (y_right[~rcond].shape[0])**2) / n**2

    return -(x_left.shape[0] / n) * hleft - (x_right.shape[0] / n) * hright

results = np.empty(len(S))
for i in range(len(S)):
    results[i] = f(X[:S[i]], Y[:S[i]], X[S[i]:], Y[S[i]:])

The array results must contain results of f on each split from S.

len(results) == len(S)

My question is how to perform my calculations in vectorised way, using numpy, to make this code faster?

vectorization python-3.x numpy arrays python
1 Answers

First, let's make your function a bit more efficient. You are doing some unnecessary indexing operations: instead of y_left[lcond].shape[0] you just need lcond.sum(), or len(lcond.nonzero()[0]) which seems to be faster.

Here's an improved loopy version of your code (complete with dummy input):

import numpy as np           

n = 1000                     
X = np.random.randint(0,n,n) 
Y = np.random.randint(0,n,n) 
S = np.random.choice(n//2, n)

def f2(x, y, s):                                     
    """Same loopy solution as original, only faster"""
    n = x.size                                       
    isone = y == 1                                   
    lval = len(isone[:s].nonzero()[0])               
    rval = len(isone[s:].nonzero()[0])               

    hleft = 1 - (lval**2 + (s - lval)**2) / n**2     
    hright = 1 - (rval**2 + (n - s - rval)**2) / n**2

    return - s / n * hleft - (n - s) / n * hright

def time_newloop():                                   
    """Callable front-end for timing comparisons"""   
    results = np.empty(len(S))                        
    for i in range(len(S)):                           
        results[i] = f2(X, Y, S[i])                   
    return results

The changes are fairly straightforward.

Now, it turns out that we can indeed vectorize your loops. For this we have to compare using each element of S at the same time. The way we can do this is creating a 2d mask of shape (nS, n) (where S.size == nS) which cuts off the values up to the corresponding element of S. Here's how:

def f3(X, Y, S):                                     
    """Vectorized solution working on all the data at the same time"""
    n = X.size                                                        
    leftmask = np.arange(n) < S[:,None] # boolean, shape (nS, n)      
    rightmask = ~leftmask # boolean, shape (nS, n)              

    isone = Y == 1 # shape (n,)                                 
    lval = (isone & leftmask).sum(axis=1) # shape (nS,)         
    rval = (isone & rightmask).sum(axis=1) # shape (nS,)        

    hleft = 1 - (lval**2 + (S - lval)**2) / n**2                
    hright = 1 - (rval**2 + (n - S - rval)**2) / n**2           

    return - S / n * hleft - (n - S) / n * hright # shape (nS,) 

def time_vector():                                             
    """Trivial front-end for fair timing"""                    
    return f3(X,Y,S)

Defining your original solution to be run as time_orig() we can check that the results are the same:

>>> np.array_equal(time_orig(), time_newloop()), np.array_equal(time_orig(), time_vector())
(True, True)

And the runtimes with the above random inputs:

>>> %timeit time_orig()
... %timeit time_newloop()
... %timeit time_vector()
... 
... 
19 ms ± 501 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
11.4 ms ± 214 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.93 ms ± 37.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

This means that the loopy version above is almost twice as fast as the original loopy version, and the vectorized version is another factor of three faster. Of course the cost of the latter improvement is an increased memory need: instead of arrays of shape (n,) you now have arrays of shape (nS, n) which can get quite big if your input arrays are huge. But as they say there's no free lunch, with vectorization you often trade runtime for memory.

user_50673110
Related