Graph visualization (#11)

* add simple graph * add node colors * update lockfile * Color detections Co-authored-by: afmagee42 <[email protected]> --------- Co-authored-by: afmagee42 <[email protected]>
CDCgov · Dec 17, 2024 · fbe36b7 · fbe36b7
1 parent 48bb870
commit fbe36b7
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 4 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,6 +9,7 @@ readme = "README.md"
 python = "^3.10"
 numpy = "^2.2.0"
 streamlit = "^1.41.0"
+graphviz = "^0.20.3"
 
 
 [tool.poetry.group.dev.dependencies]

diff --git a/ringvax/app.py b/ringvax/app.py
@@ -1,8 +1,28 @@
+import graphviz
 import streamlit as st
 
 from ringvax import Simulation
 
 
+def make_graph(sim: Simulation):
+    graph = graphviz.Digraph()
+    for infectee in sim.query_people():
+        infector = sim.get_person_property(infectee, "infector")
+
+        color = (
+            "black"
+            if not sim.get_person_property(infectee, "detected")
+            else "#068482"
+        )
+
+        graph.node(str(infectee), color=color)
+
+        if infector is not None:
+            graph.edge(str(infector), str(infectee))
+
+    return graph
+
+
 def app():
     st.title("Ring vaccination")
 
@@ -89,6 +109,10 @@ def app():
     s = Simulation(params=params, seed=seed)
     s.run()
 
+    st.header("Graph of infections")
+    st.graphviz_chart(make_graph(s))
+
+    st.header("Raw results")
     for id, content in s.infections.items():
         st.text(id)
         st.text(content)