From e4f9e15588ec796f73c000a683cc9152454a913c Mon Sep 17 00:00:00 2001
From: David Luevano <55825613+luevano@users.noreply.github.com>
Date: Sat, 28 Dec 2019 11:12:36 -0700
Subject: Fix bugs

---
 lj_matrix/__main__.py | 10 +++++++++-
 lj_matrix/do_ml.py    |  7 +++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/lj_matrix/__main__.py b/lj_matrix/__main__.py
index 98f341e1e..811024ff0 100644
--- a/lj_matrix/__main__.py
+++ b/lj_matrix/__main__.py
@@ -24,6 +24,14 @@ from lj_matrix.do_ml import do_ml
 # from lj_matrix.misc import plot_benchmarks
 
 if __name__ == '__main__':
-    do_ml(min_training_size=1500, max_training_size=3000)
+    do_ml(min_training_size=1500,
+          max_training_size=2000,
+          training_increment_size=500,
+          test_size=None,
+          ljm_diag_value=None,
+          ljm_sigma=1.0,
+          ljm_epsilon=1.0,
+          save_benchmarks=False,
+          show_msgs=True)
     # plot_benchmarks()
     print('OK!')
diff --git a/lj_matrix/do_ml.py b/lj_matrix/do_ml.py
index 45dc7a5f0..da9386bf7 100644
--- a/lj_matrix/do_ml.py
+++ b/lj_matrix/do_ml.py
@@ -114,6 +114,7 @@ def ml(desc_data,
 def do_ml(min_training_size,
           max_training_size=None,
           training_increment_size=500,
+          test_size=None,
           ljm_diag_value=None,
           ljm_sigma=1.0,
           ljm_epsilon=1.0,
@@ -128,6 +129,8 @@ def do_ml(min_training_size,
     min_training_size: minimum training size.
     max_training_size: maximum training size.
     training_increment_size: training increment size.
+    test_size: size of the test set to use. If no size is given,
+        the last remaining molecules are used.
     ljm_diag_value: if a special diagonal value should be used in lj matrix.
     ljm_sigma: sigma value for lj matrix.
     ljm_epsilon: epsilon value for lj matrix.
@@ -170,7 +173,7 @@ def do_ml(min_training_size,
                            i,
                            'CM',
                            cm_send,
-                           max_training_size,
+                           test_size,
                            sigma,
                            show_msgs))
         procs.append(p1)
@@ -184,7 +187,7 @@ def do_ml(min_training_size,
                            i,
                            'L-JM',
                            ljm_send,
-                           max_training_size,
+                           test_size,
                            sigma,
                            show_msgs))
         procs.append(p2)
-- 
cgit v1.2.3-54-g00ecf