Как заполнить 3D-массив в numpy функциональными значениями с той же скоростью, что и Java?

Пожалуйста, не используйте .vectorize(..), это неэффективно, так как он все равно будет выполнять обработку на уровне Python. .vectorize() следует использовать только в крайнем случае, если, например, функция не может быть вычислена "в большом объеме", потому что ее "структура" слишком сложна.

Но вам не нужно использовать здесь .vectorize, вы можете реализовать свою функцию для работы с массивами с помощью:

r_xy = np.sqrt(xx ** 2 + yy ** 2)
density = (r_xy <= R0) & (-hz <= zz) & (zz <= hz)

или даже немного быстрее:

r_xy = xx * xx + yy * yy
density = (r_xy <= R0 * R0) & (-hz <= zz) & (zz <= hz)

Это создаст массив логических значений 2000200020. Мы можем использовать:

intdens = density.astype(int)

чтобы построить массив из ints.

Распечатать массив здесь довольно сложно, но всего он содержит 2'356'047 единиц:

>>> density.astype(int).sum()
2356047

Контрольные показатели. Если я запускаю это локально 10 раз, я получаю:

>>> timeit(f, number=10)
18.040479518999973
>>> timeit(f2, number=10)  # f2 is the optimized variant
13.287886952000008

Таким образом, в среднем мы вычисляем эту матрицу (включая приведение к ints) за 1,3–1,8 секунды.

Вы также можете использовать скомпилированную версию функции для расчета плотности. Вы можете использовать для этого cython или numba. Я использую numba для jit-компиляции функции вычисления плотности в ans, поскольку это так же просто, как вставить декоратор .

Плюсы:

• Вы можете написать if условия, которые упоминаете в своих комментариях.
• Немного быстрее, чем версия numpy, упомянутая в ответах @Willem Van Onsem, поскольку нам нужно перебирать логический массив для вычисления суммы в density.astype(int).sum().

Минусы:

• Напишите уродливую трехуровневую петлю. Теряет красоту уникального решения для лайнеров.

Код:

import numba as nb
@nb.jit(nopython=True, cache=True)
def calc_density(xx, yy, zz, R0, hz):
    threshold = R0 * R0
    dimensions = xx.shape

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):
                r_xy = xx[i][j][k] ** 2 + yy[i][j][k] ** 2
    
                if(r_xy <= threshold and -hz <= zz[i][j][k] <= hz):
                    density+=1
    return density

Продолжительность:

Решение Виллема Ван Онсема, вариант f2: 1,28 с без суммы, 2,01 с суммой.

Решение Numba (calc_de density, при втором запуске, чтобы уменьшить время компиляции): 0,48 с.

Как предлагается в комментариях, нам также не нужно рассчитывать сетку. Мы можем напрямую передать x, y, z функции. Таким образом:

@nb.jit(nopython=True, cache=True)
def calc_density2(x, y, z, R0, hz):
    threshold = R0 * R0
    dimensions = len(x), len(y), len(z)

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):
                
                r_xy = x[i] ** 2 + y[j] ** 2
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1
    return density

Теперь, для честного сравнения, мы также включаем время np.meshgrid в ответ @Willem Van Onsem. Продолжительность:

Решение Виллема Ван Онсема, вариант f2 (включая время np.meshgrid): 2,24 с

Решение Numba (calc_de density2, при втором запуске, чтобы уменьшить время компиляции): 0,079 с.

Это подразумевается как пространный комментарий к ответу Дипака Шайни.

Главное изменение - не использовать координаты, сгенерированные np.meshgrid, которые содержат ненужные повторения. Это не рекомендуется, если вы можете этого избежать (как с точки зрения использования памяти, так и производительности).

Код

import numba as nb
import numpy as np

@nb.jit(nopython=True,parallel=True)
def calc_density_2(x, y, z,R0,hz):
    threshold = R0 * R0

    density = 0
    for i in nb.prange(y.shape[0]):
        for j in range(x.shape[0]):
            r_xy = x[j] ** 2 + y[i] ** 2
            for k in range(z.shape[0]):
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1

    return density

Время

R0 = 500
hz = 1

x = np.arange(-1000, 1000, 1)
y = np.arange(-1000, 1000, 1)
z = np.arange(-10, 10, 1)

xx, yy, zz = np.meshgrid(x, y, z)

#after the first call (compilation overhead)
#calc_density_2          9.7 ms
#calc_density_2 parallel 3.9 ms
#@Deepak Saini           115 ms