add simple graph

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]

ringvax/app.py

@@ -1,8 +1,23 @@
+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")
+
+        if infector is None:
+            # this is the index infection
+            graph.node(str(infectee))
+        else:
+            graph.edge(str(infector), str(infectee))
+
+    return graph
+
+
 def app():
     st.title("Ring vaccination")
 
@@ -78,6 +93,11 @@ 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")
     st.text(s.infections)