Issue with vectorization in python-CodePudding

I have an application where I need to merge two solutions obtained from the scipy.integrate.solve_ivp in python. The actual application is a bit more complicated, but the following example shows the idea:

from scipy.integrate import solve_ivp import numpy as np

def lotkavolterra(t, z, a, b, c, d):

    x, y = z

    return [a*x - b*x*y, -c*y   d*x*y]


sol_1 = solve_ivp(lotkavolterra, [0, 10], [10, 5], args=(1.5, 1, 3, 1), dense_output=True).sol  
sol_2 = solve_ivp(lotkavolterra, [10, 15], [10, 5], args=(1.5, 1, 3, 1), dense_output=True).sol

def sol_comb(t):

    if t <= 10:
        return sol_1(t)
    else:
        return sol_2(t)

I want to be able to use the merged or combined solution sol_comb on numpy arrays. Hence I tried to define a vectorized solution as follows:

sol_comb_vect = np.vectorize(sol_comb)

The following code, where I only call the functions on scalars, works fine:

print("sol_1 for t = 2",sol_1(2)) 
print("sol_2 for t = 11",sol_2(11))
print("sol_comb for t = 11",sol_comb(11))
print("sol_comb_vect for t = 11",sol_comb_vect(11))

The individual solutions sol_1 and sol_2 are apparently vectorized, since the following works fine:

print("sol_1 for t = [2,3]",sol_1(np.array([2])))
print("sol_2 for t = [11,13]",sol_2(np.array([11,13])))

However, if I call the non-vectorized function sol_comb on an array, as in the following example, I get the expected ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all():

print("sol_comb for t = [11,13]",sol_comb(np.array([11,13])))

I was expecting the vectorized version sol_comb_vect to work. However, in the following, I get the error ValueError: setting an array element with a sequence.

print("sol_comb_vect for t = [11,13]",sol_comb_vect(np.array([11,13])))

Any ideas how to fix this?

I would also be happy to merge the two OdeSolution instances in a cleaner way. In principle I think this should be possible, by using the time values and interpolants for sol_1 and sol_2, respectively.

CodePudding user response：

I think you need to specify the signature for your output when you vectorize your function since by default the pyfunc you pass to np.vectorize() is assumed to take scalars as input and output see doc. And I assume that your ValueError is caused by that. So try this:

sol_comb_vect = np.vectorize(sol_comb, signature='()->(n)')
sol_comb_vect(np.array([2, 11, 13]))

output:

array([[0.60031288, 0.09618044],
       [0.21298705, 1.36999868],
       [2.58274789, 0.01857732]])

I don't know if this is the expected output tho. I hope this answers your question.