Блог

Я пытаюсь написать эффективную векторизованную пользовательскую функцию numpy, которая по существу выполняет поэлементное среднее значение без -1, если оно вообще присутствует.

Идея заключается в том, что список / кортеж ndarray того же размера в качестве входных данных создает один ndarray того же размера, который фактически является (поэлементным) средним значением всего ndarray, предоставленного в качестве входных данных. Единственными возможными значениями в этом массиве могут быть [-1, 0, 1]. Но когда вычисляется это пользовательское среднее значение, я хочу игнорировать -1 при вычислении среднего значения, если присутствуют другие значения. Метод для этого пользовательского значения mean_with_dont_knows_int .

Я пробовал универсальные функции, а также операции вдоль оси (после укладки массивов in), но по мере увеличения формы ввода время вычислений увеличивается, достигая нескольких минут при вводе относительно небольшой 500, 500, 10 формы.

Ниже приведено то, к чему я добрался до сих пор:

 from typing import Tuple

import numpy as np
import pytest


def mean_with_dont_knows_int(*argv: int):
    arr = np.array(argv, dtype=np.int8)
    return mean_with_dont_knows_from_1d(arr)


def mean_with_dont_knows_from_1d(arr: np.ndarray) -> int:
    arr = np.delete(arr, np.argwhere(arr == -1))
    if arr.size == 0:
        return -1
    return int(arr.sum() / len(arr) >= 0.5)


def mean_with_dont_knows(*argv: np.ndarray, use_ufunc: bool = True):
    if use_ufunc:
        return np.frompyfunc(mean_with_dont_knows_int, len(argv), 1)(*argv)
    return np.apply_along_axis(mean_with_dont_knows_from_1d, -1, np.stack(argv, axis=-1))


@pytest.mark.parametrize("shape", [(10, 10, 5), (50, 50, 10), (500, 500, 10), (500, 500, 20), ])
def test_mean_with_ufunc(shape: Tuple[int, int, int]):
    mask1 = np.ones(shape, dtype=np.int8)
    mask2 = np.zeros(shape, dtype=np.int8)
    mask3 = np.zeros(shape, dtype=np.int8)
    expected = np.zeros(shape)
    expected[0, 0, :] = 1
    mask3[0, 0, :] = -1
    result = mean_with_dont_knows(mask1, mask2, mask3)
    assert np.array_equal(result, expected) is True


@pytest.mark.parametrize("shape", [(10, 10, 5), (50, 50, 10), (500, 500, 10), (500, 500, 20), ])
def test_mean_with_axis(shape: Tuple[int, int, int]):
    mask1 = np.ones(shape, dtype=np.int8)
    mask2 = np.zeros(shape, dtype=np.int8)
    mask3 = np.zeros(shape, dtype=np.int8)
    expected = np.zeros(shape)
    expected[0, 0, :] = 1
    mask3[0, 0, :] = -1
    result = mean_with_dont_knows(mask1, mask2, mask3, use_ufunc=False)
    assert np.array_equal(result, expected) is True

  

Я наблюдаю снижение производительности по мере увеличения размера массива. Фактически линейное увеличение производительности .. как показано в результатах ниже:

 python -m pytest --durations=0  --capture=no  test.py                                                                              Exe on: 12:33:26 on 2020-08-16
=================================================================================================== test session starts ===================================================================================================
platform darwin -- Python 3.8.3, pytest-5.4.3, py-1.9.0, pluggy-0.13.1
rootdir: /Users/suneeta.mall/Documents/nearmap/data-science/planck
plugins: nbval-0.9.6, cov-2.10.0
collected 8 items

test.py ........

================================================================================================= slowest test durations ==================================================================================================
107.01s call     test.py::test_mean_with_axis[shape3]
104.04s call     test.py::test_mean_with_ufunc[shape3]
52.77s call     test.py::test_mean_with_axis[shape2]
52.32s call     test.py::test_mean_with_ufunc[shape2]
0.58s call     test.py::test_mean_with_ufunc[shape1]
0.54s call     test.py::test_mean_with_axis[shape1]
0.02s call     test.py::test_mean_with_ufunc[shape0]
0.01s call     test.py::test_mean_with_axis[shape0]
  

Очевидно, мне не хватает чего-то, что объяснит, почему я потратил 104,04 секунды на выполнение простого среднего значения над довольно маленьким массивом .. у кого-нибудь есть идеи / предложения / улучшения.

Редактировать:

Пересмотрели mean_with_dont_knows метод, основанный на предложении @hpaulj, который, похоже, отлично работает с массивами формы (500,500,10), но с (2000, 2000, 80) shape я вернулся к просмотру времени вычисления 55,5 x секунд.

 def mean_with_dont_knows(*argv: np.ndarray) -> np.ndarray:
    arr = np.stack(argv, axis=-1).astype(np.float)
    dont_know_mean = arr.mean(axis=-1).astype(np.int8)

    arr[arr == -1] = np.nan
    arr = np.nanmean(arr, axis=-1)
    arr = (arr >= .5).astype(np.int8)

    arr[dont_know_mean == -1] = -1
    return arr
  

Я немного упростил mean_with_dont_knows_from_1d :

 def mean_1(arr):
    arr1 = arr[arr!=-1]    # simpler than np.delete
    n = arr1.size
    if n==0:
        return -1
    return int(arr1.mean() >= 0.5)
  

Применяется с:

 def foo(*argv):
    temp = np.stack(argv, axis=-1).reshape(-1, len(argv))
    res = np.reshape([mean_1(row) for row in temp], argv[0].shape)
    return res
  

С помощью reshape я упростил 3D-итерацию до итерации по строкам 2d-массива.

Для shape = (10,10,5)

 In [30]: np.array_equal(expected, foo(mask1,mask2,mask3))                                            
Out[30]: True
  

время:

 In [31]: timeit foo(mask1,mask2,mask3)                                                               
11.8 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  

по сравнению с вашим:

 In [12]: timeit res = mean_with_dont_knows(mask1, mask2, mask3)                                      
29.5 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [13]: timeit res1 = mean_with_dont_knows(mask1, mask2, mask3, use_ufunc=False)                    
31.3 ms ± 25.8 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
  

Большая часть улучшений исходит от mean_1 функции.

 def foo1(*argv):
    temp = np.stack(argv, axis=-1).astype(float)
    temp[temp==-1] = np.nan
    res = np.nanmean(args, axis=-1)
    return res >=.5
  

Это использует преимущества np.nanmean одной из коллекций функций, которые маскируют np.nan значения. Вероятно, я мог бы замаскировать это -1 напрямую, используя идеи из np.nanmean , а пока просто использую существующий код. Значительное улучшение скорости:

 In [46]: np.array_equal(foo1(mask1,mask2,mask3), expected)                                           
Out[46]: True
In [47]: timeit foo1(mask1,mask2,mask3)                                                              
155 µs ± 217 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
  

В принципе, я избежал повторения 500 раз (для простого ‘mean’ из 3 элементов).