Реализуйте параметр оси в пользовательской функции

Question

Я пишу довольно тривиальную функцию для выполнения интеграции с применением правила трапеции в пространстве журнала.

Я бы хотел добавить параметр оси для достижения функциональности, аналогичной функции 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

Я решил эту проблему копированием подхода, который используется в 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 используются, явно создавая векторы "shape" с желаемой выходной формой.

Я думал, что это может быть достигнуто гораздо более коротким и аккуратным способом, но, очевидно, этот способ реализован в numpyсам по себе.

Answer 2

Чтобы проиллюстрировать исследование 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 общим.

Например, обобщенная индексация трехмерного массива:

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])