Заполнить массив numy на основе индексов клеток

Question

Я пытаюсь создать 3d-массив, чьи записи в ячейках рассчитываются по индексам ячеек.В частности, я хочу эту ячейку (i,j,k) = sqrt(i+j+k).

Это очень легко сделать с помощью следующего для циклов:

N=10
A=np.zeros((N,N,N))

for i in range(N):
    for j in range(N):
        for k in range(N):
            A[i][j][k] = np.sqrt(i+j+k)

Мне было интересно, есть ли у numpy встроенные функции, которые делают их вложенными в циклыизбыточный.

Answer 1

вы можете использовать np.arange, а затем np.newaxis, чтобы создать различные размеры.С помощью простых sum и np.sqrt для выполнения работы после:

arr = np.arange(N)
A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])

Вы получите тот же результат:

N = 10
arr = np.arange(N)
A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])
B = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))
print ((A==B).all())
#True

Этот метод немного быстрее, чем при использовании np.ogrid:

N = 10
%timeit arr = np.arange(N); A = np.sqrt(arr + arr[:,np.newaxis]+ arr[:,np.newaxis,np.newaxis])
#18.6 µs ± 3.29 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))
#58.5 µs ± 8.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 2

Это быстрее для больших N, но может считаться мошенничеством; -)

Для полного использования квадратного корня используются все преимущества очень регулярного и повторяющегося шаблона.

def cheat(N):
    values = np.sqrt(np.arange(3*N-2))
    result = np.lib.stride_tricks.as_strided(values, (N, N, N), 3*values.strides)
    return np.ascontiguousarray(result)

Если вы можете жить с несмежным представлением только для чтения (всеми практическими средствами), это может быть значительно быстрее:

def cheat_nc_view(N):
    values = np.sqrt(np.arange(3*N-2))
    return np.lib.stride_tricks.as_strided(values, (N, N, N), 3*values.strides)

Для справки:

def cheek(N):
    arr = np.arange(N)
    return np.sqrt(arr + arr[:,np.newaxis] + arr[:,np.newaxis,np.newaxis])

>>> np.all(cheek(20) == cheat(20))
True
>>> np.all(cheek(200) == cheat_nc_view(200))
True

Сроки:

>>> timeit(lambda: cheek(20), number=1000)
0.05387042500660755
>>> timeit(lambda: cheat(20), number=1000)
0.020798540994292125
>>> timeit(lambda: cheat_nc_view(20), number=1000)
0.010791150998556986

>>> timeit(lambda: cheek(200), number=100)
6.823299437994137
>>> timeit(lambda: cheat(200), number=100)
2.0583883369981777
>>> timeit(lambda: cheat_nc_view(200), number=100)
0.0014881940151099116

Answer 3

Самый простой и производительный вариант - с открытыми сетками, используя np.ogrid, а затем выполняя соответствующую операцию (и) -

I,J,K = np.ogrid[:N,:N,:N]
A = np.sqrt(I+J+K)

Или с np.sum для широковещательных сумм открытых сеток для одноголайнер -

A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))

Соответствует: General workflow on vectorizing loops involving range iterators.