sir model with age groups implemented in python

sir_python

#!/usr/bin/env python

import numpy as np
from pylab import *
import scipy.integrate as spi

#Parameter Values
PopIn= {'g1s_young':1/12,'g1i_young':1/12,'g1r_young':1/12,'g1s_mid':1/12,'g1i_mid':1/12,'g1r_mid':1/12,'g1s_old':1/12,'g1i_old':1/12,'g1r_old':1/12,'g1s_vold':1/12,'g1i_vold':1/12,'g1r_vold':1/12}
arr = [1/12]*12
beta= 0.50
gamma=1/10.



class G1():
    def __init__(self):
        self.age_groups = ["young","mid","old","vold"]
        self.betas = {"young":0.1,"mid":0.2,"old":0.3,"vold":0.4}
        self.gamma = 0.3




t_end = 100
t_start = 1
t_step = .02
t_interval = np.arange(t_start, t_end, t_step)
g1 = G1()
#Solving the differential equation. Solves over t for initial conditions PopIn

def eq_system(PolpIn,t):
    '''Defining SIR System of Equations'''
    #Creating an array of equations
    Eqs= np.zeros((9))
    d_Eqs = {}
    for group in g1.age_groups: #S equations:
        #Eqs[0]= -beta * PopIn[0]*PopIn[1]
        d_Eqs[group+"_S"] = - PopIn['g1s_'+group]*sum([g1.betas[age]*PopIn['g1s_'+age] for age in g1.age_groups])
    for group in g1.age_groups: #I equations:
        #Eqs[1]= beta * PopIn[0]*PopIn[1] - gamma*PopIn[1]
        d_Eqs[group + "_I"] = PopIn['g1s_'+group]*sum([g1.betas[age]*PopIn['g1s_'+age] for age in g1.age_groups]) - g1.gamma*PopIn['g1s_'+group]
    for group in g1.age_groups: #R equations:
        #Eqs[2]= gamma*PopIn[1]
        d_Eqs[group + "_R"] = g1.gamma * PopIn['g1s_' + group]
    Eqs = list(d_Eqs.values())
    return Eqs

SIR = spi.odeint(eq_system, arr, t_interval)

#Splitting out the curves for S, I and R from each other, in case they need
#to be used seperately
S=(SIR[:,0])
I=(SIR[:,1])
R=(SIR[:,2])

#Create a new array of the same length to be used as the x-axis for a plot
x=arange(len(SIR),dtype=float)

#Scale x-axis array by the step size
for i in x:
    x[i]=(x[i]*t_step)

#Stack S, I and R with the x-axis
SIR_plot= vstack([S,I,R,x])

#Graph!
fig= figure()
ax = fig.add_subplot(111)
plot(SIR_plot[3],SIR_plot[0],'g--',SIR_plot[3],SIR_plot[1],'r-',SIR_plot[3],SIR_plot[2],'-.b',linewidth=3)
xlabel("Time (days)")
ylabel("Percent of Population")
title("SIR Epidemic")
grid(True)
legend(("S", "I", "R"), shadow=True, fancybox=True)
show()

for row in SIR:
    print(sum(row))

sir_python_with_groups.py

#!/usr/bin/env python

import numpy as np
from pylab import *
import scipy.integrate as spi

#Parameter Values
PopIn= {'g1s_young':1/24,'g1i_young':1/24,'g1r_young':1/24,'g1s_mid':1/24,'g1i_mid':1/24,'g1r_mid':1/24,'g1s_old':1/24,'g1i_old':1/24,'g1r_old':1/24,'g1s_vold':1/24,'g1i_vold':1/24,'g1r_vold':1/24,'g2s_young':1/24,'g2i_young':1/24,'g2r_young':1/24,'g2s_mid':1/24,'g2i_mid':1/24,'g2r_mid':1/24,'g2s_old':1/24,'g2i_old':1/24,'g2r_old':1/24,'g2s_vold':1/24,'g2i_vold':1/24,'g2r_vold':1/24}
arr = [1/24]*24
beta= 0.50
gamma=1/10.



class G1():
    def __init__(self):
        self.age_groups = ["young","mid","old","vold"]
        self.betas = {"young":0.1,"mid":0.2,"old":0.3,"vold":0.4}
        self.gamma = 0.3
        self.G1_cm = np.array([[1/16]*4]*4)
    def __str__(self):
        return 'g1'

class G2():
    def __init__(self):
        self.age_groups = ["young","mid","old","vold"]
        self.betas = {"young":0.2,"mid":0.3,"old":0.4,"vold":0.5}
        self.gamma = 0.4
        self.G1_cm = np.array([[1/16]*4]*4)
    def __str__(self):
        return 'g2'




t_end = 100
t_start = 1
t_step = .02
t_interval = np.arange(t_start, t_end, t_step)
g1 = G1()
g2 = G2()
e_groups = [g1,g2]
#Solving the differential equation. Solves over t for initial conditions PopIn

def eq_system(PolpIn,t):
    '''Defining SIR System of Equations'''
    #Creating an array of equations
    Eqs= np.zeros((18))
    d_Eqs = {}
    for e_group in e_groups:
        for age_group in e_group.age_groups: #S equations:
            #Eqs[0]= -beta * PopIn[0]*PopIn[1]
            d_Eqs[str(e_group)+'_'+age_group+"_S"] = - PopIn[str(e_group)+'s_'+age_group]*sum([sum([group.betas[age]*PopIn[str(e_group)+'s_'+age] for age in group.age_groups]) for group in e_groups])
        #I equations:
            #Eqs[1]= beta * PopIn[0]*PopIn[1] - gamma*PopIn[1]
            d_Eqs[str(e_group)+'_'+age_group + "_I"] = PopIn[str(e_group)+'s_'+age_group]*sum([sum([group.betas[age]*PopIn[str(e_group)+'s_'+age] for age in group.age_groups]) for group in e_groups]) - g1.gamma*PopIn[str(e_group)+'s_'+age_group]
        #R equations:
            #Eqs[2]= gamma*PopIn[1]
            d_Eqs[str(e_group)+'_'+age_group + "_R"] = g1.gamma * PopIn[str(e_group)+'s_' + age_group]
    Eqs = list(d_Eqs.values())
    return Eqs

SIR = spi.odeint(eq_system, arr, t_interval)

#Splitting out the curves for S, I and R from each other, in case they need
#to be used seperately
S=(SIR[:,0])
I=(SIR[:,1])
R=(SIR[:,2])

#Create a new array of the same length to be used as the x-axis for a plot
x=arange(len(SIR),dtype=float)

#Scale x-axis array by the step size
for i in x:
    x[i]=(x[i]*t_step)

#Stack S, I and R with the x-axis
SIR_plot= vstack([S,I,R,x])

#Graph!
fig= figure()
ax = fig.add_subplot(111)
plot(SIR_plot[3],SIR_plot[0],'g--',SIR_plot[3],SIR_plot[1],'r-',SIR_plot[3],SIR_plot[2],'-.b',linewidth=3)
xlabel("Time (days)")
ylabel("Percent of Population")
title("SIR Epidemic")
grid(True)
legend(("S", "I", "R"), shadow=True, fancybox=True)
show()

for row in SIR:
    print(sum(np.abs(row)))

print(SIR)

This piece code is giving me an error

for i in x:
    x[i]=(x[i]*t_step)

IndexError: only integers, slices (:), ellipsis (...), numpy.newaxis (None) and integer or boolean arrays are valid indices

any idea how I can solve this?

EDIT:
I changed the line above in x=arange(len(SIR))

The problem in this script seems to be that the values in PopIn won't change during the time, so that d_Eqs will always be the same...

FYI: I couldn't make this script work. I managed to make a SIR model with agegroups, but in a bit different way with lists
https://github.com/rcsmit/COVIDcases/blob/main/SIR_age_structured_streamlit.py