Refactor code and fix bug

3 files changed, 39 insertions, 8 deletions

diff --git a/lj_matrix/do_ml.py b/lj_matrix/do_ml.py
index 8724e6831..45dc7a5f0 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,
+          ljm_diag_value=None,
           ljm_sigma=1.0,
           ljm_epsilon=1.0,
           save_benchmarks=False,
@@ -127,6 +128,7 @@ def do_ml(min_training_size,
     min_training_size: minimum training size.
     max_training_size: maximum training size.
     training_increment_size: training increment size.
+    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.
     save_benchmarks: if benchmarks should be saved.
@@ -147,6 +149,7 @@ def do_ml(min_training_size,
     # Matrices calculation.
     cm_data, ljm_data = parallel_create_matrices(molecules,
                                                  nuclear_charge,
+                                                 ljm_diag_value,
                                                  ljm_sigma,
                                                  ljm_epsilon,
                                                  max_len,
diff --git a/lj_matrix/lj_matrix.py b/lj_matrix/lj_matrix.py
index 0c16b5686..c3b61becb 100644
--- a/lj_matrix/lj_matrix.py
+++ b/lj_matrix/lj_matrix.py
@@ -29,6 +29,7 @@ from lj_matrix.misc import printc
 
 def lj_matrix(mol_data,
               nc_data,
+              diag_value=None,
               sigma=1.0,
               epsilon=1.0,
               max_len=25,
@@ -38,6 +39,7 @@ def lj_matrix(mol_data,
     Creates the Lennard-Jones Matrix from the molecule data given.
     mol_data: molecule data, matrix of atom coordinates.
     nc_data: nuclear charge data, array of atom data.
+    diag_value: if special diagonal value is to be used.
     sigma: sigma value.
     epsilon: epsilon value.
     max_len: maximum amount of atoms in molecule.
@@ -86,7 +88,10 @@ def lj_matrix(mol_data,
                         z = (z_i-z_j)**2
 
                         if i == j:
-                            lj[i, j] = (0.5*Z_i**2.4)
+                            if not diag_value:
+                                lj[i, j] = (0.5*Z_i**2.4)
+                            else:
+                                lj[i, j] = diag_value
                         else:
                             # Calculations are done after i==j is checked
                             # so no division by zero is done.
@@ -144,7 +149,10 @@ def lj_matrix(mol_data,
                     z = (z_i-z_j)**2
 
                     if i == j:
-                        lj_row.append(0.5*Z_i**2.4)
+                        if not diag_value:
+                            lj_row.append(0.5*Z_i**2.4)
+                        else:
+                            lj_row.append(diag_value)
                     else:
                         # Calculations are done after i==j is checked
                         # so no division by zero is done.
@@ -173,6 +181,7 @@ def lj_matrix(mol_data,
 def lj_matrix_multiple(mol_data,
                        nc_data,
                        pipe=None,
+                       diag_value=None,
                        sigma=1.0,
                        epsilon=1.0,
                        max_len=25,
@@ -184,6 +193,9 @@ def lj_matrix_multiple(mol_data,
     nc_data: nuclear charge data, array of atom data.
     pipe: for multiprocessing purposes. Sends the data calculated
         through a pipe.
+    diag_value: if special diagonal value is to be used.
+    sigma: sigma value.
+    epsilon: epsilon value.
     max_len: maximum amount of atoms in molecule.
     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.
@@ -193,6 +205,7 @@ def lj_matrix_multiple(mol_data,
 
     ljm_data = np.array([lj_matrix(mol,
                                    nc,
+                                   diag_value,
                                    sigma,
                                    epsilon,
                                    max_len,
diff --git a/lj_matrix/parallel_create_matrices.py b/lj_matrix/parallel_create_matrices.py
index 0ab691525..cd5ef5c8e 100644
--- a/lj_matrix/parallel_create_matrices.py
+++ b/lj_matrix/parallel_create_matrices.py
@@ -27,8 +27,9 @@ from lj_matrix.lj_matrix import lj_matrix_multiple
 
 def parallel_create_matrices(mol_data,
                              nc_data,
-                             sigma=1.0,
-                             epsilon=1.0,
+                             ljm_diag_value=None,
+                             ljm_sigma=1.0,
+                             ljm_epsilon=1.0,
                              max_len=25,
                              as_eig=True,
                              bohr_radius_units=False):
@@ -36,8 +37,9 @@ def parallel_create_matrices(mol_data,
     Creates the Coulomb and L-J matrices in parallel.
     mol_data: molecule data, matrix of atom coordinates.
     nc_data: nuclear charge data, array of atom data.
-    sigma: sigma value for L-J matrix.
-    epsilon: epsilon value for L-J matrix.
+    ljm_diag_value: if special diagonal value is to be used for lj matrix.
+    ljm_sigma: sigma value for lj matrix.
+    ljm_epsilon: psilon value for lj matrix.
     max_len: maximum amount of atoms in molecule.
     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.
@@ -49,14 +51,27 @@ def parallel_create_matrices(mol_data,
 
     cm_recv, cm_send = Pipe(False)
     p1 = Process(target=c_matrix_multiple,
-                 args=(mol_data, nc_data, cm_send))
+                 args=(mol_data,
+                       nc_data,
+                       cm_send,
+                       max_len,
+                       as_eig,
+                       bohr_radius_units))
     procs.append(p1)
     pipes.append(cm_recv)
     p1.start()
 
     ljm_recv, ljm_send = Pipe(False)
     p2 = Process(target=lj_matrix_multiple,
-                 args=(mol_data, nc_data, ljm_send, sigma, epsilon))
+                 args=(mol_data,
+                       nc_data,
+                       ljm_send,
+                       ljm_diag_value,
+                       ljm_sigma,
+                       ljm_epsilon,
+                       max_len,
+                       as_eig,
+                       bohr_radius_units))
     procs.append(p2)
     pipes.append(ljm_recv)
     p2.start()