Fix infection channel plot if called with prepared data. (#119)

covid-19-impact-lab · May 10, 2021 · dec9b8b · dec9b8b
1 parent 2ed8725
commit dec9b8b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/sid/plotting.py b/src/sid/plotting.py
@@ -200,7 +200,9 @@ def _is_data_prepared_for_heatmap(df):
     return (
         isinstance(df, pd.DataFrame)
         and df.columns.isin(["date", "channel_infected_by_contact", "share"]).all()
-        and not df["channel_infected_by_contact"].isin("not_infected_by_contact").any()
+        and not df["channel_infected_by_contact"]
+        .isin(["not_infected_by_contact"])
+        .any()
     )