Update app.py

app.py

@@ -78,7 +78,7 @@ def plotdata(t, s, i,r,R0, e=None):
     return fig
  
 #final model
-def SIR(country,R0,t_infective):
+def SIR(country,R0,t_infective,pop):
   #R0 = 0.57 - 1.25
 
   # parameter values
@@ -86,17 +86,12 @@ def SIR(country,R0,t_infective):
   t_infective = t_infective
     
   # initial number of infected and recovered individuals
-  i_initial = all_location[country]['total_cases'].iloc[0]/pop_dict[country]
+  i_initial = all_location[country]['total_cases'].iloc[0]/pop
   r_initial = 0.00
   s_initial = 1 - i_initial - r_initial
     
   gamma = 1/t_infective
   beta = R0*gamma
-
-  # initial number of infected and recovered individuals
-  i_initial = all_location[country]['new_cases'].sum()/pop_dict[country]
-  r_initial = 0.00
-  s_initial = 1 - i_initial - r_initial
   
   t = np.linspace(0, 3000, 3000)
   x_initial = s_initial, i_initial, r_initial
@@ -104,7 +99,7 @@ def SIR(country,R0,t_infective):
   s, i, r = soln.T
   e = None
   
-  scaler = all_location[country]['total_cases'].apply(lambda x : x/pop_dict[country])
+  scaler = all_location[country]['total_cases'].apply(lambda x : x/pop)
   rangee =  len(all_location[country]['total_cases'])
   rmpe = mean_absolute_percentage_error(scaler,i[0:rangee])
 
@@ -148,7 +143,7 @@ def main():
         if clicked:
         
             # Show SIR
-            SIR_param = SIR(country_code,R0,recovery)
+            SIR_param = SIR(country_code,R0,recovery,pop)
             
             st.pyplot(SIR_param[-1])
             st.pyplot(compare_plt(country_code,pop))