可视化修改的SIR mod

Question

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

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)

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

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')