Add nancodes to cdc_covidnet

Author: dshemetov

cdc_covidnet/delphi_cdc_covidnet/update_sensor.py

@@ -11,7 +11,7 @@
 
 import numpy as np
 import pandas as pd
-from delphi_utils import GeoMapper, add_prefix
+from delphi_utils import GeoMapper, add_prefix, Nans
 
 from .api_config import APIConfig
 from .constants import SIGNALS
@@ -47,6 +47,19 @@ def write_to_csv(data: pd.DataFrame, out_name: str, output_path: str):
         sub_df.drop("epiweek", axis=1).to_csv(filename, na_rep="NA")
 
 
+def add_nancodes(df: pd.DataFrame) -> pd.DataFrame:
+    # Default missing code
+    df["missing_val"] = Nans.NOT_MISSING
+    missing_mask = ~df["val"].isnull()
+    df.loc[missing_mask, "missing_val"] = Nans.UNKNOWN
+
+    # Fill in remaining expected columns
+    df["missing_se"] = Nans.NOT_APPLICABLE
+    df["missing_sample_size"] = Nans.NOT_APPLICABLE
+
+    return df
+
+
 def update_sensor(
         state_files: List[str],
         mmwr_info: pd.DataFrame,
@@ -99,10 +112,11 @@ def update_sensor(
     assert not hosp_df.duplicated(["date", "geo_id"]).any(), "Non-unique (date, geo_id) pairs"
     hosp_df.set_index(["date", "geo_id"], inplace=True)
 
-    # Fill in remaining expected columns
     hosp_df["se"] = np.nan
     hosp_df["sample_size"] = np.nan
 
+    hosp_df = add_nancodes(hosp_df)
+
     # Write results
     signals = add_prefix(SIGNALS,
                          wip_signal=wip_signal,

cdc_covidnet/tests/test_run.py

@@ -5,6 +5,7 @@
 from pandas.testing import assert_frame_equal
 
 from delphi_cdc_covidnet.run import run_module
+from delphi_cdc_covidnet.update_sensor import add_nancodes
 
 
 class TestRun:
@@ -55,5 +56,6 @@ def test_match_old_to_new_output(self):
 
             # Contents match
             expected_df = pd.read_csv(join("receiving_test", fname))
+            expected_df = add_nancodes(expected_df)
             actual_df = pd.read_csv(join("receiving", fname))
             assert_frame_equal(expected_df, actual_df, check_less_precise=5)

cdc_covidnet/tests/test_update_sensor.py

@@ -95,7 +95,16 @@ def test_syn_update_sensor(self):
 
             for i, exp_file in enumerate(expected_files):
                 data = pd.read_csv(exp_file)
-                assert (data.columns == ["geo_id", "val", "se", "sample_size"]).all()
+                expected_columns = [
+                    "geo_id",
+                    "val",
+                    "se",
+                    "sample_size",
+                    "missing_val",
+                    "missing_se",
+                    "missing_sample_size"
+                ]
+                assert (data.columns == expected_columns).all()
 
                 # Check data for NA
                 assert (~pd.isna(data["geo_id"])).all()