Use maths formatting for "R0" in plot axis labels

tests/test_sirv.py

@@ -105,6 +105,13 @@ def verify_prior_samples(contexts):
         assert np.array_equal(baseline_vals, vaccination_vals)
 
 
+def get_scenario_label(scenario_id):
+    """
+    Use maths formatting for "R0" in plot axis labels.
+    """
+    return scenario_id.replace('R0', '$R_0$')
+
+
 def plot_epi_curves(results, png_file):
     """
     Plot the ensemble of daily incidence curves for each scenario.
@@ -122,7 +129,8 @@ def plot_epi_curves(results, png_file):
         inc_curves = daily_inc['value'].reshape((len(time), -1))
 
         ax.plot(time, inc_curves, alpha=0.02, color='black')
-        ax.set_title(f'Scenario: {scenario_id}')
+        scenario_label = get_scenario_label(scenario_id)
+        ax.set_title(f'Scenario: {scenario_label}')
 
         # Add the x-axis label to the bottom subplot.
         if ix == len(results) - 1:
@@ -139,7 +147,6 @@ def plot_final_sizes(results, png_file):
     """
     Plot the median and 95% credible interval for final size in each scenario.
     """
-    scenario_ids = results.keys()
     ys, y_mins, y_maxs = get_final_sizes(results)
     fig, ax = plt.subplots(layout='constrained')
 
@@ -150,7 +157,8 @@ def plot_final_sizes(results, png_file):
 
     # Plot the median and credible interval for each scenario.
     y_errs = np.array([ys - y_mins, y_maxs - ys])
-    ax.errorbar(scenario_ids, ys, yerr=y_errs, fmt='o')
+    labels = [get_scenario_label(key) for key in results.keys()]
+    ax.errorbar(labels, ys, yerr=y_errs, fmt='o')
 
     ax.set_xlabel('Scenario')
     ax.set_ylabel('Total infections')
@@ -235,7 +243,8 @@ def plot_prevented_infections(results, png_file):
 
     # Plot the median and credible interval for each scenario.
     y_errs = np.array([ys - y_mins, y_maxs - ys])
-    ax.errorbar(prevented_infs.keys(), ys, yerr=y_errs, fmt='o')
+    labels = [get_scenario_label(key) for key in prevented_infs.keys()]
+    ax.errorbar(labels, ys, yerr=y_errs, fmt='o')
 
     ax.set_xlabel('Scenario')
     # Expand the x-axis limits to avoid plotting the two sets of results
@@ -286,9 +295,8 @@ def plot_prevented_infection_correlations(contexts, results, png_file):
     ax.set_title('Correlation with infections prevented')
     ax.set_xlabel('Scenario')
     ax.set_ylabel('Parameter')
-    ax.set_xticks(
-        np.arange(len(prevented_infs)), labels=prevented_infs.keys()
-    )
+    labels = [get_scenario_label(key) for key in prevented_infs.keys()]
+    ax.set_xticks(np.arange(len(prevented_infs)), labels=labels)
     ax.set_yticks(
         np.arange(len(samples)),
         labels=['$R_0$', r'$\gamma$', r'$\beta$'],