add legend and cases w/o cluster estimate

clusterDynamics.py

@@ -327,7 +327,7 @@ def cost(P, t, k, n):
 rates = {}
 n_bootstraps=100
 #for a simpler plot of most interesting countries use this:
-for coun in ['Switzerland', 'England', 'Scotland', 'Wales', 'Spain']:
+for coun in ['Switzerland', 'England', 'Scotland', 'Wales', 'Spain', 'United Kingdom']:
     week_as_date, cluster_count, total_count = non_zero_counts(cluster_data, coun)
     days = np.array([x.toordinal() for x in week_as_date])
     mean_upper_lower = []
@@ -356,9 +356,10 @@ def cost(P, t, k, n):
     rates[coun]['upper'] = scoreatpercentile(bootstraps, 75)
     rates[coun]['t50'] = center_fit['x'][1]
 
-    plt.plot(week_as_date, logistic(days, center_fit['x'][0], center_fit['x'][1]),
-             c=country_styles[coun]['c'], ls=country_styles[coun]['ls'],
-             label = f"{coun}, growth rate: {rates[coun]['center']*700:1.1f}({rates[coun]['lower']*700:1.1f}-{rates[coun]['upper']*700:1.1f})%/week")
+    if coun is not 'United Kingdom':
+        plt.plot(week_as_date, logistic(days, center_fit['x'][0], center_fit['x'][1]),
+                 c=country_styles[coun]['c'], ls=country_styles[coun]['ls'],
+                 label = f"{coun}, growth rate: {rates[coun]['center']*700:1.1f}({rates[coun]['lower']*700:1.1f}-{rates[coun]['upper']*700:1.1f})%/week")
     print(f"{coun} growth rate: {rates[coun]['center']*700:1.2f}% per week")
 
 plt.legend()
@@ -383,7 +384,7 @@ def cost(P, t, k, n):
 seqs_week = {}
 cases_week = {}
 
-for coun in ['Switzerland', 'Norway', 'Spain', 'United Kingdom']:
+for coun in ['Switzerland', 'Spain', 'United Kingdom', 'Norway']:
     #read in case data
     cases = pd.read_csv(case_data_path+case_files[coun], sep='\t', index_col=False, skiprows=3)
 
@@ -431,7 +432,8 @@ def cost(P, t, k, n):
     color='tab:blue'
     ax1.set_ylabel('New Cases', color=color)
     lines.append(ax1.plot(case_week_as_date , case_data[coun], color=color, label='cases per week')[0])
-    # lines.append(ax1.plot(case_week_as_date , case_data[coun]*(1 - logistic(days, rates[coun]['center'], rates[coun]['t50']) ), color=color, label='cases per week w/o cluster')[0])
+    if coun is not 'Norway':
+        lines.append(ax1.plot(case_week_as_date , case_data[coun]*(1 - logistic(days, rates[coun]['center'], rates[coun]['t50']) ), color=color, ls='--', label='cases per week w/o cluster')[0])
     ax1.tick_params(axis='y', labelcolor=color)
     ax1.set_yscale("log")
 
@@ -446,7 +448,10 @@ def cost(P, t, k, n):
     ax2.set_yscale("log")
 
     fig.autofmt_xdate(rotation=30)
-    plt.legend(lines, ['cases per week', 'total sequences', 'sequences in cluster'], loc=3)
+    if coun is 'Norway':
+        plt.legend(lines, ['cases per week', 'total sequences', 'sequences in cluster'], loc=3)
+    else:
+        plt.legend(lines, ['cases per week', 'cases per week w/o cluster', 'total sequences', 'sequences in cluster'], loc=3)
     fig.tight_layout()
     plt.show()
     plt.savefig(figure_path+f"{coun}-newcases-seqs.{fmt}")