NumPy - Векторизация циклов с использованием итераторов диапазона

Question

Есть ли способ заставить это работать без циклов for?

import import numpy as np
import matplotlib.pyplot as plt    

L = 1
N = 255
dh = 2*L/N
dh2 = dh*dh

phi_0 = 1
c = int(N/2)
r_0 = L/2

arr = np.empty((N, N))

for i in range(N):

    for j in range(N):

            arr[i, j] = phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0


plt.imshow(arr)

Я пробовал вызывать функцию (x [None ,:], y [:, None]), где:

function(i, j):

    return phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0

но для этого нужны методы list .any или .all.Я ищу специально бездействующий метод (без функции и векторизации).Большое спасибо!

Answer 1

Векторизованное решение с использованием открытых сеток

Мы могли бы использовать два открытых массива диапазона / сетки для N, имитирующих то же поведение, что и итераторы -

I = np.arange(N)
mask = (I[:,None] - c)**2 + (I - c)**2 < r_0**2/dh2
out = np.where(mask,phi_0,0)

Для общего диапазона в двух циклах

Для общего случая, когда мы будем повторять два цикла, которые продолжаются до, скажем, M и N соответственно, мы можем использовать np.ogrid для создания этих открытых сеток и последующего использования в тех же строках -

I,J = np.ogrid[:M,:N]
mask = (I - c)**2 + (J - c)**2 < r_0**2/dh2

Для общего числа циклов

Дляобщее число циклов, просто создайте столько переменных, сколько число циклов.Следовательно, для трех циклов:

for i in range(M):
    for j in range(N):
        for k in range(P):

, мы бы имели:

I,J,K = np.ogrid[:M,:N,:P]

, затем использовали бы I,J,K вместо i,j,k соответственно для поэлементных операций, как у нас здесь.

---

Альтернатива для замены последнего шага для этого конкретного случая

Последний шаг также может быть реализован с поэлементным умножением путем масштабирования до phi_0 с maskв качестве части else устанавливается значение 0s -

out = mask*phi_0

Answer 2

Если ваша настоящая цель не в том, чтобы избежать циклов, а в том, чтобы получить хорошую производительность (в данном случае 670-кратное ускорение), простой подход заключается в использовании компилятора.В этом примере я использую Numba, но вы также можете использовать Cython, что немного сложнее (объявление типа, ...)

Пример

import numpy as np
import numba as nb
import matplotlib.pyplot as plt    

L = 1
N = 255
dh = 2*L/N
dh2 = dh*dh

phi_0 = 1
c = int(N/2)
r_0 = L/2

@nb.njit()
def create_arr(N,phi_0,c,r_0,dh2):
  arr = np.empty((N, N))
  for i in range(N):
    for j in range(N):
      if (i - c)**2 + (j - c)**2 < r_0**2/dh2:
        arr[i,j]=phi_0
      else:
        arr[i,j]=0.
  return arr

arr=create_arr(N,phi_0,c,r_0,dh2)

Сроки

#Pure Python:   58 ms
#Numba version: 0.086 ms (the first call takes longer and isn't included in the timings)

Answer 3

Если вы хотите использовать номера строк и столбцов для расчета, необходим цикл.Вы можете использовать одну петлю.У Numpy есть атрибут ndenumerate, который перебирает вашу матрицу.

def function(i, j):
    return phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0

for (i,j), value in np.ndenumerate(arr):
    arr[i, j] = function(i, j)