Merge pull request #223 from cognizant-ai-labs/fix-tests

JamiesonWarner · web-flow · commit 98abed8cfcd0 · 2023-06-05T11:59:18.000-05:00
reduce the size of the test cases to automated testing runs faster
diff --git a/covid_xprize/examples/predictors/conditional_lstm/conditional_xprize_predictor.py b/covid_xprize/examples/predictors/conditional_lstm/conditional_xprize_predictor.py
@@ -46,16 +46,21 @@ def __init__(self, path_to_model_weights, data_url):
 
         self.df = prepare_cases_dataframe(data_url, threshold_min_cases=True)
 
-    def train(self, return_results=False) -> Union[Model, tuple[Model, dict]]:
+    def train(self,
+              return_results=False,
+              nb_training_geos: int = NB_TRAINING_DAYS,
+              nb_testing_geos: int = NB_TESTING_GEOS,
+              nb_trials: int = NUM_TRIALS,
+              nb_epochs: int = NUM_EPOCHS,) -> Union[Model, tuple[Model, dict]]:
         best_model, results_df = train_predictor(
             training_data=self.df,
             nb_lookback_days=NB_LOOKBACK_DAYS,
             nb_training_days=NB_TRAINING_DAYS,
             nb_test_days=NB_TEST_DAYS,
-            nb_training_geos=NB_TRAINING_GEOS,
-            nb_testing_geos=NB_TESTING_GEOS,
-            nb_trials=NUM_TRIALS,
-            nb_epochs=NUM_EPOCHS,
+            nb_training_geos=nb_training_geos,
+            nb_testing_geos=nb_testing_geos,
+            nb_trials=nb_trials,
+            nb_epochs=nb_epochs,
             lstm_size=LSTM_SIZE,
         )
         if return_results:
diff --git a/covid_xprize/examples/predictors/conditional_lstm/tests/fixtures/trained_model_weights_for_tests.h5 b/covid_xprize/examples/predictors/conditional_lstm/tests/fixtures/trained_model_weights_for_tests.h5
diff --git a/covid_xprize/examples/predictors/conditional_lstm/tests/test_conditional_xprize_predictor.py b/covid_xprize/examples/predictors/conditional_lstm/tests/test_conditional_xprize_predictor.py
@@ -1,6 +1,7 @@
 # Copyright 2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
 
 import os
+from pathlib import Path
 import unittest
 import urllib.request
 
@@ -9,11 +10,11 @@
 from covid_xprize.examples.predictors.conditional_lstm.conditional_xprize_predictor import ConditionalXPrizePredictor
 from covid_xprize.oxford_data import load_oxford_data_trimmed
 
-ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-FIXTURES_PATH = os.path.join(ROOT_DIR, 'fixtures')
-EXAMPLE_INPUT_FILE = os.path.join(ROOT_DIR, "../../../../validation/data/2020-09-30_historical_ip.csv")
-DATA_FILE = os.path.join(FIXTURES_PATH, "OxCGRT_trimmed.csv")
-PREDICTOR_WEIGHTS = os.path.join(FIXTURES_PATH, "trained_model_weights_for_tests.h5")
+ROOT_DIR = Path(__file__).parent
+FIXTURES_PATH = ROOT_DIR / 'fixtures'
+EXAMPLE_INPUT_FILE = (ROOT_DIR / "../../../../validation/data/2020-09-30_historical_ip.csv").absolute()
+DATA_FILE = FIXTURES_PATH / "OxCGRT_trimmed.csv"
+PREDICTOR_WEIGHTS = FIXTURES_PATH / "trained_model_weights_for_tests.h5"
 
 TRAINING_START_DATE = "2020-06-01"
 TRAINING_END_DATE = "2020-07-31"
@@ -25,12 +26,16 @@ class TestConditionalXPrizePredictor(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        df = load_oxford_data_trimmed(end_date=TRAINING_END_DATE, start_date=TRAINING_START_DATE)
-        df.to_csv(DATA_FILE, index=False)
+        FIXTURES_PATH.mkdir(exist_ok=True)
+        if not DATA_FILE.exists():
+            df = load_oxford_data_trimmed(end_date=TRAINING_END_DATE, start_date=TRAINING_START_DATE)
+            df.to_csv(DATA_FILE, index=False)
 
     def test_train_and_predict(self):
         predictor = ConditionalXPrizePredictor(None, DATA_FILE)
-        model = predictor.train()
+
+        # Testing on a small number of epochs, trials and geos to make sure everything is wired correctly
+        model = predictor.train(nb_epochs=2, nb_trials=2, nb_testing_geos=2, nb_training_geos=2)
         model.save_weights(PREDICTOR_WEIGHTS)
         self.assertIsNotNone(model)
 
diff --git a/covid_xprize/examples/predictors/conditional_lstm/train_predictor.py b/covid_xprize/examples/predictors/conditional_lstm/train_predictor.py
@@ -149,24 +149,22 @@ def train_predictor(training_data: pd.DataFrame,
     # Gather test info
     country_indeps = []
     country_predss = []
-    country_casess = []
     for model in models:
-        country_indep, country_preds, country_cases = _lstm_get_test_rollouts(model,
-                                                                              df,
-                                                                              test_countries,
-                                                                              country_samples,
-                                                                              context_column)
+        country_indep, country_preds = _lstm_get_test_rollouts( model,
+                                                                df,
+                                                                test_countries,
+                                                                country_samples,
+                                                                context_column)
         country_indeps.append(country_indep)
         country_predss.append(country_preds)
-        country_casess.append(country_cases)
 
     # Compute daily smooth cases per 100K mae
     test_case_maes = []
     for m in range(len(models)):
         total_loss = 0.
         for c in test_countries:
             true_cases = np.array(df[df.GeoID == c].SmoothNewCasesPer100K)[-nb_test_days:]
-            pred_cases = country_casess[m][c][-nb_test_days:]
+            pred_cases = country_predss[m][c][-nb_test_days:]
             if true_cases.shape != pred_cases.shape: # Insufficient data
                 continue
             total_loss += np.mean(np.abs(true_cases - pred_cases))
@@ -228,7 +226,6 @@ def _lstm_roll_out_predictions(model, initial_context_input, initial_action_inpu
 def _lstm_get_test_rollouts(model, df, top_countries, country_samples, context_column):
     country_indep = {}
     country_preds = {}
-    country_cases = {}
     for c in top_countries:
         X_test_context = country_samples[c]['X_test_context']
         X_test_action = country_samples[c]['X_test_action']
@@ -250,4 +247,4 @@ def _lstm_get_test_rollouts(model, df, top_countries, country_samples, context_c
                                            future_action_sequence)
         country_preds[c] = preds
 
-    return country_indep, country_preds, country_cases
+    return country_indep, country_preds
diff --git a/covid_xprize/examples/predictors/lstm/tests/test_xprize_predictor.py b/covid_xprize/examples/predictors/lstm/tests/test_xprize_predictor.py
@@ -1,21 +1,23 @@
 # Copyright 2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
 
 import os
+from pathlib import Path
 import unittest
 import urllib.request
 
 import pandas as pd
 
 from covid_xprize.examples.predictors.lstm.xprize_predictor import XPrizePredictor
+from covid_xprize.oxford_data import load_oxford_data_trimmed
 
-ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-FIXTURES_PATH = os.path.join(ROOT_DIR, 'fixtures')
-EXAMPLE_INPUT_FILE = os.path.join(ROOT_DIR, "../../../../validation/data/2020-09-30_historical_ip.csv")
-DATA_FILE = os.path.join(FIXTURES_PATH, "OxCGRT_latest.csv")
-DATA_URL =\
-    "https://raw.githubusercontent.com/OxCGRT/covid-policy-tracker-legacy/main/legacy_data_202207/OxCGRT_latest.csv"
-PREDICTOR_WEIGHTS = os.path.join(FIXTURES_PATH, "trained_model_weights_for_tests.h5")
+ROOT_DIR = Path(__file__).parent
+FIXTURES_PATH = ROOT_DIR / 'fixtures'
+EXAMPLE_INPUT_FILE = (ROOT_DIR / "../../../../validation/data/2020-09-30_historical_ip.csv").absolute()
+DATA_FILE = FIXTURES_PATH / "OxCGRT_trimmed.csv"
+PREDICTOR_WEIGHTS = FIXTURES_PATH / "trained_model_weights_for_tests.h5"
 
+TRAINING_START_DATE = "2020-06-01"
+TRAINING_END_DATE = "2020-07-31"
 START_DATE = "2020-08-01"
 END_DATE = "2020-08-04"
 
@@ -24,15 +26,19 @@ class TestXPrizePredictor(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        # Download and cache the raw data
-        urllib.request.urlretrieve(DATA_URL, DATA_FILE)
+        FIXTURES_PATH.mkdir(exist_ok=True)
+        if not DATA_FILE.exists():
+            df = load_oxford_data_trimmed(end_date=TRAINING_END_DATE, start_date=TRAINING_START_DATE)
+            df.to_csv(DATA_FILE, index=False)
+
+    def test_train_and_predict(self):
+        predictor = XPrizePredictor(None, DATA_FILE)
+
+        # Testing on a small number of epochs and trials to make sure everything is wired correctly
+        model = predictor.train(num_trials=2, num_epochs=2)
+        model.save_weights(PREDICTOR_WEIGHTS)
+        self.assertIsNotNone(model)
 
-    def test_predict(self):
         predictor = XPrizePredictor(PREDICTOR_WEIGHTS, DATA_FILE)
         pred_df = predictor.predict(START_DATE, END_DATE, EXAMPLE_INPUT_FILE)
         self.assertIsInstance(pred_df, pd.DataFrame)
-
-    def test_train(self):
-        predictor = XPrizePredictor(None, DATA_FILE)
-        model = predictor.train()
-        self.assertIsNotNone(model)
diff --git a/covid_xprize/examples/predictors/lstm/xprize_predictor.py b/covid_xprize/examples/predictors/lstm/xprize_predictor.py
@@ -27,6 +27,7 @@
 WINDOW_SIZE = 7
 US_PREFIX = "United States / "
 NUM_TRIALS = 1
+NUM_EPOCHS = 1000
 LSTM_SIZE = 32
 MAX_NB_COUNTRIES = 20
 
@@ -213,7 +214,7 @@ def _convert_ratios_to_total_cases(self,
     def _smooth_case_list(case_list, window):
         return pd.Series(case_list).rolling(window).mean().to_numpy()
 
-    def train(self):
+    def train(self, num_trials=NUM_TRIALS, num_epochs=NUM_EPOCHS):
         print("Creating numpy arrays for Keras for each country...")
         geos = self._most_affected_geos(self.df, MAX_NB_COUNTRIES, NB_LOOKBACK_DAYS)
         country_samples = create_country_samples(self.df, geos, CONTEXT_COLUMN, NB_TEST_DAYS, NB_LOOKBACK_DAYS)
@@ -256,14 +257,14 @@ def train(self):
         train_losses = []
         val_losses = []
         test_losses = []
-        for t in range(NUM_TRIALS):
+        for t in range(num_trials):
             print('Trial', t)
             X_context, X_action, y = self._permute_data(X_context, X_action, y, seed=t)
             model, training_model = self._construct_model(nb_context=X_context.shape[-1],
                                                           nb_action=X_action.shape[-1],
                                                           lstm_size=LSTM_SIZE,
                                                           nb_lookback_days=NB_LOOKBACK_DAYS)
-            history = self._train_model(training_model, X_context, X_action, y, epochs=1000, verbose=0)
+            history = self._train_model(training_model, X_context, X_action, y, epochs=num_epochs, verbose=0)
             top_epoch = np.argmin(history.history['val_loss'])
             train_loss = history.history['loss'][top_epoch]
             val_loss = history.history['val_loss'][top_epoch]
diff --git a/covid_xprize/oxford_data/oxford_data.py b/covid_xprize/oxford_data/oxford_data.py
@@ -320,6 +320,8 @@ def create_prediction_initial_context_and_action_vectors(
     df = df[df.Date <= start_date]
     country_samples = create_country_samples(df, countries, context_column, nb_lookback_days=nb_lookback_days)
     for c in countries:
+        if c not in country_samples.keys():
+            continue
         context_vectors[c] = country_samples[c]['X_test_context'][-1]
         action_vectors[c] =  country_samples[c]['X_test_action'][-1]
     return context_vectors, action_vectors
