Add printc function and reformat

1 files changed, 15 insertions, 25 deletions

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')