在自定义功能中实现轴参数

时间:2019-02-20 19:32:57

标签: python numpy numpy-broadcasting

我正在编写一个相当琐碎的函数,以在日志空间中应用梯形规则进行积分。

我想添加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中,但我不确定如何实现

2 个答案:

答案 0 :(得分:0)

为了说明在交互式会话中进行的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])

答案 1 :(得分:0)

我解决了复制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)

为了实现“可广播性”,使用reshapeslice的混合,显式创建具有所需输出形状的“形状”矢量。

我本以为可以用更短,更整洁的方式实现这一目标,但显然这是在numpy本身中实现的方式。