diff --git a/Analysis/prep_model.py b/Analysis/prep_model.py
index aefac3d..7990f20 100644
--- a/Analysis/prep_model.py
+++ b/Analysis/prep_model.py
@@ -24,7 +24,8 @@ node_pos = {} # used for drawing/graphing
 
 # Compartments
 S = 0
-I = 0
+I_R = 0
+I_D = 0
 E = 0
 R = 0
 D = 0
@@ -68,30 +69,33 @@ for node in A:
     else:
         color_map.append(color)
         color_d[node] = color
-        I = I + 1
         # Check if node dies
         out_edges = list(G.out_edges(node))
         if not out_edges:
             D = D + 1
+            I_D = I_D + 1
+        else:
+            I_R = I_R + 1
         # Check if imported
         in_edges = list(G.in_edges(node))
         if not in_edges:
             ep_tmp = ep_tmp + 1
 
 # Params
-beta = I/len(A) # rate of infec (I/total?)
+beta = (I_R+I_D)/len(A) # rate of infec (I/total?)
 delta = E/len(A) # symptom appearance rate (E/total?)
 gamma_r = R/len(A) # recov rate (R/total?)
 gamma_d = D/len(A) # death rate (D/total?)
-mu = D/I # fatality ratio (D/I)
+mu = D/(I_R+I_D) # fatality ratio (D/I)
 epsilon = ep_tmp/len(A) # infected import rate
 omega = 0 # waning immunity rate
 
 
 print("Model Compartments:")
 print("S:", str(S))
-print("I:", str(I))
 print("E:", str(E))
+print("I_R:", str(I_R))
+print("I_D:", str(I_D))
 print("R:", str(R))
 print("D:", str(D))
 print("\n")