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

Question

Я попытался смоделировать воксели трехмерного цилиндра с помощью следующего кода:

import math
import numpy as np

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)


def density_f(x, y, z):
    r_xy = math.sqrt(x ** 2 + y ** 2)
    if r_xy <= R0 and -hz <= z <= hz:
        return 1
    else:
        return 0


density = np.vectorize(density_f)(xx, yy, zz)

, и для его вычисления потребовалось много минут.

Эквивалентный субоптимальный код Java выполняется 10-15 секунд.

Как заставить Python вычислять воксели с одинаковой скоростью?Где оптимизировать?

Answer 1

Вы также можете использовать скомпилированную версию функции для расчета плотности. Вы можете использовать Cython или Numba для этого. Я использую numba , чтобы скомпилировать функцию вычисления плотности в ANS, так же просто, как вставить декоратор.

Плюсы :

• Вы можете написать if условия, которые упоминаете в комментариях
• Чуть быстрее, чем просто обалденная версия, упомянутая в ответе @ Виллем Ван Онсем, так как мы должны перебрать логический массив для рассчитать сумму в 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_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. Продолжительность :

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

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

Answer 2

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

Основное изменение состоит в том, чтобы не использовать координаты, генерируемые 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

Answer 3

Пожалуйста, сделайте , а не , используйте .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)

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

intdens = density.astype(int)

для создания массива int s.

Печать массива здесь довольно сложна, но в общей сложности она содержит 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

Итак, в среднем мы вычисляем эту матрицу (включая приведение ее кint с) за 1,3-1,8 с.