How to implement 'solve_ivp' with vectorized='True' in python

Question

I've been trying to solve a set of differential equations using solve_ivp. The Jacobian matrix of the system is the A as you can see below. I wanted to enable the option vectorized='True' but unfortunately i do not know how to modify the present code to vectorize the Jacobian matrix A. Does anyone know how this can be done?

# imports
import numpy as np
import scipy.sparse as sp
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# grid sizing
R=0.05  #sphere radius
N=1000#number of points
D=0.00002 #diffusion coefficient
k=10  # Arrhenius 
Cs=1.0 # Boundary concentration
C0=0.0 # Initial concentration
time_constant=R**2.0/D

dr=R/(N-1)

# Algebra simplification
a=D/dr**2

Init_conc=np.linspace(0,0,N)
B=np.zeros(N)
B[N-1]=Cs*(a+a/(N-1))
#
e1 = np.ones(N)
e2 = np.ones(N)
e3 = np.ones(N)
#
#
#
e1[0]=-k-6*a
e1[1:]=-k-2*a
# 
#
e2[1]=6*a
for i in range(2,N) :
 e2[i]=a+a/(i-1)
#
#
#
for i in range (0,N-1) :
 e3[i]=a-a/(i+1)


A = sp.spdiags([e3,e1,e2],[-1,0,1],N,N,format="csc")


def dc_dt(t,C)   :
    dc=A.dot(C)+B
    return dc

# Solving the system, I want to implement the same thing with vectorized='True'
OutputTimes=np.linspace(0,0.2*time_constant,100)
ans=solve_ivp(dc_dt,(0,0.2*time_constant),Init_conc,method='RK45',t_eval=OutputTimes,vectorized='False')
print (ans)

Potential [duplicate](https://stackoverflow.com/questions/54991056/scipy-integrate-solve-ivp-vectorized). — Patol75, Nov 29 '19 at 05:09

score 2 · Accepted Answer · answered Nov 29 '19 at 05:13

Please have a look at this answer, the explanation is thorough. For your code in particular, please see below for updated snippet and figure. It is not obvious that vectorize is providing any speed-up. However, providing A for the keyword jac makes a difference. But I guess it is only valid if A is constant?

# imports
import numpy as np
import scipy.sparse as sp
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt  # noqa


def dc_dt(t, C):
    print(C.shape)
    if len(C.shape) == 1:
        return np.squeeze(A.dot(C)) + B
    else:
        return A.dot(C) + np.transpose(np.tile(B, (C.shape[1], 1)))
    # return np.squeeze(A.dot(C)) + B


# grid sizing
R = 0.05  # sphere radius
N = 1000  # number of points
D = 0.00002  # diffusion coefficient
k = 10  # Arrhenius
Cs = 1.0  # Boundary concentration
C0 = 0.0  # Initial concentration
time_constant = R**2.0 / D
dr = R / (N - 1)
# Algebra simplification
a = D / dr**2
Init_conc = np.repeat(0, N)
B = np.zeros(N)
B[-1] = Cs * (a + a / (N - 1))
e1 = np.ones(N)
e2 = np.ones(N)
e3 = np.ones(N)
e1[0] = -k - 6 * a
e1[1:] = -k - 2 * a
e2[1] = 6 * a
for i in range(2, N):
    e2[i] = a + a / (i - 1)
for i in range(0, N - 1):
    e3[i] = a - a / (i + 1)
A = sp.spdiags([e3, e1, e2], [-1, 0, 1], N, N, format="csc")
# Solving the system, I want to implement the same thing with vectorized='True'
OutputTimes = np.linspace(0, 0.2 * time_constant, 10000)
ans = solve_ivp(dc_dt, (0, 0.2 * time_constant), Init_conc,
                method='BDF', t_eval=OutputTimes, jac=A, vectorized=True)
plt.plot(np.arange(N), ans.y[:, 0])
plt.plot(np.arange(N), ans.y[:, 1])
plt.plot(np.arange(N), ans.y[:, 10])
plt.plot(np.arange(N), ans.y[:, 20])
plt.plot(np.arange(N), ans.y[:, 50])
plt.plot(np.arange(N), ans.y[:, -1])
plt.show()

Yes, A is constant. That solves my problem. Thanks for the suggestion — Giannis1996, Nov 29 '19 at 06:09

How to implement 'solve_ivp' with vectorized='True' in python

1 Answers1