scipy.integrate.solve_ivp vectorized

Question

Trying to use the vectorized option for solve_ivp and strangely it throws an error that y0 must be 1 dimensional. MWE :

from scipy.integrate import solve_ivp
import numpy as np
import math

def f(t, y):
    theta = math.pi/4
    ham = np.array([[1,0],[1,np.exp(-1j*theta*t)]])
    return-1j * np.dot(ham,y)


def main():

    y0 = np.eye(2,dtype= np.complex128)
    t0 = 0
    tmax = 10**(-6)
    sol=solve_ivp( lambda t,y :f(t,y),(t0,tmax),y0,method='RK45',vectorized=True)
    print(sol.y)

if __name__ == '__main__':
    main()

The calling signature is fun(t, y). Here t is a scalar, and there are two options for the ndarray y: It can either have shape (n,); then fun must return array_like with shape (n,). Alternatively it can have shape (n, k); then fun must return an array_like with shape (n, k), i.e. each column corresponds to a single column in y. The choice between the two options is determined by vectorized argument (see below). The vectorized implementation allows a faster approximation of the Jacobian by finite differences (required for stiff solvers).

Error :

ValueError: y0 must be 1-dimensional.

Python 3.6.8

scipy.version '1.2.1'

Shields

Answer 1

The meaning of vectorize here is a bit confusing. It doesn't mean that y0 can be 2d, but rather that y as passed to your function can be 2d. In other words that func may be evaluated at multiple points at once, if the solver so desires. How many points is up to the solver, not you.

Change the f to show the shape a y at each call:

def f(t, y):
    print(y.shape)
    theta = math.pi/4
    ham = np.array([[1,0],[1,np.exp(-1j*theta*t)]])
    return-1j * np.dot(ham,y)

A sample call:

In [47]: integrate.solve_ivp(f,(t0,tmax),[1j,0],method='RK45',vectorized=False) 
(2,)
(2,)
(2,)
(2,)
(2,)
(2,)
(2,)
(2,)
Out[47]: 
  message: 'The solver successfully reached the end of the integration interval.'
     nfev: 8
     njev: 0
      nlu: 0
      sol: None
   status: 0
  success: True
        t: array([0.e+00, 1.e-06])
 t_events: None
        y: array([[0.e+00+1.e+00j, 1.e-06+1.e+00j],
       [0.e+00+0.e+00j, 1.e-06-1.e-12j]])

Same call, but with vectorize=True:

In [48]: integrate.solve_ivp(f,(t0,tmax),[1j,0],method='RK45',vectorized=True)  
(2, 1)
(2, 1)
(2, 1)
(2, 1)
(2, 1)
(2, 1)
(2, 1)
(2, 1)
Out[48]: 
  message: 'The solver successfully reached the end of the integration interval.'
     nfev: 8
     njev: 0
      nlu: 0
      sol: None
   status: 0
  success: True
        t: array([0.e+00, 1.e-06])
 t_events: None
        y: array([[0.e+00+1.e+00j, 1.e-06+1.e+00j],
       [0.e+00+0.e+00j, 1.e-06-1.e-12j]])

With False, the y passed to f is (2,), 1d; with True it is (2,1). I'm guessing it could be (2,2) or even (2,3) if the solver method so desires. That could speed up the execution, with fewer calls to f. In this case, it doesn't matter.

quadrature has a similar vec_func boolean parameter:

Numerical Quadrature of scalar valued function with vector input using scipy

A related bug/issue discussion:

https://github.com/scipy/scipy/issues/8922

Rodrickrodrigez

Thank you very much. That helps clear the confusion. So if I want to repeat the procedure for a large number of one-dimensional arrays, is looping through my best bet? It would be great if you can point out something that would work faster than looping through each array. I have (243) one-dimensional arrays each of size 243. So it is consuming a lot of time in this fashion. Thank you once again! – Shields

currently unanswered and links back here: What does the letter k mean in the documentation of solve_ivp function of Scipy? – Lamanna

@uhoh, as I demonstrate, if vectorized, k is at least 1, and may be larger. Other than making sure your function returns a like dimensioned result, there isn't much you can or have to do. If you come across a case where it is larger than 1 (eg in the jacobian evaluation) post an answer for all to learn from. – Rodrickrodrigez

@Rodrickrodrigez thanks, I was looking forward to some way that y0 could be 2d. When solve_ivp was first announced I got excited because I'd read something about dimensions not being 1d only, but now that I've started using it I see that that's not the case. – Lamanna

@uhoh, if for each initial value array the optimal ODE stepping is different, then solve_ivp can't perform the calculation in a 2d manner. Each has to be solved separately. Or you could make y0 n*k long. For some problem it might be faster that way, for others k iterations on the size n y0 might be. How were you hoping the 2d case would work? – Rodrickrodrigez

@Rodrickrodrigez A simple example and a more complicated/representative one For strictly convenience/housekeeping reasons I'd like my orbital mechanics state vectors to have the shape (2, 3, N) where the first index is for position vs velocity, the second index is the number of spatial dimensions (2 or 3) and the third dimension is the number of particles or planets. I've always wondered if the .reshape() lines inside deriv() were impacting speed since they are called frequently but have no evidence. – Lamanna

I changed my answer to What does the letter k mean in the documentation of solve_ivp function of Scipy?. Hope this is helpful. – Horwitz