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

Question

Я пытаюсь написать функцию, которая вычисляет интеграл Симпсона для заданных входных данных. Части этих входов могут быть в форме float или np.array.

В настоящее время моя функция выглядит следующим образом:

s = np.array([0.2,0.3,0.4])

@np.vectorize
def recovery(s,R0,n):

    x = np.linspace(s,1,n)
    y = 1/(x*R0(x))
    ans = simp(y,x)

    return ans

Это работает нормально и возвращает:

array([a, b, c]) ## Where a, b and c are the correct answers

Однако, если я хочу получить доступ только к первому элементу в массиве внутри функции, с помощью:

return ans[0]

я получаю сообщение об ошибке:

IndexError: invalid index to scalar variable

Я также получаю всевозможные ошибки при попытке поиграть с выводом «ans», например:

return type(ans)
ValueError: invalid __array_struct__

return len(ans)
TypeError: object of type 'numpy.float64' has no len()

Цель функции - вернуть float, если вход s равен число с плавающей точкой или array, если s ввод np.ndarray.

Я полный python новичок, поэтому любая помощь будет принята с благодарностью.

Answer 1

Чтобы получить доступ только к первому элементу возвращаемого массива, вам нужно сделать это вне функции, а не внутри, потому что внутри функции ans по-прежнему является скаляром.

Итак, while

return ans[0]

не работает,

recovery(s,R0,n)[0]

или

ans = recovery(s,R0,n)
ans[0]

будет работать.

Непонятно то, что декоратор @np.vectorize меняется функция из одного, который отображает скаляры в скаляры, в функцию, которая отображает векторы в векторы. Но эта манипуляция выполняется «вне» функции recovery; внутри исходного исходного кода функции все выглядит так, как будто это все еще функция на скалярах.

Вот почему, когда вы написали return ans[0], Python жаловались, что вы пытались проиндексировать скаляр; внутри функции , ans - скаляр.