Vectorizing code that includes conditional statement

Viewed 84

I've been using for loops for a script I've been working on, and I think the for loops are making the script run too slow. I think I could pick up a lot of speed with vectorization, but I'm having a little trouble figuring out how to vectorize my code. I'm copying some simple sample code that more or less mimics what I'm trying to accomplish in my actual code to make the question easier to follow. (Sorry if it's not the most elegant or technically sound. I'm still getting experience with Python.)

import numpy as np

def create_combo(input1,input2):

    combo = input1 + input2
    return combo

def another_combo(monitor,scalar_1,scalar_2):

    add_more = monitor + scalar_1 + scalar_2

    return add_more

# Initialize monitor
monitor = 0

# Initialize a threshold variable
threshold = 15

# Create input arrays
primary_1 = [1, 3, 5, 7, 9]
primary_2 = [2, 4, 6, 8, 10]

primary_1 = np.array(primary_1)
primary_2 = np.array(primary_2)

storage_vec = []
for i in range(5):
    # Create input variables
    scalar_1 = 0.5
    scalar_2 = 2

    # Call the create_combo function 
    combination = create_combo(primary_1[i],primary_2[i])

    # call the another_combo function
    add_on = another_combo(monitor,scalar_1,scalar_2)

    monitor = combination + add_on

    # Check if monitor exceeds threshold
    if monitor > threshold:
        # Store the index if so
        storage_vec.append(i)
        # Reset the variable, monitor
        monitor = 0

I can understand how it would be easy to vectorize the function named create_combo. And I can also see that it would be simple to make the variable named monitor a vector as well.

What I'm having trouble wrapping my head around is how I could reset that monitor variable at specific points in the vector and continue doing computations with monitor after it has been reset. So, I've been working with monitor as a scalar and doing computations on each element of the inputs to the functions. However, this seems to be way too slow, and I know vectorizing typically speeds things up.

I read that perhaps the np.where() function could help, but I'm also not sure how to apply that in this situation.

Any guidance?

1 Answers

Your example if probably over complicated (or your example is not complete).

But if we want to create the storage_vec easily we can noticed that monitor increase by 21.5 after each step until it reach the threshold.

So the vector storage_vec can simply be computed with:

#Increase value:
inc = primary_1[-1]+primary_2[-1]+scalar_1+scalar_2 # = 21.5 
#floor division
fd = threshold//inc
#ceil division
cd = -(-threshold//inc)
#storage_vec 
storage_vec = np.r_[fd:10:cd]
Related