fix transmission tree error
@@ -1346,6 +1346,8 @@ class TransTree(sc.prettyobj):
ttlist.append(tdict)
df = pd.DataFrame(ttlist).rename(columns={'date': 'Day'})
+
+ if len(df): # Don't proceed if there were no infections
df = df.loc[df['layer'] != 'seed_infection']
df['Stage'] = 'Symptomatic'
@@ -1355,6 +1357,8 @@ class TransTree(sc.prettyobj):
df['Severity'] = 'Mild'
df.loc[df['s_sev'], 'Severity'] = 'Severe'
df.loc[df['s_crit'], 'Severity'] = 'Critical'
+ else:
+ print('Warning: transmission tree is empty since there were no infections.')
self.df = df
