Add a bit of error handling

1 files changed, 12 insertions, 0 deletions

diff --git a/do_ml.py b/do_ml.py
index 561b33904..e0df62be3 100644
--- a/do_ml.py
+++ b/do_ml.py
@@ -58,6 +58,18 @@ def do_ml(desc_data,
     if not desc_type:
         desc_type = 'NOT SPECIFIED'
 
+    d_len = len(desc_data)
+    e_len = len(energy_data)
+    if d_len != e_len:
+        printc(''.join(['ERROR. Descriptor data size different ',
+                        'than energy data size.']), Fore.RED)
+        return None
+
+    if training_size > d_len or test_size > d_len:
+        printc('ERROR. Training or test size greater than data size.',
+               Fore.RED)
+        return None
+
     tic = time.perf_counter()
     if show_msgs:
         printc('{} ML started, with parameters:'.format(desc_type), Fore.CYAN)