Numpy.apply_along_axis 在应用带有 if else 条件的函数时意外工作

Numpy.apply_along_axis works unexpectedly when applying a function with if else condition

我无法获得意想不到的结果。下面的代码以最简单的方式重现了结果(f只是一个测试函数):


#returns absolute difference between last and first element in an array
def f(arr): 
    return 0 if arr[-1] == arr[0] else abs(arr[-1]-arr[0])

def test_vectorized(test_arr, window = 2):
    T = test_arr.shape[0]
    #create sliding windows
    slide_windows = np.expand_dims(np.arange(window+1), axis=0) + np.expand_dims(np.arange(T - window), axis=0).T
    print(slide_windows)
    slide_values = test_arr[slide_windows]
    print(slide_values)
    #apply function to each sliding window
    return np.apply_along_axis(f, axis=1, arr=slide_values)
#testing
test_arr = np.array([27.75, 27.71, 28.05, 27.75, 26.55,27.18])
test_vectorized(test_arr, window=3)
#Output
[[0 1 2 3]
 [1 2 3 4]
 [2 3 4 5]]
[[27.75 27.71 28.05 27.75]
 [27.71 28.05 27.75 26.55]
 [28.05 27.75 26.55 27.18]]
Out[238]:
array([0, 1, 0])

代码应该returnarray([0,1.16,0.87]),即每个滑动数组中第一个元素和最后一个元素之间的绝对差。
我正在使用带有 python 3.8.2 的 Jupyter 笔记本。调试了一个多小时,代码本身好像没有问题。有人能帮忙吗?非常感谢。

你的函数 f returns 整数。 你必须使用:

def f(arr):
    return float(0 if arr[-1] == arr[0] else abs(arr[-1]-arr[0]))

[[0 1 2 3]
 [1 2 3 4]
 [2 3 4 5]]
[[27.75 27.71 28.05 27.75]
 [27.71 28.05 27.75 26.55]
 [28.05 27.75 26.55 27.18]]
[0.   1.16 0.87]

P.S

您的函数 f 可以概括为简单的 return abs(arr[-1]-arr[0]),因为它涵盖了 0 的情况。您不需要 if 语句。