Как заполнить двумерный массив numpy без циклов?

Question

Я хочу быстро и без циклов заполнить массив numpy N, этот массив хранит в себе сколько "соседей" у клетки в игре жизнь, вот код чистого python:

def compute_neighbours(Z):
    start_time = pygame.time.get_ticks()
    shape = Z.shape
    N = [[0 for i in range(shape[1])] for i in range(shape[0])]
    
    for x in range(1, shape[0]-1):
        for y in range(1, shape[1]-1):
            N[x][y] = Z[x - 1][y - 1] + Z[x][y - 1] + Z[x + 1][y - 1] \
                    + Z[x - 1][y    ]               + Z[x + 1][y    ] \
                    + Z[x - 1][y + 1] + Z[x][y + 1] + Z[x + 1][y + 1]

    print("compute neighbours ", (pygame.time.get_ticks() - start_time)/1000)
    return N

Я хочу сделать тоже самое, но без циклов и быстрее чистого питона, но я незнаю как это сделать, я попробовал с циклами, но это почему то работает в разы медленнее

N = np.asarray([[np.sum(Z[x-1:x+2, y-1:y+2])-Z[x, y] for y in range(0, shape[1])] for x in range(0, shape[0])])

Answer 1

Нужно именно при помощи numpy? Или и другой библиотекой подойдёт?
Например, scipy.

import numpy as np
from scipy.signal import convolve2d

def compute_neighbours_scipy(Z):
    return convolve2d(Z, np.ones((3,3), dtype=int), 'same') - Z

Причём эта функция рассчитает правильно и для элементов на границах матрицы.

UPDATE
Варианты на чистом numpy.

def compute_neighbours_numpy2(Z):
    mask = np.asarray([[1,1,1], [1,0,1], [1,1,1]])
    s = mask.shape + tuple(np.subtract(Z.shape, mask.shape) + 1)
    sub = np.lib.stride_tricks.as_strided(Z, shape = s, strides = Z.strides * 2)
    return np.pad(np.einsum('ij,ijkl->kl', mask, sub), [(1, 1), (1, 1)], mode='constant')

Вариант от @CrazyElf c корректным расчётом для границ матрицы.

def compute_neighbours_numpy_CrazyElf(Z):
    ZZ = np.pad(Z, [(1, 1), (1, 1)], mode='constant')
    N = ZZ[:-2,  :-2] + ZZ[1:-1,  :-2] + ZZ[2:,  :-2] \
      + ZZ[:-2, 1:-1]                  + ZZ[2:, 1:-1] \
      + ZZ[:-2, 2:  ] + ZZ[1:-1, 2:  ] + ZZ[2:, 2:  ]
    return N

Измерения скорости выполнения:

import numpy as np
from scipy.signal import convolve2d
import numba
import timeit

def compute_neighbours(Z):
    shape = Z.shape
    N = [[0 for i in range(shape[1])] for i in range(shape[0])]
    
    for x in range(1, shape[0]-1):
        for y in range(1, shape[1]-1):
            N[x][y] = Z[x - 1][y - 1] + Z[x][y - 1] + Z[x + 1][y - 1] \
                    + Z[x - 1][y    ]               + Z[x + 1][y    ] \
                    + Z[x - 1][y + 1] + Z[x][y + 1] + Z[x + 1][y + 1]
    return N

@numba.njit
def compute_neighbours_numba(Z):
    shape = Z.shape
    N = [[0 for i in range(shape[1])] for i in range(shape[0])]
    
    for x in range(1, shape[0]-1):
        for y in range(1, shape[1]-1):
            N[x][y] = Z[x - 1][y - 1] + Z[x][y - 1] + Z[x + 1][y - 1] \
                    + Z[x - 1][y    ]               + Z[x + 1][y    ] \
                    + Z[x - 1][y + 1] + Z[x][y + 1] + Z[x + 1][y + 1]
    return N

def compute_neighbours_numpy(Z):
    return np.asarray([[np.sum(Z[x-1:x+2, y-1:y+2])-Z[x, y] for y in range(0, Z.shape[1])] for x in range(0, Z.shape[0])])

@numba.njit
def compute_neighbours_numpy_numba(Z):
    return np.asarray([[np.sum(Z[x-1:x+2, y-1:y+2])-Z[x, y] for y in range(0, Z.shape[1])] for x in range(0, Z.shape[0])])

def compute_neighbours_numpy2(Z):
    mask = np.asarray([[1,1,1], [1,0,1], [1,1,1]])
    s = mask.shape + tuple(np.subtract(Z.shape, mask.shape) + 1)
    sub = np.lib.stride_tricks.as_strided(Z, shape = s, strides = Z.strides * 2)
    return np.pad(np.einsum('ij,ijkl->kl', mask, sub), [(1, 1), (1, 1)], mode='constant')

def compute_neighbours_numpy_CrazyElf(Z):
    ZZ = np.pad(Z, [(1, 1), (1, 1)], mode='constant')
    N = ZZ[:-2,  :-2] + ZZ[1:-1,  :-2] + ZZ[2:,  :-2] \
      + ZZ[:-2, 1:-1]                  + ZZ[2:, 1:-1] \
      + ZZ[:-2, 2:  ] + ZZ[1:-1, 2:  ] + ZZ[2:, 2:  ]
    return N

def compute_neighbours_scipy(Z):
    return convolve2d(Z, np.ones((3,3), dtype=int),'same') - Z


matrix = np.random.randint(2, size=(2000,2000))
print("list          :", timeit.timeit("compute_neighbours(matrix)"               , globals=globals(), number=2))
print("list + numba  :", timeit.timeit("compute_neighbours_numba(matrix)"         , globals=globals(), number=2))
print("numpy         :", timeit.timeit("compute_neighbours_numpy(matrix)"         , globals=globals(), number=2))
print("numpy + numba :", timeit.timeit("compute_neighbours_numpy_numba(matrix)"   , globals=globals(), number=2))
print("numpy2        :", timeit.timeit("compute_neighbours_numpy2(matrix)"        , globals=globals(), number=2))
print("numpy CrazyElf:", timeit.timeit("compute_neighbours_numpy_CrazyElf(matrix)", globals=globals(), number=2))
print("scipy         :", timeit.timeit("compute_neighbours_scipy(matrix)"         , globals=globals(), number=2))

# matrix = np.random.randint(2, size=(10,10))
# print(matrix)
# print(*compute_neighbours(matrix), sep="\n")
# print(compute_neighbours_numpy2(matrix))
# print(compute_neighbours_numpy_CrazyElf(matrix))
# print(compute_neighbours_numpy(matrix))
# print(compute_neighbours_scipy(matrix))

Результаты измерений:

list          : 19.943651329
list + numba  : 0.7527782960000025
numpy         : 62.82599582
numpy + numba : 0.8693800110000041
numpy2        : 0.2166252049999997
numpy CrazyElf: 0.1425725370000066
scipy         : 0.3862877149999946

Answer 2

Ну вот чистый numpy, хотя границы тут не считаются, а, наверное, должны. Я просто переделал вашу же функцию для чистого питона в векторизованный вариант:

def compute_neighbours_numpy(Z):
    N = np.zeros_like(Z)
    N[1:-1,1:-1] = Z[:-2,:-2] + Z[1:-1,:-2] + Z[:-2,1:-1] \
                 + Z[2:,:-2]                + Z[2:,1:-1] \
                 + Z[:-2,2:]  + Z[1:-1,2:]  + Z[2:,2:]
    return N

Нужно всего то было пересчитать индексы в срезы правильным образом и выкинуть циклы.Выполняется порядка 0.1с вообще в тех же условиях, что другие функции у GrAnd. Но с краями нужно что-то делать.

БЛОГ НА HUSL

Как заполнить двумерный массив numpy без циклов?

Ответы (2 шт):