Also generate by_location_prevalence-violin plot

naobservatory · Aug 7, 2023 · dddf4d7 · dddf4d7
1 parent ab18f7e
commit dddf4d7
Showing 1 changed file with 25 additions and 1 deletion.
diff --git a/plot_summaries.py b/plot_summaries.py
@@ -264,8 +264,13 @@ def plot_three_virus(
         # TODO Get these values automatically
         num_spurbeck = 10
         num_rothman = 8
+        num_crits_christoph = 4
         ax.hlines(
-            [num_spurbeck - 0.5, num_spurbeck + num_rothman - 0.5],
+            [
+                num_spurbeck - 0.5,
+                num_spurbeck + num_rothman - 0.5,
+                num_spurbeck + num_rothman + num_crits_christoph - 0.5,
+            ],
             *ax.get_xlim(),
             linestyle="solid",
             color="k",
@@ -286,6 +291,14 @@ def plot_three_virus(
                 "Crits-Christoph",
                 va="top",
             )
+            if predictor_type == "prevalence":
+                ax.text(
+                    x_text,
+                    num_spurbeck + num_rothman + num_crits_christoph - 0.4,
+                    "Brinch",
+                    va="top",
+                )
+
         adjust_axes(ax, predictor_type=predictor_type)
         ax.set_title(pathogen)
         ax.get_legend().remove()
@@ -353,6 +366,17 @@ def start() -> None:
     save_plot(
         fig_three_virus_incidence, figdir, "by_location_incidence-violin"
     )
+    prevalence_viruses = {
+        "JCV": (-14.0, -7.0),
+        "BKV": (-14.0, -7.0),
+        "MCV": (-14.0, -7.0),
+    }
+    fig_three_virus_prevalence = plot_three_virus(
+        fits_df, input_df, prevalence_viruses, "prevalence"
+    )
+    save_plot(
+        fig_three_virus_prevalence, figdir, "by_location_prevalence-violin"
+    )
 
 
 if __name__ == "__main__":