Added variable cured rates

Author: khvilaboa

epidemic_zhong_et_al_v2/generator/model_generator.py

@@ -8,8 +8,10 @@
 parser.add_argument('-t', '--template_ma', type=str, default="template_sirs_zhong_ma", help='Path to zhong .ma template')
 parser.add_argument('-m', '--template_macros', type=str, default="template_sirs_zhong_macros", help='Path to zhong macros template')
 parser.add_argument('-o', '--out_ma_file', type=str, default="out.ma", help='Output path of the .ma file')
-parser.add_argument('-p', '--out_macros_file', type=str, default="sirs_zhong_macros.inc", help='Output path of the macros file')
-parser.add_argument('-v', '--virulences', type=str, default="0.6", help='Default virulence values per state')
+parser.add_argument('-om', '--out_macros_file', type=str, default="sirs_zhong_macros.inc", help='Output path of the macros file')
+parser.add_argument('-ov', '--out_var_file', type=str, default="sirs_zhong.var", help='Output path of the .var file')
+parser.add_argument('-v', '--virulences', type=str, default="0.6", help='Default virulence values per infected state')
+parser.add_argument('-c', '--cured_rates', type=str, default="0.0", help='Default cured rates values per infected state')
 
 args = parser.parse_args()
 
@@ -33,15 +35,26 @@
 with open(args.template_ma, "r") as f:
     content = f.read()
 
-internal_vars = ["i_%s" % s for s in states] + ["v_%d" % n for n in range(1, args.infected_steps + 1)]
+internal_vars = ["i_%s" % s for s in states] + \
+                ["v_%d" % n for n in range(1, args.infected_steps + 1)] + \
+                ["cr_%d" % n for n in range(1, args.infected_steps)]
 content = content.replace("[[internal_vars]]", " ".join(internal_vars))
 
 internal_vars_values = " ".join(steps_per_state + ["1.0"] + ["0.0"] * (len(states) - 1))
 if "," in args.virulences:
     virulences = list(args.virulences.split(","))
 else:
     virulences = [args.virulences] * args.infected_steps
+
 internal_vars_values += " %s" % " ".join(virulences)
+
+if "," in args.cured_rates:
+    cured_rates = list(args.cured_rates.split(","))
+else:
+    cured_rates = [args.cured_rates] * (args.infected_steps - 1)
+
+internal_vars_values += " %s" % " ".join(cured_rates)
+
 content = content.replace("[[internal_vars_values]]", internal_vars_values)
 update_ports_str = " ".join(["~%s := $i_%s;" % (s, s) for s in states])
 content = content.replace("[[initial_ports_set_up]]", update_ports_str)
@@ -55,7 +68,8 @@
 update_rules += "$i_%s := $i_%s + #macro(local_cured);\n\n" % (states[idx_first_rec], states[idx_first_rec-1])
 
 for i in range(idx_first_rec-1, 1, -1):
-    update_rules += "$i_%s := round(min((1 - $cured_rate) * $i_%s, 1)*100)/100;\n" % (states[i], states[i - 1])
+    cr = "$cr_%s" % states[i - 1].split("_")[-1]
+    update_rules += "$i_%s := round(min((1 - %s) * $i_%s, 1)*100)/100;\n" % (states[i], cr, states[i - 1])
 update_rules += "$i_%s := #macro(internal_infected) + #macro(external_infected);\n\n" % states[1]
 
 update_rules += "$i_%s := 1 " % states[0]
@@ -77,7 +91,8 @@
 
 local_cured_rules = []
 for i in range(idx_first_rec-2, 0, -1):
-    local_cured_rules.append("round(min($cured_rate * $i_%s, 1)*100)" % states[i])
+    cr = "$cr_%s" % states[i].split("_")[-1]
+    local_cured_rules.append("round(min(%s * $i_%s, 1)*100)" % (cr, states[i]))
 
 content = content.replace("[[local_cured]]", "(( %s ) / 100 )" % " +\n".join(local_cured_rules))
 content = content.replace("[[infected_vars]]", " + ".join(["i_%s" % s for s in states[idx_first_rec:0]]))
@@ -118,3 +133,16 @@
 
 with open(args.out_macros_file, "w") as f:
     f.write(content)
+
+
+#################################
+# GENERATION OF THE .VAR FILE
+#################################
+
+with open(args.out_var_file, "w") as f:
+    ivv = internal_vars_values.split()
+    assert ivv[2] == "1.0" or ivv[2] == "1", "Unexpected vals (original susceptible: %s)" % ivv[2]
+    ivv[2] = "0.7"
+    ivv[3] = "0.3"
+    stimulus = "(25,25) = 100 1 0.6 %s\n" % " ".join(ivv)
+    f.write(stimulus)

epidemic_zhong_et_al_v2/generator/template_sirs_zhong_ma

@@ -16,8 +16,8 @@ initialvalue : -1
 
 localtransition : sirs-zhong-rule
 
-statevariables: population connection movement cured_rate t_i t_r [[internal_vars]]
-statevalues: 100 1 0.6 0.0 [[internal_vars_values]]
+statevariables: population connection movement t_i t_r [[internal_vars]]
+statevalues: 100 1 0.6 [[internal_vars_values]]
 initialvariablesvalue: sirs_zhong.var
 
 neighborports : initial pop [[ports]]