From 640ce10392d6b30cdcd85d946300b67d45f70e3a Mon Sep 17 00:00:00 2001
From: Luis Aleixo <luis.aleixo@cern.ch>
Date: Wed, 22 Sep 2021 11:39:13 +0200
Subject: [PATCH] Probability of infection plots

---
 cara/model_scenarios_paper.py |  4 +--
 cara/plot_output.py           |  5 ++-
 cara/results_paper.py         | 63 ++++++++++++++++++++++++++++++++---
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/cara/model_scenarios_paper.py b/cara/model_scenarios_paper.py
index 4e70e5f8..828576ca 100644
--- a/cara/model_scenarios_paper.py
+++ b/cara/model_scenarios_paper.py
@@ -97,11 +97,11 @@ def exposure_vl(activity: str, expiration: str, mask: str, vl: float):
             ),
             infected=mc.InfectedPopulation(
                 number=1,
-                virus=models.Virus(
+                virus=mc.SARSCoV2(
                     viral_load_in_sputum=10**vl,
                     infectious_dose=infectious_dose_distribution,
                     viable_to_RNA_ratio=viable_to_RNA_ratio_distribution,
-                    transmissibility_factor=1.,
+                    transmissibility_factor=1.0,
                 ),
                 presence=mc.SpecificInterval(((0, 2),)),
                 mask=models.Mask.types[mask],
diff --git a/cara/plot_output.py b/cara/plot_output.py
index b0e82305..6b9fd71f 100644
--- a/cara/plot_output.py
+++ b/cara/plot_output.py
@@ -49,7 +49,10 @@ from dataclasses import dataclass
 #compare_concentration_curves()
 
 ############ Emission Rate Violin plot ############ 
-compare_viruses_vr()
+#compare_viruses_vr()
+
+############ Probability of infection vs Viral load ############ 
+plot_pi_vs_viral_load(activity='Seated', expiration='Talking', mask='No mask')
 
 ############ Used for testing ############
 #exposure_model_from_vl_talking_new_points()
diff --git a/cara/results_paper.py b/cara/results_paper.py
index e7e20aef..52e800cc 100644
--- a/cara/results_paper.py
+++ b/cara/results_paper.py
@@ -58,7 +58,6 @@ def exposure_model_from_vl(activity, expiration, mask):
 
     er_means = []
     er_means_1h = []
-    er_medians = []
     lower_percentiles = []
     upper_percentiles = []
 
@@ -76,7 +75,6 @@ def exposure_model_from_vl(activity, expiration, mask):
             emission_rate = emission_rate_when_present(exposure_model)
 
         er_means.append(np.mean(emission_rate))
-        er_medians.append(np.median(emission_rate))
         lower_percentiles.append(np.quantile(emission_rate, 0.01))
         upper_percentiles.append(np.quantile(emission_rate, 0.99))
         emission_rate_1h = emission_rate_when_present(exposure_model)
@@ -600,9 +598,66 @@ def compare_viruses_vr():
     plt.tight_layout()
     plt.show()
 
+
+######### Probability of infection vs Viral load #########
+def plot_pi_vs_viral_load(activity, expiration, mask):
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+
+    pi_means = []
+    pi_medians = []
+    lower_percentiles = []
+    upper_percentiles = []
+
+    for vl in tqdm(viral_loads):
+        exposure_mc = exposure_vl(activity, expiration, mask, vl)
+        exposure_model = exposure_mc.build_model(size=SAMPLE_SIZE)
+
+        pi = exposure_model.infection_probability()/100
+        
+        pi_means.append(np.mean(pi))
+        pi_medians.append(np.median(pi))
+        lower_percentiles.append(np.quantile(pi, 0.01))
+        upper_percentiles.append(np.quantile(pi, 0.99))
+
+    ax.plot(viral_loads, pi_means, label='Baseline')
+    ax.fill_between(viral_loads, lower_percentiles,
+                    upper_percentiles, alpha=0.2)
+
+    ############ Plot ############
+    plt.ylabel('Probability of infection', fontsize=14)
+    plt.xticks(ticks=[i for i in range(2, 13)], labels=['$10^{' + str(i) + '}$' for i in range(2, 13)])
+    plt.xlabel('NP viral load, $\mathrm{vl_{in}}$\n(RNA copies)', fontsize=14)
+
+    # add vertical lines for the critical viral loads for which pi= 5 or 95
+    left_index, right_index = 0, 0
+    for i, pi in enumerate(pi_means):
+        if pi > 0.05:
+            left_index = i
+            break
+
+    for i, pi in enumerate(pi_means[::-1]):
+        if pi < 0.95:
+            right_index = len(viral_loads) - i
+            break
+
+    left, right = viral_loads[left_index], viral_loads[right_index]
+    print('Vl_0.05 = 10^', np.round(left, 1), '\n')
+    print('Vl_0.95 = 10^', np.round(right, 1), '\n')
+
+    plt.vlines(x=(left, right), ymin=0, ymax=1,
+              colors=('grey', 'grey'), linestyles='dotted')
+    plt.text(left - 1.1, 0.80, '$vl_{0.05}$', fontsize=14,color='black')
+    plt.text(right + 0.1, 0.80, '$vl_{0.95}$', fontsize=14,color='black')
+    # add 3 shaded areas
+    plt.axvspan(2, left, alpha=0.1, color='limegreen')
+    plt.axvspan(left, right, alpha=0.1, color='orange')
+    plt.axvspan(right, 12, alpha=0.1, color='tomato')
+
+    plt.legend()
+    plt.show()
+
 ######### Auxiliar functions #########
-
-
 def get_enclosure_points(x_coordinates, y_coordinates):
     df = pd.DataFrame({'x': x_coordinates, 'y': y_coordinates})