diff --git a/ringvax/app.py b/ringvax/app.py index 6822adb..aedf3d8 100644 --- a/ringvax/app.py +++ b/ringvax/app.py @@ -9,10 +9,15 @@ def make_graph(sim: Simulation): for infectee in sim.query_people(): infector = sim.get_person_property(infectee, "infector") - if infector is None: - # this is the index infection - graph.node(str(infectee)) - else: + color = ( + "black" + if sim.get_person_property(infectee, "actually_infected") + else "gray" + ) + + graph.node(str(infectee), color=color) + + if infector is not None: graph.edge(str(infector), str(infectee)) return graph