Visualize whole tree of infections

pyproject.toml

@@ -9,7 +9,6 @@ readme = "README.md"
 python = "^3.10"
 numpy = "^2.2.0"
 streamlit = "^1.41.0"
-graphviz = "^0.20.3"
 polars = "^1.17.1"
 
 

ringvax/__init__.py

@@ -7,7 +7,9 @@
 
 class Simulation:
     PROPERTIES = {
+        "id",
         "infector",
+        "infectees",
         "generation",
         "t_exposed",
         "t_infectious",
@@ -63,6 +65,13 @@ def query_people(self, query: Optional[dict[str, Any]] = None) -> List[str]:
                 if all(person[k] == v for k, v in query.items())
             ]
 
+    def register_infectee(self, infector, infectee) -> None:
+        infectees = self.get_person_property(infector, "infectees")
+        if infectees is None:
+            self.update_person(infector, {"infectees": []})
+            infectees = self.get_person_property(infector, "infectees")
+        infectees.append(infectee)
+
     def run(self) -> None:
         """Run simulation"""
         # queue is pairs (t_exposed, infector)
@@ -133,8 +142,11 @@ def generate_infection(
             generation = 0
         else:
             generation = self.get_person_property(infector, "generation") + 1
+            self.register_infectee(infector, id)
 
-        self.update_person(id, {"infector": infector, "generation": generation})
+        self.update_person(
+            id, {"id": id, "infector": infector, "generation": generation}
+        )
 
         # disease state history in this individual
         disease_history = self.generate_disease_history(t_exposed=t_exposed)
@@ -144,6 +156,7 @@ def generate_infection(
         detection_history = self.generate_detection_history(id)
         self.update_person(id, detection_history)
 
+        t_start_infectious = disease_history["t_infectious"]
         if detection_history["detected"]:
             t_end_infectious = detection_history["t_detected"]
         else:
@@ -161,6 +174,11 @@ def generate_infection(
                 infectious_duration=(
                     t_end_infectious - disease_history["t_infectious"]
                 ),
+                t_infectious=t_start_infectious,
+            )
+            assert all(
+                (it >= t_start_infectious) and (it <= t_end_infectious)
+                for it in infection_times
             )
 
         self.update_person(id, {"infection_times": infection_times})
@@ -243,7 +261,10 @@ def generate_active_detection_delay(self) -> float:
 
     @staticmethod
     def generate_infection_times(
-        rng: numpy.random.Generator, rate: float, infectious_duration: float
+        rng: numpy.random.Generator,
+        rate: float,
+        infectious_duration: float,
+        t_infectious: float,
     ) -> np.ndarray:
         """Times from onset of infectiousness to each infection"""
         assert rate >= 0.0
@@ -255,7 +276,7 @@ def generate_infection_times(
         n_events = rng.poisson(infectious_duration * rate)
 
         # We sort these elsewhere, no need to do extra work
-        return rng.uniform(0.0, infectious_duration, n_events)
+        return t_infectious + rng.uniform(0.0, infectious_duration, n_events)
 
     def bernoulli(self, p: float) -> bool:
         return self.rng.binomial(n=1, p=p) == 1

ringvax/app.py

@@ -2,12 +2,12 @@
 from typing import List
 
 import altair as alt
-import graphviz
 import numpy.random
 import polars as pl
 import streamlit as st
 
 from ringvax import Simulation
+from ringvax.plot import plot_simulation
 from ringvax.summary import (
     get_all_person_properties,
     get_total_infection_count_df,
@@ -17,24 +17,6 @@
 )
 
 
-def make_graph(sim: Simulation) -> graphviz.Digraph:
-    """Make a transmission graph"""
-    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
-
-
 @st.fragment
 def show_graph(sims: List[Simulation], pause: float = 0.1):
     """Show a transmission graph. Wrap as st.fragment, to not re-run simulations.
@@ -49,7 +31,7 @@ def show_graph(sims: List[Simulation], pause: float = 0.1):
     )
     placeholder = st.empty()
     time.sleep(pause)
-    placeholder.graphviz_chart(make_graph(sims[idx]))
+    placeholder.pyplot(plot_simulation(sims[idx]))
 
 
 def format_control_gens(gen: int):

ringvax/plot.py

@@ -0,0 +1,225 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ringvax import Simulation
+
+stage_map = {
+    "latent": {"start": "t_exposed", "end": "t_infectious"},
+    "infectious": {"start": "t_infectious", "end": "t_recovered"},
+}
+
+
+def diamond(xc, yc, width, height, ppl):
+    half_x_l = np.linspace(xc - width / 2.0, xc, ppl).tolist()
+    half_x_r = np.linspace(xc, xc + width / 2.0, ppl).tolist()
+    x = np.array(
+        [*half_x_l[:-1], *half_x_r[:-1], *reversed(half_x_r[1:]), *reversed(half_x_l)]
+    )  # type: ignore
+
+    half_up_y = np.linspace(yc, yc + height / 2.0, ppl).tolist()
+    half_down_y = np.linspace(yc - height / 2.0, yc, ppl).tolist()
+    y = np.array(
+        [
+            *half_up_y[:-1],
+            *reversed(half_up_y[1:]),
+            *reversed(half_down_y[1:]),
+            *half_down_y,
+        ]
+    )  # type: ignore
+
+    return (
+        x,
+        y,
+    )
+
+
+def get_end(id: str, sim: Simulation, stage: str, stage_map: dict[str, dict[str, str]]):
+    """
+    Get time of end of given stage for given infection.
+
+    Returns t_detected if this is less than natural end time of the stage, even if the stage did not start.
+    """
+    if sim.get_person_property(id, "detected") and sim.get_person_property(
+        id, "t_detected"
+    ) < sim.get_person_property(id, stage_map[stage]["end"]):
+        return sim.get_person_property(id, "t_detected")
+    else:
+        return sim.get_person_property(id, stage_map[stage]["end"])
+
+
+def mark_detection(ax, id, sim: Simulation, plot_par):
+    """
+    If this infection was detected, mark that.
+    """
+    if sim.get_person_property(id, "detected"):
+        y_loc = plot_par["height"][id]
+        x_loc = sim.get_person_property(id, "t_detected")
+        x_adj = (
+            0.5
+            * (max(plot_par["x_range"]) - min(plot_par["x_range"]))
+            / (max(plot_par["y_range"]) - min(plot_par["y_range"]))
+        )
+        x, y = diamond(
+            x_loc,
+            y_loc,
+            plot_par["history_thickness"] * x_adj,
+            plot_par["history_thickness"],
+            plot_par["ppl"],
+        )
+        ax.fill(
+            x,
+            y,
+            color=plot_par["color"]["detection"][
+                sim.get_person_property(id, "detect_method")
+            ],
+        )
+
+
+def mark_infections(ax, id, sim: Simulation, plot_par):
+    """
+    Put down tick marks at every time a new infection arises caused by given infection.
+    """
+    y_loc = plot_par["height"][id]
+    if (sim.get_person_property(id, "infection_times")).size > 0:
+        for t in sim.get_person_property(id, "infection_times"):
+            y = np.linspace(
+                y_loc - plot_par["history_thickness"] / 2.0,
+                y_loc + plot_par["history_thickness"] / 2.0,
+                plot_par["ppl"],
+            )
+            x = np.array([t] * len(y))
+            ax.plot(x, y, color=plot_par["color"]["infection"])
+
+
+def draw_stages(ax, id, sim: Simulation, plot_par):
+    """
+    Draw the stages (latent, infectious) of this infection
+    """
+    y_loc = plot_par["height"][id]
+    for stage in ["latent", "infectious"]:
+        x = np.linspace(
+            sim.get_person_property(id, stage_map[stage]["start"]),
+            get_end(id, sim, stage, stage_map),
+            plot_par["ppl"],
+        )
+
+        y = np.array([y_loc] * len(x))
+
+        ax.fill_between(
+            x,
+            y - plot_par["history_thickness"] / 2.0,
+            y + plot_par["history_thickness"] / 2.0,
+            color=plot_par["color"][stage],
+        )
+
+
+def connect_child_infections(ax, id, sim: Simulation, plot_par):
+    """
+    Connect this infection to its children
+    """
+    y_parent = plot_par["height"][id]
+    times_infections = get_infection_time_tuples(id, sim)
+    if times_infections is not None:
+        for t, inf in times_infections:
+            assert (
+                sim.infections[inf]["t_exposed"] == t
+            ), f"Child {inf} reports infection at time {sim.infections[inf]['t_exposed']} while parent reports time was {t}"
+            y_child = plot_par["height"][inf]
+            y = np.linspace(
+                y_child - plot_par["history_thickness"] / 2.0,
+                y_parent - plot_par["history_thickness"] / 2.0,
+                plot_par["ppl"],
+            )
+            x = np.array([t] * len(y))
+            ax.plot(x, y, color=plot_par["color"]["connection"])
+
+
+def get_infection_time_tuples(id: str, sim: Simulation):
+    """
+    Get tuple of (time, id) for all infections this infection causes.
+    """
+    infectees = sim.get_person_property(id, "infectees")
+    if infectees is None or len(infectees) == 0:
+        return None
+
+    return [(sim.get_person_property(inf, "t_exposed"), inf) for inf in infectees]
+
+
+def order_descendants(sim: Simulation):
+    """
+    Get infections in order for plotting such that the tree has no crossing lines.
+
+    We order such that, allowing space for all descendants thereof, the most recent
+    infection caused by any infection is closest to it.
+    """
+    order = ["0"]
+    _order_descendants("0", sim, order)
+    return order
+
+
+def _order_descendants(id: str, sim: Simulation, order: list[str]):
+    """
+    Add this infections descendants in order
+    """
+    assert id in order, f"Cannot append infection {id} to ordering list {order}"
+    times_infections = get_infection_time_tuples(id, sim)
+    if times_infections is not None:
+        for _, inf in times_infections:
+            order.insert(order.index(id) + 1, inf)
+            _order_descendants(inf, sim, order)
+
+
+def make_plot_par(sim: Simulation):
+    """
+    Get parameters for plotting this simulation
+    """
+    plot_order = order_descendants(sim)
+
+    return {
+        "color": {
+            "latent": "#888888",
+            "infectious": "#c7dcdd",
+            "infection": "#888888",
+            "connection": "#888888",
+            "detection": {
+                "active": "#f78f47",
+                "passive": "#068482",
+            },
+        },
+        "history_thickness": 0.5,
+        "linewidth": {
+            "connection": 2.0,
+            "detection": 10.0,
+            "infection": 10.0,
+        },
+        "ppl": 100,
+        "height": {
+            inf: len(plot_order) - 1.0 * height for height, inf in enumerate(plot_order)
+        },
+        "x_range": [
+            0.0,
+            max(
+                get_end(id, sim, "infectious", stage_map)
+                for id in sim.infections.keys()
+            ),
+        ],
+        "y_range": [0.0, len(sim.infections)],
+    }
+
+
+def plot_simulation(sim: Simulation):
+    plot_par = make_plot_par(sim)
+
+    fig, ax = plt.subplots()
+
+    for inf in sim.query_people():
+        draw_stages(ax, inf, sim, plot_par)
+
+        mark_detection(ax, inf, sim, plot_par)
+
+        mark_infections(ax, inf, sim, plot_par)
+
+        connect_child_infections(ax, inf, sim, plot_par)
+
+    ax.set_axis_off()
+    return fig

ringvax/summary.py

@@ -7,7 +7,9 @@
 
 infection_schema = pl.Schema(
     {
+        "id": pl.String,
         "infector": pl.String,
+        "infectees": pl.List(pl.String),
         "generation": pl.Int64,
         "t_exposed": pl.Float64,
         "t_infectious": pl.Float64,