Add prediction script

Author: Colin Grambow

predict.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import reacdiff.parsing as parsing
+import reacdiff.train.predict as predict
+
+
+if __name__ == '__main__':
+    args = parsing.parse_predict_args()
+    predict.predict(args)

reacdiff/parsing.py

@@ -11,6 +11,23 @@ def parse_dataprep_args():
     return parser.parse_args()
 
 
+def parse_predict_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument('--data_path', type=str, required=True,
+                        help='Path to data containing states for prediction task')
+    parser.add_argument('--model', type=str, required=True,
+                        help='Path to trained model')
+    parser.add_argument('--data_path2', type=str,
+                        help='Path to additional observable states for prediction')
+    parser.add_argument('--save_path', type=str, default=os.path.join(os.getcwd(), 'preds.csv'),
+                        help='Path to save predictions to')
+    parser.add_argument('--batch_size', type=int, default=32,
+                        help='Batch size')
+    return parser.parse_args()
+
+
 def parse_train_args():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter

reacdiff/train/predict.py

@@ -0,0 +1,24 @@
+import os
+
+import keras
+
+import reacdiff.data.data as datamod
+import reacdiff.utils as utils
+
+
+def predict(args):
+    # Load data
+    print('Loading data')
+    data = datamod.Dataset(
+        datamod.load_data(args.data_path),
+        data2=None if args.data_path2 is None else datamod.load_data(args.data_path2)
+    )
+
+    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+
+    # Load model
+    model = keras.models.load_model(args.model, custom_objects={'rmse': utils.rmse, 'mae': utils.mae})
+
+    # Predict
+    preds = model.predict(data.get_data(), batch_size=args.batch_size, verbose=1)
+    datamod.save_csv(preds, args.save_path)