Melhoria de desempenho na transformação numpy array

Oct 31 2020

Dados três numpyarrays 1D, quero transformá-los da seguinte maneira:

import numpy as np

Xd = np.asarray([0, 0,   1,   1,   0.5])
Yd = np.asarray([0, 0,   0,   2.5, 2.5])
Zd = np.asarray([0, 1.5, 1.5, 1.5, 1.5])

points = np.stack([Xd, Yd, Zd], axis=1).reshape(-1, 1, 3)
segments = np.concatenate([points[:-1], points[1:]], axis = 1)    

print(segments.shape)
print(segments)

Resultado:

(4, 2, 3)
[[[0.  0.  0. ]
  [0.  0.  1.5]]

 [[0.  0.  1.5]
  [1.  0.  1.5]]

 [[1.  0.  1.5]
  [1.  2.5 1.5]]

 [[1.  2.5 1.5]
  [0.5 2.5 1.5]]]

Existe uma maneira de melhorar o desempenho dessa transformação?

fundo

Esta transformação é necessária para usar as XYZcoordenadas em matplotlibcom Line3DCollection. Até agora, só vi variações do código acima, mas com milhares de coordenadas ou dados interpolados para melhor resolução, uma abordagem otimizada é necessária.

Resumo

Graças a @Mercury , pode-se concluir que para arrays mais curtos (<1k de comprimento) a resposta de @Miguel tem melhor desempenho, mas a abordagem de @mathfux escala muito melhor quando os arrays ficam mais longos.

Respostas

3 mathfux Oct 31 2020 at 13:03

Parece que você está tentando rolar uma janela de forma (2, 3)em um array 2D. Isso é semelhante à convolução da imagem, que pode ser feita de np.lib.stride_tricksuma forma muito eficiente.

a = np.transpose([Xd, Yd, Zd])
window = (2, 3)
view_shape = (len(a) - window[0] + 1,) + window # (4,2,3) if len(a) == 5
sub_matrix = np.lib.stride_tricks.as_strided(a, shape = view_shape, strides = (a.itemsize,) + a.strides)
>>> sub_matrix
array([[[0. , 0. , 0. ],
        [0. , 0. , 1.5]],

       [[0. , 0. , 1.5],
        [1. , 0. , 1.5]],

       [[1. , 0. , 1.5],
        [1. , 2.5, 1.5]],

       [[1. , 2.5, 1.5],
        [0.5, 2.5, 1.5]]])

Observe que np.lib.stride_trickstem um ótimo desempenho em relação a quaisquer formas alternativas.

4 Miguel Oct 31 2020 at 12:22

Como um conselho geral, quando você deseja velocidade, geralmente deve tentar evitar empilhar e concatenar, pois isso geralmente significa copiar os mesmos dados várias vezes.

Enfim, aqui está como eu faria, um código um pouco mais longo, mas não faz mais trabalho do que o necessário

n = len(Xd)
segments = np.empty((n-1, 2, 3))

segments[:,0,0] = Xd[:-1]
segments[:,1,0] = Xd[1:]

segments[:,0,1] = Yd[:-1]
segments[:,1,1] = Yd[1:]

segments[:,0,2] = Zd[:-1]
segments[:,1,2] = Zd[1:]

[EDIT] - O seguinte foi feita para a ciência / fun, não reproduc e

Então, tentei ver se conseguia um pouco mais de desempenho com a resposta de @mathfux , e obtive este código feio:

a = np.empty(3*n)
a[:n]    = Xd
a[n:n+n] = Yd
a[n+n:]  = Zd

interface = dict(a.__array_interface__)
interface['shape'] = (n-1, 2, 3)
interface['strides'] = (a.itemsize, a.itemsize, n*a.itemsize)
segments= np.array(np.lib.stride_tricks.DummyArray(interface, base=a), copy=False)

Na minha máquina, é mensuravelmente mais rápido (até ~ 30% dependendo do tamanho da entrada). Os ganhos são em parte devido à construção ae omissão das verificações deas_strided

1 Mercury Oct 31 2020 at 13:03

Aqui estão alguns testes de tempo, em matrizes maiores, o que torna a diferença mais clara.

import numpy as np
from timeit import timeit

# original
def f1(x, y, z):
    points = np.stack([x, y, z], axis=1).reshape(-1, 1, 3)
    return np.concatenate([points[:-1], points[1:]], axis = 1)

# preallocating and then assigning
def f2(x, y, z):
    segments = np.empty((len(x)-1, 2, 3))

    segments[:,0,0] = x[:-1]
    segments[:,1,0] = x[1:]

    segments[:,0,1] = y[:-1]
    segments[:,1,1] = y[1:]

    segments[:,0,2] = z[:-1]
    segments[:,1,2] = z[1:]
    return segments

# stacking, but in one go
def f3(x, y, z):
    segments = np.stack([x[:-1], y[:-1], z[:-1], x[1:], y[1:],z[1:]], axis=1)
    return segments.reshape(-1, 2, 3)

# list comparison
def f4(x, y, z):
    z_ = [i for i in zip(x,y,z)]
    return [[[z_[i]],[z_[i+1]]] for i in range(len(z_)-1)]

#np.lib.stride_tricks approach
def f5(x, y, z):
    a = np.transpose([x, y, z])
    window = (2, 3)
    view_shape = (len(a) - window[0] + 1,) + window # (4,2,3) if len(a) == 5
    return np.lib.stride_tricks.as_strided(a, shape = view_shape, strides = (a.itemsize,) + a.strides)
    

ntime = 5000 #number of test runs
nxd = 500    #array length

Xd = np.random.randn(nxd)
Yd = np.random.randn(nxd)
Zd = np.random.randn(nxd)

print(timeit(lambda: f1(Xd, Yd, Zd), number=ntime))
#0.11369249999999999

print(timeit(lambda: f2(Xd, Yd, Zd), number=ntime))
#0.0480651

print(timeit(lambda: f3(Xd, Yd, Zd), number=ntime))
#0.10202380000000003

print(timeit(lambda: f4(Xd, Yd, Zd), number=ntime))
#1.8407391

print(timeit(lambda: f5(Xd, Yd, Zd), number=ntime))
#0.09132560000000023
    
ntime = 50     #number of test runs
nxd = 500000   #array length

Xd = np.random.randn(nxd)
Yd = np.random.randn(nxd)
Zd = np.random.randn(nxd)

print(timeit(lambda: f1(Xd, Yd, Zd), number=ntime))
#1.7519548999999999

print(timeit(lambda: f2(Xd, Yd, Zd), number=ntime))
#1.504727

print(timeit(lambda: f3(Xd, Yd, Zd), number=ntime))
#1.5010566

print(timeit(lambda: f4(Xd, Yd, Zd), number=ntime))
#22.6208157

print(timeit(lambda: f5(Xd, Yd, Zd), number=ntime))
#0.46465339999999955

Como você pode ver, o jeito do @Miguel é o melhor: pré-alocar o array e então atribuir é a maneira mais eficiente de fazer isso. Mesmo que você os empilhe de maneira mais inteligente, como em f3 (), ainda é mais lento do que f2 (). Mas nada supera f5 () quando o comprimento do array aumenta substancialmente.

venky__ Oct 31 2020 at 12:31

Achei isso mais rápido do que o código do @Miguel.

z = [i for i in zip(Xd,Yd,Zd)]
segments = [[[z[i]],[z[i+1]]] for i in range(len(z)-1)]