generations-termination was being mistaken for extinction

CDCgov · Dec 18, 2024 · 117a0ad · 117a0ad
1 parent fe4c2c5
commit 117a0ad
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/ringvax/__init__.py b/ringvax/__init__.py
@@ -61,6 +61,14 @@ def run(self) -> None:
             id = self.create_person()
             self.generate_infection(id=id, t_exposed=t_exposed, infector=infector)
 
+            if (
+                len(infection_queue) == 0
+                and len(self.get_person_property(id, "infection_times")) == 0
+            ):
+                termination["criterion"] = "extinction"
+                termination["last_complete_generation"] = max(
+                    self.infections[id]["generation"] for id in self.infections
+                )
             # add the infections they caused to the end of the queue,
             # unless we are past the number of generations
             if (
@@ -72,12 +80,6 @@ def run(self) -> None:
                 for t in self.get_person_property(id, "infection_times"):
                     bisect.insort_right(infection_queue, (t, id), key=lambda x: x[0])
 
-            if len(infection_queue) == 0:
-                termination["criterion"] = "extinction"
-                termination["last_complete_generation"] = max(
-                    infection["generation"] for infection in self.infections
-                )
-
             if len(self.query_people()) >= self.params["max_infections"]:
                 termination["max_infections"] = True
                 min_in_progress = min(
@@ -90,6 +92,8 @@ def run(self) -> None:
             termination["criterion"] = "max_generations"
             termination["last_complete_generation"] = self.params["n_generations"]
 
+        self.termination = termination
+
     def generate_infection(
         self, id: str, t_exposed: float, infector: Optional[str]
     ) -> None: