refactored reading into seperate function

Author: aysim319

doctor_visits/delphi_doctor_visits/config.py

@@ -26,6 +26,9 @@ class Config:
     HRR_COLS = ["Pat HRR Name", "Pat HRR ID"]
     ID_COLS = [DATE_COL] + [GEO_COL] + HRR_COLS + [AGE_COL]
     FILT_COLS = ID_COLS + COUNT_COLS
+    # as of 2020-05-11, input file expected to have 10 columns
+    # id cols: ServiceDate, PatCountyFIPS, PatAgeGroup, Pat HRR ID/Pat HRR Name
+    # value cols: Denominator, Covid_like, Flu_like, Flu1, Mixed
     DTYPES = {
         "ServiceDate": str,
         "PatCountyFIPS": str,

doctor_visits/delphi_doctor_visits/update_sensor.py

@@ -9,10 +9,8 @@
 """
 
 # standard packages
-from datetime import timedelta
+from datetime import timedelta, datetime
 from multiprocessing import Pool, cpu_count
-from pathlib import Path
-
 # third party
 import dask.dataframe as dd
 import numpy as np
@@ -68,29 +66,23 @@ def write_to_csv(output_df: pd.DataFrame, geo_level, se, out_name, logger, outpu
                 out_n += 1
     logger.debug(f"wrote {out_n} rows for {geo_level}")
 
-
-def update_sensor(
-        filepath, startdate, enddate, dropdate, geo, parallel,
-        weekday, se, logger
-):
-    """Generate sensor values.
-
-    Args:
+#TODO clean the date params
+def process_csv(filepath: str, startdate: datetime, enddate: datetime, dropdate: datetime) -> pd.DataFrame:
+    '''
+    Reads csv using Dask and filters out based on date range and currently unused column,
+    then converts back into pandas dataframe.
+    Parameters
+    ----------
       filepath: path to the aggregated doctor-visits data
       startdate: first sensor date (YYYY-mm-dd)
       enddate: last sensor date (YYYY-mm-dd)
       dropdate: data drop date (YYYY-mm-dd)
-      geo: geographic resolution, one of ["county", "state", "msa", "hrr", "nation", "hhs"]
-      parallel: boolean to run the sensor update in parallel
-      weekday: boolean to adjust for weekday effects
-      se: boolean to write out standard errors, if true, use an obfuscated name
-      logger: the structured logger
-    """
-    # as of 2020-05-11, input file expected to have 10 columns
-    # id cols: ServiceDate, PatCountyFIPS, PatAgeGroup, Pat HRR ID/Pat HRR Name
-    # value cols: Denominator, Covid_like, Flu_like, Flu1, Mixed
-    filename = Path(filepath).name
-    data = pd.read_csv(
+
+    Returns
+    -------
+    cleaned dataframe
+    '''
+
     ddata = dd.read_csv(
         filepath,
         compression="gzip",
@@ -101,39 +93,46 @@ def update_sensor(
     ddata = ddata.dropna()
     ddata = ddata.rename(columns=Config.DEVIANT_COLS_MAP)
     ddata = ddata[Config.FILT_COLS]
+    ddata[Config.DATE_COL] = dd.to_datetime(ddata[Config.DATE_COL])
 
+    # restrict to training start and end date
+    startdate = startdate - Config.DAY_SHIFT
 
-    data = ddata.compute()
+    assert startdate > Config.FIRST_DATA_DATE, "Start date <= first day of data"
+    assert startdate < enddate, "Start date >= end date"
+    assert enddate <= dropdate, "End date > drop date"
 
-    # data.dropna(inplace=True)  # drop rows with any missing entries
+    date_filter = ((ddata[Config.DATE_COL] >= Config.FIRST_DATA_DATE) & (ddata[Config.DATE_COL] < dropdate))
 
-    # data.columns = data.columns.to_series().replace(Config.DEVIANT_COLS_MAP)
-    #
-    # data = data[Config.FILT_COLS]
-    #
-    # # drop HRR columns - unused for now since we assign HRRs by FIPS
-    # data.drop(columns=Config.HRR_COLS, inplace=True)
-    # data.dropna(inplace=True)  # drop rows with any missing entries
+    return ddata[date_filter].compute()
+
+def update_sensor(
+        filepath:str, startdate:datetime, enddate:datetime, dropdate:datetime, geo:str, parallel: bool,
+        weekday:bool, se:bool, logger
+):
+    """Generate sensor values.
+
+    Args:
+      filepath: path to the aggregated doctor-visits data
+      startdate: first sensor date (YYYY-mm-dd)
+      enddate: last sensor date (YYYY-mm-dd)
+      dropdate: data drop date (YYYY-mm-dd)
+      geo: geographic resolution, one of ["county", "state", "msa", "hrr", "nation", "hhs"]
+      parallel: boolean to run the sensor update in parallel
+      weekday: boolean to adjust for weekday effects
+      se: boolean to write out standard errors, if true, use an obfuscated name
+      logger: the structured logger
+    """
+    data = process_csv(filepath, startdate, enddate, dropdate)
 
-    data[Config.DATE_COL] = data[Config.DATE_COL].apply(pd.to_datetime)
     # aggregate age groups (so data is unique by service date and FIPS)
     data = data.groupby([Config.DATE_COL, Config.GEO_COL]).sum(numeric_only=True).reset_index()
     assert np.sum(data.duplicated()) == 0, "Duplicates after age group aggregation"
     assert (data[Config.COUNT_COLS] >= 0).all().all(), "Counts must be nonnegative"
 
-    ## collect dates
-    # restrict to training start and end date
     drange = lambda s, e: np.array([s + timedelta(days=x) for x in range((e - s).days)])
-    startdate = pd.to_datetime(startdate) - Config.DAY_SHIFT
-    burnindate = startdate - Config.DAY_SHIFT
-    enddate = pd.to_datetime(enddate)
-    dropdate = pd.to_datetime(dropdate)
-    assert startdate > Config.FIRST_DATA_DATE, "Start date <= first day of data"
-    assert startdate < enddate, "Start date >= end date"
-    assert enddate <= dropdate, "End date > drop date"
-    data = data[(data[Config.DATE_COL] >= Config.FIRST_DATA_DATE) & \
-                (data[Config.DATE_COL] < dropdate)]
     fit_dates = drange(Config.FIRST_DATA_DATE, dropdate)
+    burnindate = startdate - Config.DAY_SHIFT
     burn_in_dates = drange(burnindate, dropdate)
     sensor_dates = drange(startdate, enddate)
     # The ordering of sensor dates corresponds to the order of burn-in dates