Plotting the graphs of many differential equations using JiTCODE in Python

Question

I would like to plot the graphs of a system of differential equations for individuals connected in a network in Python 3.6. The system of equations is as follows:

dx_i/dt = omega_i - mu_i * x_i + epsilon_i * x_i * y_i

dy_i/dt = r_i * y_i - gamma * x_i * y_i + \sum_{j!=1} A_{ji} *y_i


x_i(t) is the antibody response in the i-th individual
y_i(t) is the viral charge in that individual where i = 1,....,n
omega_i is the rate of production and/or transport of antibodies
mu_i is the death rate of antibodies
epsilon_i is the rate of production of antibodies induced by a unit viral 
population
r_i is the intrinsic growth rate of viral population
gamma_i is the rate of destruction of viruses by a unit antibody population
A_{ji} is the ji-th of a matrix representing the strength of transmission from j to i

I have written the code for the immune response to viral invasion for n individuals connected in a network.

The model represents a system of coupled equations representing the antibody and viral population in a network of connected individuals.

from jitcode import jitcode, y
import numpy as np
import sympy
import matplotlib.pyplot as plt
from numpy.random import uniform
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D


n = 5
alpha = 0.05


beta = uniform(0.01,3.0,n)
beta.sort()
mu = uniform(0.01,3.0,n)
mu.sort()
epsilon = uniform(0.01,3.0,n)
epsilon.sort()
gamma = uniform(0.01,3.0,n)
gamma.sort()
omega = uniform(0.01,3.0,n)
omega.sort()
r = uniform(0.01,3.0,n)
r.sort()


# Knonecker delta

M = np.einsum('ij-> ij',np.eye(n,n))

print(M)

# The transmission matrix A whose elements represent the strength of 
# transmission from j to i depending of spatial factors. 

A = beta * M * sympy.exp(-alpha) 
      
print(A)



def X(i): return y(2*i)
def Y(i): return y(2*i+1)


def f():
    for i in range(n):
    
        coupling_sum = sum(A[i,j]*Y(j) for j in range(n) if j!=i ) 
   
        yield omega[i] - mu[i] * X(i) + epsilon[i] * X(i) * Y(i)
        yield r[i] * Y(i) - gamma[i] * X(i) * Y(i) + coupling_sum
    
    

    
#integrate
#---------------

t = np.linspace(0, 100)
T = np.row_stack([t, t])

initial_state = np.random.random(2*n)

ODE = jitcode(f, n=2*n)
ODE.set_integrator("dopri5", atol=1e-6,rtol=0)
ODE.generate_f_C(simplify=False, do_cse=False, chunk_size=150)
ODE.set_initial_value(initial_state,time=0.0)

#data structure: x[0], w[0], v[0], z[0], ..., x[n], w[n], v[n], z[n]
data = []
data = np.vstack(ODE.integrate(T) for T in range(10, 100, 10))
print(data)


# Plotting the graphs
#-----------------------

plt.show()
plt.savefig('tmp.pdf'); plt.savefig('tmp.png')
plt.title("The Immunoepidemiological model")
plt.plot(t, f)
plt.xlabel('t')
plt.ylabel('f')
fig = plt.figure()

I am expecting to get graphs of antibody and viral population over time t. However, i am getting the following error message.

ValueError: x and y must have same first dimension, but have shapes (50,) and (1,)

Answer 1

Most of your code is fine except you are not plotting your the correct thing. f is an abstract description of your differential equations, not your solution.

Your integrated data is contained in data variable, which has the shape (timepoints, system_variables). Now, for instance, if you want to plot the antibody response and viral population as a function of time for the first individual, your plot command should be: plt.plot(t, data[:,0:2]).

Moreover, you are not using t as the actual times for integration. Following the Jitcode documentation, it should be something like:

t = np.arrange(0,100,0.1)
data = np.vstack(ODE.integrate(time) for time in t)