Variable type of y and subscripting y in Scipy's solve_ivp

Question

I am trying to solve a matrix differential equation in Python using solve_ivp from Scipy, and I have run into some problems.

The code returning dy/dt that I have written is the following:

def fun(y, t, J, param):
    [g, k, p, N] = param
    x = np.reshape(y[:N], (N,))
    A = np.reshape(y[N+1:], (N, N))
    W = J + A
    phiX = np.tanh(x)
    dx = -x + np.matmul(W, phiX)
    dA = -A + (k/N) * np.outer(phiX, phiX)
    return np.concatenate((dx, (1/p) * dA), axis=None)

The code to get the solution was the following (where J and param were defined earlier):

y0 = np.random.uniform(0, 1, (N*(N + 1),))
sol = solve_ivp(fun, (0, 300), y0 = y0, first_step = 0.1, args = (J, param))

This was done because, so far as I know, solve_ivp cannot take matrices as y. To fix this problem, I converted an N-dimensional vector and an (N, N) matrix into an N * (N+1) dimensional vector. However, when I tried to reconstruct x and A back from y, I got an error saying 'float' object is not subscriptable. Initial conditions y0 are also set as an N * (N+1) dimensional vector.

I have tried to figure out what I was doing wrong, because I know that Scipy can solve N-dimensional differential equations, but I could not figure out why. Why isn't the type of y np.ndarray and why cannot I splice the input y?

Answer 1

The signature of the model function has changed from odeint and solve_ivp:

def fun(t, y, J, param):
    [g, k, p, N] = param
    x = np.reshape(y[:N], (N,))
    A = np.reshape(y[N+1:], (N, N))
    W = J + A
    phiX = np.tanh(x)
    dx = -x + np.matmul(W, phiX)
    dA = -A + (k/N) * np.outer(phiX, phiX)
    return np.concatenate((dx, (1/p) * dA), axis=None)

Just swap y and t in the signature.