Применение функции Numba guvectorize во временном измерении 3D-массива с помощью apply_ufunc из Xarray - PullRequest
Новая
       63

Применение функции Numba guvectorize во временном измерении 3D-массива с помощью apply_ufunc из Xarray

1 голос
/ 26 мая 2020

У меня есть некоторые проблемы с тем, чтобы это работало должным образом, и я также открыт для других предложений, поскольку я не уверен на 100%, правильно ли я иду с этим.

Вот несколько простых фиктивные данные: 

times = pd.date_range(start='2012-01-01',freq='1W',periods=25)
x = np.array([range(0,20)]).squeeze()
y = np.array([range(20,40)]).squeeze()

data = np.random.randint(3, size=(25,20,20))

ds = xr.DataArray(data, dims=['time', 'y', 'x'], coords = {'time': times, 'y': y, 'x': x})

Для каждой координаты x, y я хочу вернуть самую длинную последовательность 1 или 2 с течением времени. Итак, мой входной массив - 3D (время, x, y), а мой выходной - 2D (x, y). Код в 'seq_gufun c' вдохновлен этим потоком . Мой фактический набор данных намного больше (с классами landuse вместо 1s, 2s, et c), и это лишь небольшая часть большого рабочего процесса, где я также использую dask для параллельной обработки. Так что, в конце концов, это должно работать быстро и эффективно, поэтому я закончил тем, что попытался выяснить, как заставить numba @guvectorize и Xarray apply_ufun c работать вместе: 


@guvectorize(
    "(int64[:], int64[:])",
    "(n) -> (n)", target='parallel', nopython=True
)
def seq_gufunc(x, out):

    f_arr = np.array([False])

    bool_stack = np.hstack((f_arr, (x == 1) | (x == 2), f_arr))

    # Get start, stop index pairs for sequences
    idx_pairs = np.where(np.diff(bool_stack))[0].reshape(-1, 2)

    # Get length of longest sequence
    longest_seq = np.max(np.diff(idx_pairs))

    out[:] = longest_seq


## Input for dim would be: 'time' 
def apply_seq_gufunc(data, dim):

    return xr.apply_ufunc(seq_gufunc,
                          data,
                          input_core_dims=[[dim]],
                          exclude_dims=set((dim,)), 
                          dask="allowed")

Вероятно, есть очень очевидные ошибки, которые, надеюсь, кто-то сможет указать. Мне сложно понять, что на самом деле происходит в фоновом режиме, и как мне настроить строку макета @guvectorize и параметры apply_ufun c, чтобы он делал то, что я хочу.

EDIT2: это рабочее решение. См. Ответ @OriolAbril для получения дополнительной информации о параметрах apply_ufunc и guvectorize. Также было необходимо реализовать предложение if...else... на случай, если значения не совпадают, и избежать ошибки ValueError, которая может возникнуть. 

@guvectorize(
    "(int64[:], int64[:])",
    "(n) -> ()", nopython=True
)
def seq_gufunc(x, out):

    f_arr = np.array([False])
    bool_stack = np.hstack((f_arr, (x == 1) | (x == 2), f_arr))

    if np.sum(bool_stack) == 0:
        longest_seq = 0

    else:
        # Get start, stop index pairs for sequences
        idx_pairs = np.where(np.diff(bool_stack))[0].reshape(-1, 2)

        # Get length of longest sequence
        longest_seq = np.max(np.diff(idx_pairs))   

    out[:] = longest_seq


def apply_seq_gufunc(data, dim):

    return xr.apply_ufunc(seq_gufunc,
                          data,
                          input_core_dims=[[dim]],
                          dask="parallelized",
                          output_dtypes=['uint8']
                         )

1 Ответ

1 голос
/ 26 мая 2020

Я бы указал вам на Как применить xarray u_function к NetCDF и вернуть 2D-массив (несколько новых переменных) в DataSet , ближайшая цель не та же, но подробный описание и примеры должны прояснить проблему.

В частности, вы правы, используя time как input_core_dims (чтобы убедиться, что он перемещен в последнее измерение), и он правильно отформатирован как список списков, однако вы это делаете не нужно excluded_dims а output_core_dims==[["time"]].

Выходной файл имеет ту же форму, что и входной, однако, как объяснено в ссылке выше, apply_ufunc ожидает, что он будет иметь такую ​​же форму, как широковещательные затемнения . output_core_dims необходимо, чтобы apply_ufunc ожидал вывода с димами y, x, time.

...