Python 改进SIR模型的可视化

我试图修改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)

---

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

以下是我现在得到的输出：

您必须安装github页面上提供的EoN版本

1.0.8rc3

或更高版本（请参阅）。目前，

pip

无法安装它。我想确保我没有损坏任何东西，然后将其设置为默认安装程序

pip

这是基于你的代码。你应该仔细看看我所做的改变。它也值得一看（包括一个SIRV模型，其中疫苗接种规则与您现有的不同）

从集合导入defaultdict
导入EoN
将networkx导入为nx
随机输入
将matplotlib.pyplot作为plt导入
def测试_变速箱（u、v、p）：
return random.random（）目前，您需要修改代码和display函数。我希望很快做修改，这样你就可以不用做重大的改变。我稍后会尝试给出一个答案（目前没有时间）给出更多细节。@Joel我对代码做了一些更改，看起来模拟是正确的，但当然它只显示SIR图而不是SIRV。我会发布我的更改，如果您能告诉我要对显示函数进行哪些更改，以包括将要使用的V参数helpful@Joel你有时间调查吗？。你也可以告诉我需要做什么。我很乐意为EON做贡献。嘿-对不起，我没有机会做。需要做的第一个更改是，当一个节点变成
'V'
时，您需要更新该节点的
节点历史记录。第二个变化是，类
模拟调查（G、节点\u历史、传输）
的
\uuu初始化\uuuuuuuu
代码需要修改，以处理节点是否为
'V'
，处理方式与处理S、I和R的方式相同。最后需要修改
\uplot\ucode>方法。我希望将我编写的代码变得更加通用，这样人们可以在任何新状态下执行此操作，但我暂时没有时间。@Joel将尝试执行此操作，并将其发布到此处或github，供您查看。
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')

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