python - 可视化修改后的 SIR 模型

Question

我尝试从 eon 包中修改 SIR 模型并对其进行了一些更改。它附加了一个新的疫苗接种参数，其中包含新参数 beta、omega 和 VI，我的代码是-

def test_transmission(u, v, p):

return random.random()<p


def discrete_SIR(G,
            initial_infecteds,beta,
            w,Vl,return_full_data=True):

if G.has_node(initial_infecteds):
    initial_infecteds=[initial_infecteds]           

if return_full_data:
    node_history = defaultdict(lambda : ([tmin], ['S']))
    transmissions = []
    for node in initial_infecteds:
        node_history[node] = ([tmin], ['I'])
        transmissions.append((tmin-1, None, node))



node_history = defaultdict(lambda : ([tmin], ['S']))
#    transmissions = []
for node in initial_infecteds:
    node_history[node] = ([tmin], ['I'])
    #transmissions.append((tmin-1, None, node))


N=G.order()
t = [tmin]
S = [N-len(initial_infecteds)]
I = [len(initial_infecteds)]
R = [0]
V = [0]

susceptible = defaultdict(lambda: True)  
#above line is equivalent to u.susceptible=True for all nodes.

for u in initial_infecteds:
    susceptible[u] = False


infecteds = set(initial_infecteds)

while infecteds and t[-1]<tmax :
    new_infecteds = set()
    vaccinated= set()
    infector = {}  #used for returning full data.  a waste of time otherwise
    for u in infecteds:
     #   print('u-->' +str(u))
        for v in G.neighbors(u):
     #       print('v --> '+ str(v))
        ##vaccination
            if len(vaccinated)+V[-1]< (Vl*N)  : #check if vaccination over or not
                #print(len(vaccinated),Vl*N)
                #print("HI")

                if susceptible[v] and test_transmission(u, v, w): 
                    vaccinated.add(v)
                    susceptible[v] = False
     #               print('transmitting vaccination')

                elif susceptible[v] and test_transmission(u,v,beta):
                    new_infecteds.add(v)
                    susceptible[v]=False
                    infector[v] = [u]
     #               print('transmitting infection')
            else:

    #            print("BYE")
                if susceptible[v] and test_transmission(u, v,beta): 
                    new_infecteds.add(v)
                    susceptible[v] = False
                    infector[v] = [u]

           #infector[v] = [u]


    if return_full_data:
        for v in infector.keys():
            transmissions.append((t[-1], random.choice(infector[v]), v))
        next_time = t[-1]+1
        if next_time <= tmax:
            for u in infecteds:
                node_history[u][0].append(next_time)
                node_history[u][1].append('R')
            for v in new_infecteds:
                node_history[v][0].append(next_time)
                node_history[v][1].append('I')



    infecteds = new_infecteds

    R.append(R[-1]+I[-1])
    V.append(len(vaccinated)+V[-1])
    I.append(len(infecteds))
    S.append(N-V[-1]-I[-1]-R[-1])
    #S.append(S[-1]-V[-1]-I[-1])
    t.append(t[-1]+1)


print(str(R[-1])+','+str(V[-1])+','+str(I[-1])+','+str(S[-1]))  

if not return_full_data:
    return scipy.array(t), scipy.array(S), scipy.array(I), \
           scipy.array(R)
else:
    return EoN.Simulation_Investigation(G, node_history, transmissions)

现在我想在它上面运行可视化，就像在 EON 包中一样-

m=5

G=nx.grid_2d_graph(m,m,periodic=True)
initial_infections = [(u,v) for (u,v) in G if u==int(m/2) and v==int(m/2)]
sim = EoN.basic_discrete_SIR(G,0.5,initial_infecteds = initial_infections,
               return_full_data=True, tmax = 25)

pos = {node:node for node in G}
sim.set_pos(pos)
sim.display(0, node_size = 40) #display time 6
plt.show()
plt.savefig('SIR_2dgrid.png')

我需要在代码中进行哪些更改才能使显示功能正常工作，或者我还需要对显示功能进行更改吗？

Answer 1

这是我现在得到的输出:

您必须安装 EoN 版本 1.0.8rc3 或更高版本，该版本可在 github 页面上找到(请参阅 installation instructions )。目前 pip 无法安装它。我想确保在将其设置为 pip 默认安装之前没有破坏任何内容。

这是基于您的代码。你应该看看我所做的改变。也值得一看the examples I've put in the documentation (包括 SIRV 模型，其中疫苗接种规则与您所拥有的不同)。

from collections import defaultdict
import EoN
import networkx as nx
import random
import matplotlib.pyplot as plt

def test_transmission(u, v, p):
    return random.random()<p


def discrete_SIRV(G, initial_infecteds,beta,
            w,Vl,tmin=0,tmax=float('Inf'), return_full_data=True):

    if G.has_node(initial_infecteds):
        initial_infecteds=[initial_infecteds]           

    if return_full_data:
        node_history = defaultdict(lambda : ([tmin], ['S']))
        transmissions = []
        for node in initial_infecteds:
            node_history[node] = ([tmin], ['I'])
            transmissions.append((tmin-1, None, node))



    '''
    node_history = defaultdict(lambda : ([tmin], ['S']))
    #    transmissions = []
    for node in initial_infecteds:
        node_history[node] = ([tmin], ['I'])
    #transmissions.append((tmin-1, None, node))
    '''

    N=G.order()
    t = [tmin]
    S = [N-len(initial_infecteds)]
    I = [len(initial_infecteds)]
    R = [0]
    V = [0]

    susceptible = defaultdict(lambda: True)  
    #above line is equivalent to u.susceptible=True for all nodes.

    for u in initial_infecteds:
        susceptible[u] = False


    infecteds = set(initial_infecteds)

    while infecteds and t[-1]<tmax :
        new_infecteds = set()
        vaccinated= set()
        infector = {}  #used for returning full data.  a waste of time otherwise
        for u in infecteds:
            #   print('u-->' +str(u))
            for v in G.neighbors(u):
                #       print('v --> '+ str(v))
                ##vaccination
                if len(vaccinated)+V[-1]< (Vl*N)  : #check if vaccination over or not
                    #print(len(vaccinated),Vl*N)
                    #print("HI")

                    if susceptible[v] and test_transmission(u, v, w): 
                        vaccinated.add(v)
                        susceptible[v] = False

                        '''It's probably better to define a `new_vaccinated`
                        set and then do the `return_full_data` stuff later
                        where all the others are done.'''
                        if return_full_data:
                            node_history[v][0].append(t[-1]+1)
                            node_history[v][1].append('V')
     #               print('transmitting vaccination')

                    elif susceptible[v] and test_transmission(u,v,beta):
                        new_infecteds.add(v)
                        susceptible[v]=False
                        infector[v] = [u]
                        #               print('transmitting infection')
                else:

                    #            print("BYE")
                    if susceptible[v] and test_transmission(u, v,beta): 
                        new_infecteds.add(v)
                        susceptible[v] = False
                        infector[v] = [u]

                        #infector[v] = [u]


        if return_full_data:
            for v in infector.keys():
                '''This random choice is no longer needed as you've taken out 
                the possibility of multiple nodes transmitting to `v` in a given
                time step.  Now only the first one encountered does it.'''

                transmissions.append((t[-1], random.choice(infector[v]), v))
            next_time = t[-1]+1
            if next_time <= tmax:
                for u in infecteds:
                    node_history[u][0].append(next_time)
                    node_history[u][1].append('R')
                for v in new_infecteds:
                    node_history[v][0].append(next_time)
                    node_history[v][1].append('I')



        infecteds = new_infecteds

        R.append(R[-1]+I[-1])
        V.append(len(vaccinated)+V[-1])
        I.append(len(infecteds))
        S.append(N-V[-1]-I[-1]-R[-1])
        #S.append(S[-1]-V[-1]-I[-1])
        t.append(t[-1]+1)


        print(str(R[-1])+','+str(V[-1])+','+str(I[-1])+','+str(S[-1]))  

    if not return_full_data:
        return scipy.array(t), scipy.array(S), scipy.array(I), \
           scipy.array(R)
    else:
        return EoN.Simulation_Investigation(G, node_history, transmissions, possible_statuses=['S', 'I', 'R', 'V'])



print(EoN.__version__)
print("line above needs to be 1.0.8rc3 or greater or it will not work\n\n")

m=20

G=nx.grid_2d_graph(m,m,periodic=True)
initial_infections = [(u,v) for (u,v) in G if u==int(m/2) and v==int(m/2)]
beta=0.8
Vl=0.3
w=0.1

sim = discrete_SIRV(G, initial_infections, beta, w, Vl, return_full_data=True)

pos = {node:node for node in G}
sim.set_pos(pos)
sim.sim_update_colordict({'S': '#009a80','I':'#ff2000', 'R':'gray','V': '#5AB3E6'})
sim.display(6, node_size = 40) #display time 6
plt.savefig('SIRV_2dgrid.png')