Эффективное построение двумерного массива из 1d массива и функция с условиями

Question

Я создаю 2-мерный массив из функции, примененной к 1-мерному массиву (который содержит условные выражения), и хотел бы узнать более эффективный способ сделать это. В настоящее время это самая медленная часть моего кода. x является 1d массивом numpy, а вывод представляет собой 2d массив numpy. Существует переключатель для создания другого элемента массива в зависимости от того, меньше ли x или больше 0. В будущем число переключателей может быть произвольным.

def basis2(x) :

  final = []
  for i in x :
    if i > 0 :
        xr = 2.0*(i-0.5)
        final.append(np.array([0.0, 0.0, 0.0, 0.5*xr*(xr-1.0),-1.0*(xr+1)*(xr-1), 0.5*xr*(xr+1.0)]))
    else :
        xl = 2.0*(i+0.5)
        final.append(np.array([0.5*xl*(xl-1.0),-1.0*(xl+1)*(xl-1),0.5*xl*(xl+1.0),0.0,0.0,0.0]))


return np.array(final)

В идеале я мог быустранить цикл for - но до сих пор мне не удалось сделать это должным образом, используя, например, «где». Спасибо за любую помощь.

Answer 1

Вот векторизованный подход, который использует симметрию:

col = np.s_[...,None]

def basis2_v(x):
    h,w = np.arange(x.size)[col],np.arange(3,dtype=np.int8)
    # if the terms for xr are the same as the formulas for xl applied to -xr
    # we can therefore unify the code and apply it to |x|:
    # there are three factors in total each term consists of two of them
    aux = (np.abs(x[col])-w/2)*(2/(1^-(w&1)))
    # the easiest is multiplying all three and then dividing one out again
    aux = -aux.prod(-1)[col]/aux
    # fix degenerate cases
    aux[np.isnan(aux)] = 1
    # finally, we need to embed the terms in a zero matrix
    out = np.zeros((x.size,6),x.dtype)
    # the xor trick maps 2,1,0 to -3,-2,-1 if x>0
    out[h,(-(x[col]>0).view(np.int8))^w[::-1]] = aux
    return out

# make small test
x = np.random.randn(10)
# this should pass
assert np.allclose(basis2(x),basis2_v(x))

Answer 2

С вашей функцией:

In [247]: basis2(np.array([1,.5,0,-.5,-1]))                                     
Out[247]: 
[array([ 0.,  0.,  0.,  0., -0.,  1.]),
 array([ 0.,  0.,  0., -0.,  1.,  0.]),
 array([ 0., -0.,  1.,  0.,  0.,  0.]),
 array([-0.,  1.,  0.,  0.,  0.,  0.]),
 array([ 1.,  0., -0.,  0.,  0.,  0.])]
In [248]: %hist 245                                                             
basis2_1(np.array([1,.5,0,-.5,-1]))

С некоторыми незначительными изменениями:

def basis2_1(x) :
    xr = 2.0*(x[x>0]-0.5)
    res1 = np.array([0.0*xr, 0.0*xr, 0.0*xr, 0.5*xr*(xr-1.0),-1.0*(xr+1)*(xr-1), 0.5*xr*(xr+1.0)])
    xl = 2.0*(x[x<=0]+0.5)
    res2 = np.array([0.5*xl*(xl-1.0),-1.0*(xl+1)*(xl-1),0.5*xl*(xl+1.0),0.0*xl,0.0*xl,0.0*xl])
    return res1, res2

In [250]: basis2_1(np.array([1,.5,0,-.5,-1]))                                   
Out[250]: 
(array([[ 0.,  0.],
        [ 0.,  0.],
        [ 0.,  0.],
        [ 0., -0.],
        [-0.,  1.],
        [ 1.,  0.]]), 
 array([[ 0., -0.,  1.],
        [-0.,  1.,  0.],
        [ 1.,  0., -0.],
        [ 0.,  0., -0.],
        [ 0.,  0., -0.],
        [ 0.,  0., -0.]]))

Соединение двух подмассивов:

In [251]: np.hstack(_)                                                          
Out[251]: 
array([[ 0.,  0.,  0., -0.,  1.],
       [ 0.,  0., -0.,  1.,  0.],
       [ 0.,  0.,  1.,  0., -0.],
       [ 0., -0.,  0.,  0., -0.],
       [-0.,  1.,  0.,  0., -0.],
       [ 1.,  0.,  0.,  0., -0.]])

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

Например, вы можете создать массив result = np.zeros((5,x.shape[0])), просто вставив соответствующие ненулевые элементы (сохранив все эти 0.0*xr условия).

Глядя на эти блоки в Out[251]:

In [257]: x = np.array([1,.5,0,-.5,-1])                                         
In [258]: Out[251][3:,np.nonzero(x>0)[0]]                                       
Out[258]: 
array([[ 0., -0.],
       [-0.,  1.],
       [ 1.,  0.]])
In [259]: Out[251][:3,np.nonzero(x<=0)[0]]                                      
Out[259]: 
array([[ 0., -0.,  1.],
       [-0.,  1.,  0.],
       [ 1.,  0., -0.]])