Как сделать цикл for для массива только с одним элементом (массив 0-d)

Question

У меня есть функция, которая принимает аргумент y_0, который может быть списком или массивом, или может быть просто с плавающей точкой. Я хотел, чтобы это работало в случае, когда y_0 было одним числом с плавающей запятой, поэтому я попытался использовать np.asarray(y_0) в приведенном ниже коде, чтобы, если был только один элемент, цикл продолжал работать. Однако я получил ошибку TypeError: iteration over a 0-d array. Я мог бы использовать оператор if, чтобы проверить, является ли он одиночным или нет, и предпринять соответствующие действия. Тем не менее, мне было любопытно, есть ли способ заставить его перебирать один объект?

def vf_grapher(fn, t_0, t_n, dt, y_0, lintype='-r', sup_title=None,
               title=None, xlab=None, ylab=None):

    t = np.arange(t_0, t_n, dt)
    y_min = .0
    y_max = .0

    fig, axs = plt.subplots()
    fig.suptitle(sup_title)

    axs.set_title(title)
    axs.set_ylabel(ylab)
    axs.set_xlabel(xlab)

    for iv in np.asarray(y_0):
        soln = rk4(dt, t, fn, iv)
        plt.plot(t, soln, lintype)
        if y_min > np.min(soln):
            y_min = np.min(soln)
        if y_max < np.max(soln):
            y_max = np.max(soln)

Для минимального рабочего примера включите следующую функцию:

def rk4(dt, t, field, y_0):
"""
:param dt: float - the timestep
:param t: array - the time mesh
:param field: method - the vector field y' = f(t, y)
:param y_0: array - contains initial conditions
:return: ndarray - solution
"""

# Initialize solution matrix. Each row is the solution to the system
# for a given time step. Each column is the full solution for a single
# equation.
y = np.asarray(len(t) * [y_0])

for i in np.arange(len(t) - 1):
    k1 = dt * field(t[i], y[i])
    k2 = dt * field(t[i] + 0.5 * dt, y[i] + 0.5 * k1)
    k3 = dt * field(t[i] + 0.5 * dt, y[i] + 0.5 * k2)
    k4 = dt * field(t[i] + dt, y[i] + k3)
    y[i + 1] = y[i] + (k1 + 2 * k2 + 2 * k3 + k4) / 6

return y

if __name__ == '__main__':
    def f(t, x): return x**2 - x

    vf_grapher(f, 0, 4, 0.1, (-0.9, 0.5, 1.01), xlab='t', ylab='x(t)',
               sup_title=r'Solution Field for $\dot{x} = x^2 - x$')

Answer 1

Вы можете использовать аргумент ndmin для np.array, чтобы убедиться, что ваш массив фактически повторяемый:

np.array(y_0, ndmin=1, copy=False)

np.asarray - это просто псевдоним для np.array, который устанавливает некоторые аргументы по умолчаниюпо-разному.

ndmin может использоваться, чтобы дополнить вашу форму единичными размерами. Это помогает, потому что обычно нульмерные массивы эквивалентны скалярам. Один раздражающий побочный эффект этого - то, что они не повторяемы. ndmin=1 означает, что скалярные входные данные должны рассматриваться как одномерный, одноэлементный массив, что вам и нужно.

copy=False просто говорит numpy использовать существующие массивы как есть. вместо того, чтобы делать копию. Таким образом, если пользователь передает фактический массив (в отличие от списка или скаляра), он будет использоваться без дублирования данных. Я часто связываю этот аргумент с subok=True, который также будет проходить через подклассы ndarray без копирования.

Answer 2

Я не уверен, почему вы можете перебирать список одного объекта, но не массив одного объекта, но я нашел способ определить, является ли элемент итеративным или нет в этом ответе на другой вопрос:https://stackoverflow.com/a/1952481/3696204

Затем я использовал попытку, кроме блока следующим образом:

try:
    iter(y_0)
except TypeError:
    y_0 = list([y_0])

for iv in y_0:
    soln = rk4(dt, t, fn, iv)
    plt.plot(t, soln, lintype)
    if y_min > np.min(soln):
        y_min = np.min(soln)
    if y_max < np.max(soln):
        y_max = np.max(soln)

Спасибо всем за полезные комментарии.