Итерация по последовательным 1-D слайсам вдоль произвольной оси массива NumPy - PullRequest
Новая
       47

Итерация по последовательным 1-D слайсам вдоль произвольной оси массива NumPy

0 голосов
/ 29 апреля 2019

Я пишу пакет Python, который выполняет различные сложные задачи статистического анализа вдоль произвольной оси массива numpy произвольной формы.

В настоящее время, так что форма и ось массива могут быть произвольными, япросто переставьте массив, чтобы ось интереса была помещена в дальнюю RHS, и раздавите оси LHS в одну.Например, если форма массива (3,4,5), и мы хотим выполнить некоторую операцию вдоль оси 1, она преобразуется в форму (15,4), операции выполняются вдоль оси -1, затем она преобразуется обратно.в форме (3,4,5) и возвращается функцией.

Я чувствую, что этот подход может быть излишне медленным из-за всех этих манипуляций с массивами.Есть ли способ, которым я могу чисто перебрать все, кроме одного измерения массива?То есть в приведенном выше примере это будет идти [0,:,0], [0,:,1], ..., [2,:,3], [2,:,4], но опять же это должно работать для произвольной формы массива и положения оси.

Может бытьnp.ndenumerate, np.ndindex и np.take могут быть использованы для этого каким-либо образом?

Редактировать: Есть ли способ сделать это с np.nditer?Возможно, это может соответствовать скорости перестановки / изменения формы.

1 Ответ

0 голосов
/ 30 апреля 2019

Оказывается, просто транспонирование и изменение формы действительно быстрее. Поэтому я думаю, что ответ ... не делайте этого, предпочтительнее переставить и изменить, как я уже делал.

Вот код из моего проекта. 

# Benchmark
f = lambda x: x # can change this to any arbitrary function
def test1(data, axis=-1):
    # Test the lead flatten approach
    data, shape = lead_flatten(permute(data, axis))
    output = np.empty(data.shape)
    for i in range(data.shape[0]): # iterate along first dimension; each row is an autocor
        output[i,:] = f(data[i,:]) # arbitrary complex equation
    return unpermute(lead_unflatten(output, shape), axis)
def test2(data, axis=-1):
    # Test the new approach
    output = np.empty(data.shape)
    for d,o in zip(iter_1d(data, axis), iter_1d(output, axis)):
        o[...] = f(d)
    return output

# Iterator class
class iter_1d(object):
    def __init__(self, data, axis=-1):
        axis = (axis % data.ndim) # e.g. for 3D array, -1 becomes 2
        self.data = data
        self.axis = axis
    def __iter__(self):
        shape = (s for i,s in enumerate(self.data.shape) if i!=self.axis)
        self.iter = np.ndindex(*shape)
        return self
    def __next__(self):
        idx = self.iter.next()
        idx = [*idx]
        idx.insert(self.axis, slice(None))
        return self.data[idx]

# Permute and reshape functions
def lead_flatten(data, nflat=None):
    shape = list(data.shape)
    if nflat is None:
        nflat = data.ndim-1 # all but last dimension
    if nflat<=0: # just apply singleton dimension
        return data[None,...], shape
    return np.reshape(data, (np.prod(data.shape[:nflat]).astype(int), *data.shape[nflat:]), order='C'), shape # make column major

def lead_unflatten(data, shape, nflat=None):
    if nflat is None:
        nflat = len(shape) - 1 # all but last dimension
    if nflat<=0: # we artificially added a singleton dimension; remove it
        return data[0,...]
    if data.shape[0] != np.prod(shape[:nflat]):
        raise ValueError(f'Number of leading elements {data.shape[0]} does not match leading shape {shape[nflat:]}.')
    if not all(s1==s2 for s1,s2 in zip(data.shape[1:], shape[nflat:])):
        raise ValueError(f'Trailing dimensions on data, {data.shape[1:]}, do not match trailing dimensions on new shape, {shape[nflat:]}.')
    return np.reshape(data, shape, order='C')

def permute(data, source=-1, destination=-1):
    data = np.moveaxis(data, source, destination)
    return data

def unpermute(data, source=-1, destination=-1):
    data = np.moveaxis(data, destination, source)
    return data

А вот результаты некоторых %timeit операций. 

import numpy as np
a = np.random.rand(10,20,30,40)
%timeit -r10 -n10 test1(a, axis=2) # around 12ms
%timeit -r10 -n10 test2(a, axis=2) # around 22ms
...