get_counts() takes date_col parameter

globaldothealth · Aug 29, 2024 · 0ab1779 · 0ab1779
1 parent abe7a36
commit 0ab1779
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/obr/outbreaks/__init__.py b/src/obr/outbreaks/__init__.py
@@ -11,7 +11,7 @@
 )
 
 outbreak_marburg = [
-    ("data", get_counts),
+    ("data", get_counts, dict(date_col="Data_up_to")),
     ("figure/epicurve", plot_epicurve),
     (
         "figure/epicurve_location_status",
@@ -41,5 +41,9 @@
     ),
 ]
 
-OUTBREAKS = {"marburg": outbreak_marburg}
+outbreak_mpox_2024 = [
+    ("data", get_counts, dict(date_col="Date_entry")),
+    ("figure/epicurve", plot_epicurve),
+]
+OUTBREAKS = {"marburg": outbreak_marburg, "mpox-2024": outbreak_mpox_2024}
 __all__ = ["OUTBREAKS"]
diff --git a/src/obr/plots.py b/src/obr/plots.py
@@ -4,10 +4,7 @@
 
 import re
 import logging
-from functools import cache
 
-import boto3
-import chevron
 import pandas as pd
 import numpy as np
 from dateutil.parser import ParserError
@@ -16,8 +13,6 @@
 from plotly.subplots import make_subplots
 
 
-pd.options.mode.chained_assignment = None
-
 from .util import percentage_occurrence, name_bin, AGE_BINS, get_age_bins
 from .theme import (
     REGEX_DATE,
@@ -32,6 +27,8 @@
     GRID_COLOR,
 )
 
+pd.options.mode.chained_assignment = None
+
 
 def get_age_bin_data(df: pd.DataFrame) -> pd.DataFrame:
     confirmed = df[df.Case_status == "confirmed"][["Age", "Gender"]]
@@ -101,13 +98,14 @@ def get_epicurve(df: pd.DataFrame, cumulative: bool = True) -> pd.DataFrame:
     return epicurve.reset_index()
 
 
-def get_counts(df: pd.DataFrame) -> dict[str, int]:
+def get_counts(df: pd.DataFrame, date_col: str) -> dict[str, int]:
     status = df.Case_status.value_counts()
     confirmed = df[df.Case_status == "confirmed"]
     return {
         "n_confirmed": int(status.confirmed),
         "n_probable": int(status.get("probable", 0)),
-        "date": str(df[~pd.isna(df.Data_up_to)].Data_up_to.max()),
+        "n_suspected": int(status.get("suspected", 0)),
+        "date": str(df[~pd.isna(df[date_col])][date_col].max()),
         "pc_valid_age_gender": percentage_occurrence(
             confirmed,
             (~confirmed.Age.isna()) & (~confirmed.Gender.isna()),
@@ -232,7 +230,7 @@ def plot_timeseries_location_status(
     return fig
 
 
-def plot_epicurve(df: pd.DataFrame, cumulative: bool = True):
+def plot_epicurve(df: pd.DataFrame, non_confirmed_col: str, cumulative: bool = True):
     data = get_epicurve(df, cumulative=cumulative)
     fig = go.Figure()
 
@@ -248,7 +246,7 @@ def plot_epicurve(df: pd.DataFrame, cumulative: bool = True):
     fig.add_trace(
         go.Scatter(
             x=data.Date_onset_estimated,
-            y=data.probable,
+            y=data[non_confirmed_col],
             name="probable",
             line_color=SECONDARY_COLOR,
             line_width=3,