Reformat matrices calculation

3 files changed, 18 insertions, 17 deletions

diff --git a/c_matrix.py b/c_matrix.py
index f6f18976e..2bc4d4c0c 100644
--- a/c_matrix.py
+++ b/c_matrix.py
@@ -20,6 +20,8 @@ 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.
 """
+import time
+from misc import printc
 import math
 import numpy as np
 from numpy.linalg import eig
@@ -159,7 +161,13 @@ def c_matrix_multiple(mol_data,
     as_eig: if data should be returned as matrix or array of eigenvalues.
     bohr_radius_units: if units should be in bohr's radius units.
     """
+    printc('Coulomb Matrices calculation started.', 'CYAN')
+    tic = time.perf_counter()
+
     cm_data = np.array([c_matrix(mol, nc, max_len, as_eig, bohr_radius_units)
                        for mol, nc in zip(mol_data, nc_data)])
 
+    toc = time.perf_counter()
+    printc('\tCM calculation took {:.4f} seconds.'.format(toc - tic), 'GREEN')
+
     return cm_data
diff --git a/lj_matrix.py b/lj_matrix.py
index 353631633..6769bc0c3 100644
--- a/lj_matrix.py
+++ b/lj_matrix.py
@@ -20,6 +20,8 @@ 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.
 """
+import time
+from misc import printc
 import math
 import numpy as np
 from numpy.linalg import eig
@@ -177,7 +179,13 @@ def lj_matrix_multiple(mol_data,
     as_eig: if data should be returned as matrix or array of eigenvalues.
     bohr_radius_units: if units should be in bohr's radius units.
     """
+    printc('L-J Matrices calculation started.', 'CYAN')
+    tic = time.perf_counter()
+
     ljm_data = np.array([lj_matrix(mol, nc, max_len, as_eig, bohr_radius_units)
                         for mol, nc in zip(mol_data, nc_data)])
 
+    toc = time.perf_counter()
+    printc('\tL-JM calculation took {:.4f} seconds.'.format(toc-tic), 'GREEN')
+
     return ljm_data
diff --git a/main.py b/main.py
index 4354e4063..31a12e7e8 100644
--- a/main.py
+++ b/main.py
@@ -63,27 +63,12 @@ printc('\tData reading took {:.4f} seconds.'.format(toc-tic), 'GREEN')
 # Go back to main folder.
 os.chdir(init_path)
 
-printc('Coulomb Matrices calculation started.', 'CYAN')
-tic = time.perf_counter()
-
+# Matrices calculation.
 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),
-       'GREEN')
-
-printc('L-J Matrices calculation started.', 'CYAN')
-tic = time.perf_counter()
-
 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),
-       'GREEN')
-
-# Coulomb Matrix ML
+# ML calculation.
 do_ml(cm_data, energy_pbe0, 1000, 100, sigma=1000.0, desc_type='CM')
-# Lennard-Jones Matrix ML
 do_ml(ljm_data, energy_pbe0, 1000, 100, sigma=1000.0, desc_type='L-JM')
 
 # End of program