SIRモデル作成中の KeyError の解決

前提

Pythonで複数のウイルスかつネットワーク化されたSIRモデルを作成しています．
一つのウイルスでの作成はうまくいったのですが，三つに拡張した場合の処理が上手くいきません

実現したいこと

エラーの解決
モデルの完成

発生している問題・エラーメッセージ

一つ目のウイルスでは正しく設定できているのですが，二つ目三つ目以降が上手く行きません
キーの設定に問題が発生しているようなのですがどこに問題があるのか分かりません．

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_15596/301256213.py in <module>
      1 if __name__ == '__main__':
----> 2     sir = SirModel(81306)
      3     sir.spread(gamma1=0.05, gamma2=0.04, gamma3=0.03, beta1=0.6, beta2=0.5, beta3=0.4)
      4     sir.show_SIR_detail()

~\AppData\Local\Temp/ipykernel_15596/204124025.py in __init__(self, N, m)
     31                        "Susceptible": '#808080'}
     32         self.init_source1(1)
---> 33         self.init_source2(1)
     34         self.init_source3(1)
     35 

~\AppData\Local\Temp/ipykernel_15596/204124025.py in init_source2(self, n, S_rate)
     65         for i in range(n):
     66             node = degree_cent[i][1]
---> 67             self.network.nodes[node]['state'] = State.I2
     68             self.infective3_nodes.append(node)
     69         self.update_SIR_nums(int(self.N * S_rate), 0, n, 0, 0)

~\anaconda3\lib\site-packages\networkx\classes\reportviews.py in __getitem__(self, n)
    191                 f"try list(G.nodes)[{n.start}:{n.stop}:{n.step}]"
    192             )
--> 193         return self._nodes[n]
    194 
    195     # Set methods

KeyError: 0.02816554947420208

該当のソースコード①

ウイルスが三つの場合

Python
1from enum import Enum
2import random
3import matplotlib.pyplot as plt
4import networkx as nx
5import numpy as np
6
7
8class State(Enum):
9    I1 = 'Infective1'
10    I2 = 'Infective2'
11    I3 = 'Infective3'
12    S = 'Susceptible'
13    R = 'Removal'
14
15
16
17class SirModel:
18
19    def __init__(self, N, m=2):
20        self.iter = 0
21        self.N = N
22        self.network = nx.barabasi_albert_graph(N, m)
23        self.infective1_nodes = []
24        self.infective2_nodes = []
25        self.infective3_nodes = []
26        self.removal_nodes = []
27
28        self.S_list = []
29        self.I1_list = []
30        self.I2_list = []
31        self.I3_list = []
32        self.R_list = []
33        self.pos = None
34        self.colors = {"Removal": '#228B22', 
35                       "Infective1": '#B22222', 
36                       "Infective2": '#AC83f3', 
37                       "Infective3": '#DDAA22',  
38                       "Susceptible": '#808080'}
39        self.init_source1(1)
40        self.init_source2(1)
41        self.init_source3(1)
42
43
44    def update_SIR_nums(self, s, i1, i2, i3, r):
45        self.S_list.append(s)
46        self.I1_list.append(i1)
47        self.I2_list.append(i2)
48        self.I3_list.append(i3)
49        self.R_list.append(r)
50        
51        
52    """初期感染者の決定"""
53    def init_source1(self, n, S_rate=1):
54        for i in range(int(self.N * S_rate)):
55            self.network.nodes[i]['state'] = State.S
56
57        degree_cent = nx.algorithms.centrality.degree_centrality(self.network)
58        degree_cent = sorted(degree_cent.items(), key=lambda x: x[1], reverse=True)
59        for i in range(n):
60            node = degree_cent[i][0]
61            self.network.nodes[node]['state'] = State.I1
62            self.infective1_nodes.append(node)
63        self.update_SIR_nums(int(self.N * S_rate), n, 0, 0, 0)
64        self.pos = nx.random_layout(self.network)
65        
66    def init_source2(self, n, S_rate=1):
67        for i in range(int(self.N * S_rate)):
68            self.network.nodes[i]['state'] = State.S
69
70        degree_cent = nx.algorithms.centrality.degree_centrality(self.network)
71        degree_cent = sorted(degree_cent.items(), key=lambda x: x[1], reverse=True)
72        for i in range(n):
73            node = degree_cent[i][1]
74            self.network.nodes[node]['state'] = State.I2
75            self.infective3_nodes.append(node)
76        self.update_SIR_nums(int(self.N * S_rate), 0, n, 0, 0)
77        self.pos = nx.random_layout(self.network)
78        
79    def init_source3(self, n, S_rate=1):
80        for i in range(int(self.N * S_rate)):
81            self.network.nodes[i]['state'] = State.S
82
83        degree_cent = nx.algorithms.centrality.degree_centrality(self.network)
84        degree_cent = sorted(degree_cent.items(), key=lambda x: x[1], reverse=True)
85        for i in range(n):
86            node = degree_cent[i][2]
87            self.network.nodes[node]['state'] = State.I3
88            self.infective3_nodes.append(node)
89        self.update_SIR_nums(int(self.N * S_rate), 0, 0, n, 0)
90        self.pos = nx.random_layout(self.network)
91    
92    
93    """ダイナミクスの計算"""
94    def spread(self, gamma1=0.05, gamma2=0.04, gamma3=0.03, beta1=0.6, beta2=0.5, beta3=0.4, max_iter=100):
95        cur_iter = 0
96        img_data_list = []
97        while len(self.infective1_nodes) > 0 or len(self.infective2_nodes) > 0 or len(self.infective3_nodes) > 0:
98            cur_iter += 1
99            if max_iter and cur_iter > max_iter:
100                break
101
102            new_S_num = self.S_list[-1]
103            new_I1_num = self.I1_list[-1]
104            new_I2_num = self.I2_list[-1]
105            new_I3_num = self.I3_list[-1]
106            new_R_num = self.R_list[-1]
107            
108            for i in self.infective1_nodes:
109                if random.random() < gamma1 and self.network.nodes[i]['state'] == State.I1:
110                    self.removal_nodes.append(i)
111                    self.infective1_nodes.remove(i)
112
113                    self.network.nodes[i]['state'] = State.R
114                    new_R_num += 1
115                    new_I1_num -= 1
116
117            for i in self.infective2_nodes:   
118                if random.random() < gamma2 and self.network.nodes[i]['state'] == State.I2:
119                    self.removal_nodes.append(i)
120                    self.infective2_nodes.remove(i)
121
122                    self.network.nodes[i]['state'] = State.R
123                    new_R_num += 1
124                    new_I2_num -= 1
125                    
126            for i in self.infective3_nodes:    
127                if random.random() < gamma3 and self.network.nodes[i]['state'] == State.I3:
128                    self.removal_nodes.append(i)
129                    self.infective3_nodes.remove(i)
130
131                    self.network.nodes[i]['state'] = State.R
132                    new_R_num += 1
133                    new_I3_num -= 1
134
135            tmp_infective1_nodes = nx.Graph()
136            tmp_infective2_nodes = nx.Graph()
137            tmp_infective3_nodes = nx.Graph()
138            tmp_infective1_nodes.add_nodes_from(self.infective1_nodes)
139            tmp_infective2_nodes.add_nodes_from(self.infective2_nodes)
140            tmp_infective3_nodes.add_nodes_from(self.infective3_nodes)
141            new_infective1_nodes = []
142            new_infective2_nodes = []
143            new_infective3_nodes = []
144            
145            for node in tmp_infective1_nodes.nodes:
146                for neighbor in self.network.neighbors(node):
147                    if random.random() < beta1 and self.network.nodes[neighbor]['state'] == State.S:
148                        self.network.nodes[neighbor]['state'] = State.I1
149                        new_I1_num += 1
150                        new_S_num -= 1
151                        new_infective1_nodes.append(neighbor)
152            for node in new_infective1_nodes:
153                if node not in self.infective_nodes:
154                    self.infective1_nodes.append(node)
155            
156            for node in tmp_infective2_nodes.nodes:
157                for neighbor in self.network.neighbors(node):
158                    if random.random() < beta2 and self.network.nodes[neighbor]['state'] == State.S:
159                        self.network.nodes[neighbor]['state'] = State.I2
160                        new_I2_num += 1
161                        new_S_num -= 1
162                        new_infective2_nodes.append(neighbor)
163            for node in new_infective2_nodes:
164                if node not in self.infective2_nodes:
165                    self.infective2_nodes.append(node)
166            
167            for node in tmp_infective3_nodes.nodes:
168                for neighbor in self.network.neighbors(node):
169                    if random.random() < beta3 and self.network.nodes[neighbor]['state'] == State.S:
170                        self.network.nodes[neighbor]['state'] = State.I3
171                        new_I3_num += 1
172                        new_S_num -= 1
173                        new_infective3_nodes.append(neighbor)
174            for node in new_infective3_nodes:
175                if node not in self.infective3_nodes:
176                    self.infective3_nodes.append(node)
177            
178            
179            self.update_SIR_nums(new_S_num, new_I1_num, new_I2_num, new_I3_num, new_R_num)
180
181        self.iter = cur_iter
182        
183        print("finish")
184       
185    
186    """結果をフォルダーに保存"""
187    def show_SIR_detail(self):
188        x = list(range(1, len(self.S_list) + 1))
189        plt.plot(x, self.S_list, self.colors[State.S.value], label=State.S.value)
190        plt.plot(x, self.I1_list, self.colors[State.I1.value], label=State.I1.value)
191        plt.plot(x, self.I2_list, self.colors[State.I2.value], label=State.I2.value)
192        plt.plot(x, self.I3_list, self.colors[State.I3.value], label=State.I3.value)
193        plt.plot(x, self.R_list, self.colors[State.R.value], label=State.R.value)
194        plt.legend()
195        plt.xlabel('t')
196        plt.ylabel('count')
197        plt.savefig(f"./results/sir_curve.png",dpi=400)
198        plt.close()
199
200
201if __name__ == '__main__':
202    sir = SirModel(1000)
203    sir.spread(gamma1=0.05, gamma2=0.04, gamma3=0.03, beta1=0.6, beta2=0.5, beta3=0.4)
204    sir.show_SIR_detail()