Search code examples
pythonscipyode

Mass plotting ODE curves that are solved from a list of parameter values

My code is an ODE system that I can plot curves for. However, I wish to have 100 different solutions generated, for which I can then plot all 100 as curves onto one figure. My code is:

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

N = 1000

I0, R0 = 1, 0

S0 = N - I0 - R0
J0 = I0

beta, gamma = 2/7, 1/7

t = np.linspace(0, 160, 160+1)


def deriv(y, t, N, beta, gamma):
    S, I, R, J = y
    dS = ((-beta * S * I) / N)
    dI = ((beta * S * I) / N) - (gamma * I)
    dR = (gamma * I)
    dJ = ((beta * S * I) / N)
    return dS, dI, dR, dJ


solve = odeint(deriv, (S0, I0, R0, J0), t, args=(N, beta, gamma))
S, I, R, J = solve.T

J_diff = np.diff(J)

fig = plt.figure(facecolor='w')
ax = fig.add_subplot(111, facecolor='#dddddd', axisbelow=True)
ax.plot(t[1:], J_diff, 'blue', alpha=1, lw=2, label='Daily incidence')
ax.set_xlabel('Time in days')
ax.set_ylabel('Number (1000s)')
ax.grid(b=True, which='major', c='w', lw=2, ls='-')
legend = ax.legend()
legend.get_frame().set_alpha(0.5)
plt.show()

This yields a curve such as this

enter image description here

However, I wish to calculate the ODE equation J for a list of random 100 beta values, all other parameters kept fixed, which I can plot so I get 100 curves. What I have tried is using something like:

empty = []
for i in range(100):
    empty.append(random.uniform(1.5, 2.5)*gamma)

for the new list of beta values, and then do

solns =  []
for empt in empty:
    ces = odeint(deriv, (S0, I0, R0, J0), t, args=(N, empty, gamma))
    solns.append(ces)

but this is not correct. I don't know if I'm thinking about it the wrong way or syntax is incorrect as opposed to the logic. Is it even possible to take each beta value in the list empty, calculate the ODE system for it, and the plot each solution on to one figure?

Solution

  • Your approach is correct. The problem is that you are passing a list of beta values instead of a single scalar due to a minor typo in the args argument. It should be

    solns =  []
for empt in empty:
    ces = odeint(deriv, (S0, I0, R0, J0), t, args=(N, empt, gamma))
    solns.append(ces)

    Then, you can proceed similarly to your code snippet, i.e.

    J_diffs = []
for sol in solns:
    S, I, R, J = sol.T
    J_diffs.append(np.diff(J))

# plot all the solutions
fig = plt.figure(facecolor='w')
ax = fig.add_subplot(111, facecolor='#dddddd', axisbelow=True)
ax.set_xlabel('Time in days')
ax.set_ylabel('Number (1000s)')
ax.grid(b=True, which='major', c='w', lw=2, ls='-')

for J_diff in J_diffs:
    ax.plot(t[1:], J_diff, 'blue', alpha=1, lw=2, label='Daily incidence')

# plot without legend
plt.show()