Add printc function and reformat

3 files changed, 76 insertions, 45 deletions

diff --git a/do_ml.py b/do_ml.py
index e545506f3..25cf04e01 100644
--- a/do_ml.py
+++ b/do_ml.py
@@ -21,23 +21,12 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
 import time
-from colorama import init, Fore, Style
+from misc import printc
 import numpy as np
 from gauss_kernel import gauss_kernel
 from cholesky_solve import cholesky_solve
 
 
-init()
-
-
-def printc(text, color):
-    """
-    Prints texts normaly, but in color. Using colorama.
-    text: string with the text to print.
-    """
-    print(color + text + Style.RESET_ALL)
-
-
 def do_ml(desc_data,
           energy_data,
           training_size,
@@ -65,20 +54,19 @@ def do_ml(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)
+                        'than energy data size.']), '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)
+        printc('ERROR. Training or test size greater than data size.', 'RED')
         return None
 
     tic = time.perf_counter()
     if show_msgs:
-        printc('{} ML started, with parameters:'.format(desc_type), Fore.CYAN)
-        printc('\tTraining size: {}'.format(training_size), Fore.BLUE)
-        printc('\tTest size: {}'.format(test_size), Fore.BLUE)
-        printc('\tSigma: {}'.format(sigma), Fore.BLUE)
+        printc('{} ML started, with parameters:'.format(desc_type), 'CYAN')
+        printc('\tTraining size: {}'.format(training_size), 'BLUE')
+        printc('\tTest size: {}'.format(test_size), 'BLUE')
+        printc('\tSigma: {}'.format(sigma), 'BLUE')
 
     Xcm_training = desc_data[:training_size]
     Ycm_training = energy_data[:training_size]
@@ -98,6 +86,6 @@ def do_ml(desc_data,
     tictoc = toc - tic
     if show_msgs:
         printc('\t{} ML took {:.4f} seconds.'.format(desc_type, tictoc),
-               Fore.GREEN)
+               'GREEN')
 
     return mae, tictoc
diff --git a/main.py b/main.py
index 679bb4eea..4354e4063 100644
--- a/main.py
+++ b/main.py
@@ -22,40 +22,32 @@ SOFTWARE.
 """
 import os
 import time
-from colorama import init, Fore, Style
+# from colorama import init, Fore, Style
+from misc import printc
 # import math
 # import random
-import numpy as np
+# import numpy as np
 # from numpy.linalg import cholesky, eig
 # import matplotlib.pyplot as plt
 from read_nc_data import read_nc_data
 from read_db_edata import read_db_edata
-from c_matrix import c_matrix
-from lj_matrix import lj_matrix
+from c_matrix import c_matrix_multiple
+from lj_matrix import lj_matrix_multiple
 # from frob_norm import frob_norm
 # from gauss_kernel import gauss_kernel
 # from cholesky_solve import cholesky_solve
 from do_ml import do_ml
 
 
-def printc(text, color):
-    """
-    Prints texts normaly, but in color. Using colorama.
-    text: string with the text to print.
-    """
-    print(color + text + Style.RESET_ALL)
-
-
-# Initialization.
+# Initialization time.
 init_time = time.perf_counter()
-init()
 
 # Move to data folder.
 init_path = os.getcwd()
 os.chdir('data')
 data_path = os.getcwd()
 
-printc('Data reading started.', Fore.CYAN)
+printc('Data reading started.', 'CYAN')
 tic = time.perf_counter()
 
 # Read nuclear charge data.
@@ -66,30 +58,28 @@ molecules, nuclear_charge, energy_pbe0, energy_delta = \
     read_db_edata(zi_data, data_path)
 
 toc = time.perf_counter()
-printc('\tData reading took {:.4f} seconds.'.format(toc-tic), Fore.GREEN)
+printc('\tData reading took {:.4f} seconds.'.format(toc-tic), 'GREEN')
 
 # Go back to main folder.
 os.chdir(init_path)
 
-printc('Coulomb Matrices calculation started.', Fore.CYAN)
+printc('Coulomb Matrices calculation started.', 'CYAN')
 tic = time.perf_counter()
 
-cm_data = np.array([c_matrix(mol, nc, as_eig=True)
-                    for mol, nc in zip(molecules, nuclear_charge)])
+cm_data = c_matrix_multiple(molecules, nuclear_charge, as_eig=True)
 
 toc = time.perf_counter()
 printc('\tCoulomb matrices calculation took {:.4f} seconds.'.format(toc-tic),
-       Fore.GREEN)
+       'GREEN')
 
-printc('L-J Matrices calculation started.', Fore.CYAN)
+printc('L-J Matrices calculation started.', 'CYAN')
 tic = time.perf_counter()
 
-ljm_data = np.array([lj_matrix(mol, nc, as_eig=True)
-                     for mol, nc in zip(molecules, nuclear_charge)])
+ljm_data = lj_matrix_multiple(molecules, nuclear_charge, as_eig=True)
 
 toc = time.perf_counter()
 printc('\tL-J matrices calculation took {:.4f} seconds.'.format(toc-tic),
-       Fore.GREEN)
+       'GREEN')
 
 # Coulomb Matrix ML
 do_ml(cm_data, energy_pbe0, 1000, 100, sigma=1000.0, desc_type='CM')
@@ -99,4 +89,4 @@ do_ml(ljm_data, energy_pbe0, 1000, 100, sigma=1000.0, desc_type='L-JM')
 # End of program
 end_time = time.perf_counter()
 printc('Program took {:.4f} seconds of runtime.'.format(end_time - init_time),
-       Fore.CYAN)
+       'CYAN')
diff --git a/misc.py b/misc.py
new file mode 100644
index 000000000..c50653a5c
--- /dev/null
+++ b/misc.py
@@ -0,0 +1,53 @@
+"""MIT License
+
+Copyright (c) 2019 David Luevano Alvarado
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from colorama import init, Fore, Style
+
+init()
+
+
+def printc(text, color):
+    """
+    Prints texts normaly, but in color. Using colorama.
+    text: string with the text to print.
+    color: color to be used, same as available in colorama.
+    """
+    color_dic = {'BLACK': Fore.BLACK,
+                 'RED': Fore.RED,
+                 'GREEN': Fore.GREEN,
+                 'YELLOW': Fore.YELLOW,
+                 'BLUE': Fore.BLUE,
+                 'MAGENTA': Fore.MAGENTA,
+                 'CYAN': Fore.CYAN,
+                 'WHITE': Fore.WHITE,
+                 'RESET': Fore.RESET}
+
+    color_dic_keys = color_dic.keys()
+    if color not in color_dic_keys:
+        print(Fore.RED
+              + '\'{}\' not found, using default color.'.format(color)
+              + Style.RESET_ALL)
+        actual_color = Fore.RESET
+    else:
+        actual_color = color_dic[color]
+
+    print(actual_color + text + Style.RESET_ALL)