Как уменьшить время выполнения этого кода, содержащего 2 цикла for?

Question

В следующем фрагменте кода цикл 2 for занимает в среднем около 0,05 с. data является массивом NumPy. Каждая ячейка i, j содержит кортеж, в котором хранятся значения RGB, определенные этими функциями. Позже я создаю изображение, используя эти значения RGB, и i, j являются x, y координатами пикселей.

Можно ли как-то сократить время этой операции? Или есть какой-нибудь другой способ создать изображение, указав значения RGB каждого пикселя в виде математических функций?

start = time.time()
for i in range (0, 150):
    for j in range(0, 150):
        data[i,j] = [int(math.sin(math.pi*i/300.0)*127.5 + 127.5),
                     int(math.cos(peak)*127.5 + 127.5),
                     int(math.sin(math.pi*j/300.0)*127.5 + 127.5)] 
print ('Time: ', time.time() - start)```

Answer 1

Вы можете удалить все циклы, используя ogrid, и вместо этого просто использовать три вычисления, одно для r, g и b.

---

x, y, z = data.shape
i, j = np.ogrid[:x, :y]

data[..., 0] = (np.sin(np.pi*i/300)*127.5 + 127.5).astype(int)
data[..., 1] = (np.cos(peak)*127.5 + 127.5).astype(int)
data[..., 2] = (np.sin(np.pi*j/300)*127.5 + 127.5).astype(int)

---

Производительность

def rgb_vectorized(x, y, peak=1):
    data = np.empty((x, y, 3), dtype=int)
    i, j = np.ogrid[:x, :y]
    data[..., 0] = (np.sin(np.pi*i/300)*127.5 + 127.5).astype(int)
    data[..., 1] = (np.cos(peak)*127.5 + 127.5).astype(int)
    data[..., 2] = (np.sin(np.pi*j/300)*127.5 + 127.5).astype(int)
    return data

def rgb_original(x, y, peak=1):
    data = np.empty((x, y, 3), dtype=int)
    for i in range (x):
        for j in range(y):
            data[i,j] = [int(math.sin(math.pi*i/300.0)*127.5 + 127.5),
                         int(math.cos(peak)*127.5 + 127.5),
                         int(math.sin(math.pi*j/300.0)*127.5 + 127.5)]
    return data

---

%timeit rgb_vectorized(1000, 1000)
9.85 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit rgb_original(1000, 1000)
3.4 s ± 27.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Проверка

>>> np.array_equal(rgb_vectorized(1000, 1000), rgb_original(1000, 1000))
True