在自定义函数中实现轴参数

Question

我正在编写一个相当简单的函数来执行对数 space 中应用梯形法则的积分。

我想添加 axis 参数来实现类似于 numpy.trapz 函数的功能，但我对如何正确实现它有点困惑。

不可广播函数如下所示：

import numpy as np

def logtrapz(y, x):

    logx = np.log(x)
    dlogx = np.diff(logx)

    logy = np.log(y)
    dlogy = np.diff(logy)

    b = dlogx + dlogy
    a = np.exp(logx + logy)

    dF = a[:-1] * (np.exp(b) - 1)/b * dlogx

    return np.sum(dF)

这适用于一维输入。

我认为解决方案在于 numpy.expand_dims，但我不太确定如何实施它

Answer 1

为了在交互式会话中说明 slice 探索：

In [216]: slice(None)                                                           
Out[216]: slice(None, None, None)
In [217]: slice??                                                               
Init signature: slice(self, /, *args, **kwargs)
Docstring:     
slice(stop)
slice(start, stop[, step])

Create a slice object.  This is used for extended slicing (e.g. a[0:10:2]).
Type:           type
Subclasses:     
In [218]: np.s_[:]                                                              
Out[218]: slice(None, None, None)

我没有看过 np.trapz 代码，但我知道其他 numpy 函数通常在需要 axis 通用时构造索引元组。

例如 3d 数组的通用索引：

In [221]: arr = np.arange(24).reshape(2,3,4)                                    
In [223]: idx = [slice(None) for _ in range(3)]                                 
In [224]: idx                                                                   
Out[224]: [slice(None, None, None), slice(None, None, None), slice(None, None, None)]
In [225]: idx[1]=1                                                              
In [226]: idx                                                                   
Out[226]: [slice(None, None, None), 1, slice(None, None, None)]
In [227]: tuple(idx)                                                            
Out[227]: (slice(None, None, None), 1, slice(None, None, None))
In [228]: arr[tuple(idx)]     # arr[:,1,:]                                                  
Out[228]: 
array([[ 4,  5,  6,  7],
       [16, 17, 18, 19]])
In [229]: idx[2]=2                                                              
In [230]: arr[tuple(idx)]     # arr[:,1,2]                                                  
Out[230]: array([ 6, 18])

Answer 2

我通过复制 numpy.trapz 中使用的方法解决了这个问题。这有点复杂，但效果很好。

对于未来的读者，上述函数的广播版本是

import numpy as np

def logtrapz(y, x, axis=-1):

    x = np.asanyarray(x)
    logx = np.log(x)
    if x.ndim == 1:
        dlogx = np.diff(logx)
        # reshape to correct shape
        shape1 = [1]*y.ndim
        shape1[axis] = dlogx.shape[0]
        shape2 = [1]*y.ndim
        shape2[axis] = logx.shape[0]
        dlogx = dlogx.reshape(shape1)
        logx  = logx.reshape(shape2)
    else:
        dlogx = np.diff(x, axis=axis)

    nd = y.ndim
    slice1 = [slice(None)]*nd
    slice2 = [slice(None)]*nd
    slice1[axis] = slice(None, -1)
    slice2[axis] = slice(1, None)
    slice1 = tuple(slice1)
    slice2 = tuple(slice2)

    logy = np.log(y)
    dlogy = logy[slice2] - logy[slice1]

    b = dlogx + dlogy
    a = np.exp(logx + logy)

    dF = a[slice1] * (np.exp(b) - 1)/b * dlogx

    np.sum(dF, axis=axis)

为了实现 "broadcastability"，采用了 reshape 和 slice 的混合，显式创建具有所需输出形状的 "shape" 向量。

我认为这可以通过更短更简洁的方式来实现，但显然这是在 numpy 本身中实现的方式。

在自定义函数中实现轴参数

Implement axis parameter in a custom function

python

numpy

array-broadcasting