From cc4eaa121dc340fe44650b4fac40353266b224d8 Mon Sep 17 00:00:00 2001
From: David Luevano <55825613+luevano@users.noreply.github.com>
Date: Thu, 12 Dec 2019 03:38:25 -0700
Subject: Add a bit of error handling

---
 do_ml.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

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)
-- 
cgit v1.2.3-70-g09d2