Merge pull request #228 from cognizant-ai-labs/conditional_lstm_data_processing

Conditional lstm data processing

Author: JamiesonWarner

covid_xprize/examples/predictors/conditional_lstm/conditional_xprize_predictor.py

@@ -44,7 +44,7 @@ def __init__(self, path_to_model_weights, data_url):
                                             nb_lookback_days=NB_LOOKBACK_DAYS)
             self.predictor.load_weights(path_to_model_weights)
 
-        self.df = prepare_cases_dataframe(data_url, threshold_min_cases=True)
+        self.df = prepare_cases_dataframe(data_url)
 
     def train(self,
               return_results=False,

covid_xprize/examples/predictors/conditional_lstm/train_predictor.py

@@ -21,7 +21,7 @@
 from keras.constraints import Constraint
 
 from covid_xprize.examples.predictors.conditional_lstm.conditional_lstm_model import construct_conditional_lstm_model
-from covid_xprize.oxford_data import most_affected_countries, create_country_samples
+from covid_xprize.oxford_data import most_affected_countries, create_country_samples, threshold_min_cases
 
 
 def construct_model(nb_context: int, nb_action: int, lstm_size: int = 32, nb_lookback_days: int = 21) -> Model:
@@ -71,6 +71,9 @@ def train_predictor(training_data: pd.DataFrame,
     df = training_data
     context_column = 'SmoothNewCasesPer100K'
 
+    # Only look at data with # cases above a minimum.
+    df = threshold_min_cases(df)
+
     # Create data set for training
     if nb_training_geos == None: # Use all countries
         nb_training_geos = len(df.GeoID.unique())

covid_xprize/oxford_data/oxford_data.py

@@ -177,7 +177,7 @@ def add_population_column(df: pd.DataFrame) -> pd.DataFrame:
     return df.merge(pop_df[['GeoID', 'Population']], on=['GeoID'], how='left', suffixes=('', '_y'))
 
 
-def prepare_cases_dataframe(data_url: str, threshold_min_cases=False) -> pd.DataFrame:
+def prepare_cases_dataframe(data_url: str) -> pd.DataFrame:
     """
     Loads the cases dataset from the given file, cleans it, and computes cases columns.
     :param data_url: the url containing the original data
@@ -218,10 +218,6 @@ def prepare_cases_dataframe(data_url: str, threshold_min_cases=False) -> pd.Data
     df['DeathRatio'] = df.groupby('GeoID', group_keys=False).SmoothNewDeaths.pct_change(
     ).fillna(0).replace(np.inf, 0) + 1
 
-    # Remove all rows with too few cases
-    if threshold_min_cases:
-        df.drop(df[df.ConfirmedCases < MIN_CASES].index, inplace=True)
-
     # Add column for proportion of population infected
     df['ProportionInfected'] = df['ConfirmedCases'] / df['Population']
 
@@ -234,6 +230,11 @@ def prepare_cases_dataframe(data_url: str, threshold_min_cases=False) -> pd.Data
     return df
 
 
+def threshold_min_cases(df: pd.DataFrame) -> pd.DataFrame:
+    """Remove all rows with too few cases"""
+    return df.drop(df[df.ConfirmedCases < MIN_CASES].index)
+
+
 def create_country_samples(df: pd.DataFrame,
                            countries: list[str],
                            context_column: str,