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

Question

我正在编写一个相当简单的函数来执行在日志空间中应用梯形规则的集成。

我想添加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)

为了实现“可广播性”，采用reshape和slice的混合，明确地创建具有所需输出形状的“形状”矢量。

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