CODEWITHSUNDEEP
pythonnumpyfor-loopoptimization
Viktor VN
Viktor VN

Reputation: 235

Removing for loops and appending to lists in numpy

I need to optimize my code and have been able to remove for loops almost everywhere but struggle for this small part. I've looked at numpy.where but don't think I can use that but I'm not quite sure. If anyone knows which functions I should look at to optimize this part that would be greatly appreciated. This is used quite often and for loops are very slow in pyton so i need to do some optimization.

def f(x):
    return np.pi * x * np.cos((np.pi / 2) * x ** 2)

x_samples = np.random.uniform(low=0.0, high=1.0, size=500000)
y_samples = np.random.uniform(low=0.0, high=1.0, size=500000) * 1.61
y_functie = f(x_samples)
y_hit, x_hit, x_miss, y_miss = [], [], [], []
for n in range(len(x_samples)):
    if y_samples[n] <= y_functie[n]:
        x_hit.append(x_samples[n]), y_hit.append(y_samples[n])
    else:
        x_miss.append(x_samples[n]), y_miss.append(y_samples[n])

Upvotes: 0

Views: 61

Answers (2)

Dschoni
Dschoni

Reputation: 3862

By doing indexing you could do something like this:

import numpy as np

def f(x):
return np.pi * x * np.cos((np.pi / 2) * x ** 2)

x_samples = np.random.uniform(low=0.0, high=1.0, size=500000)
y_samples = np.random.uniform(low=0.0, high=1.0, size=500000) * 1.61

y_functie = f(x_samples)

def test1():
    y_hit, x_hit, x_miss, y_miss = [], [], [], []
    for n in range(len(x_samples)):
        if y_samples[n] <= y_functie[n]:
            x_hit.append(x_samples[n]), y_hit.append(y_samples[n])
        else:
            x_miss.append(x_samples[n]), y_miss.append(y_samples[n])

def test2():
    x_hit = x_samples[y_samples<=y_functie]
    x_miss = x_samples[y_samples>y_functie]
    y_hit = y_samples[y_samples<=y_functie]
    y_miss = y_samples[y_samples>y_functie]
    
test1()
test2()

With your test data I get timing differences of 700ms for test1 and 21ms for test2. Of course you could simplify more by using only one boolean table and do everything in one go.

Depending on what you want to optimize for (speed or memory footprint) you could precalculate the boolean table with

hit_test = y_samples<=y_functie

Upvotes: 1

Michael
Michael

Reputation: 2367

You can use logical masking with the fancy indexing to achieve the desired outcome:

import numpy as np

def f(x):
    return np.pi * x * np.cos((np.pi / 2) * x ** 2)

x_samples = np.random.uniform(low=0.0, high=1.0, size=500000)
y_samples = np.random.uniform(low=0.0, high=1.0, size=500000) * 1.61

y_functie = f(x_samples)

hits = y_samples <= y_functie
misses = y_samples > y_functie

y_hit, x_hit, x_miss, y_miss = x_samples[hits], y_samples[hits], x_samples[misses], y_samples[misses]

Cheers.

Upvotes: 0

Related Questions