add node colors

CDCgov · Dec 13, 2024 · 663c286 · 663c286
1 parent a96ad9a
commit 663c286
Showing 1 changed file with 9 additions and 4 deletions.
diff --git 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