Add exception

1 files changed, 79 insertions, 79 deletions

diff --git a/ml_exp/representations.py b/ml_exp/representations.py
index 830e51707..50b78e548 100644
--- a/ml_exp/representations.py
+++ b/ml_exp/representations.py
@@ -42,83 +42,34 @@ def coulomb_matrix(coords,
     else:
         conversion_rate = 1
 
-    mol_n = len(coords)
-    mol_nr = range(mol_n)
+    n = coords.shape[0]
+    nr = range(n)
 
-    if not mol_n == len(nc):
-        print(''.join(['Error. Molecule matrix dimension is different ',
-                       'than the nuclear charge array dimension.']))
-    else:
-        if size < mol_n:
-            print(''.join(['Error. Molecule matrix dimension (mol_n) is ',
-                           'greater than size. Using mol_n.']))
-            size = None
-
-        if size:
-            cm = np.zeros((size, size))
-            ml_r = range(size)
-
-            # Actual calculation of the coulomb matrix.
-            for i in ml_r:
-                if i < mol_n:
-                    x_i = coords[i, 0]
-                    y_i = coords[i, 1]
-                    z_i = coords[i, 2]
-                    Z_i = nc[i]
-                else:
-                    break
-
-                for j in ml_r:
-                    if j < mol_n:
-                        x_j = coords[j, 0]
-                        y_j = coords[j, 1]
-                        z_j = coords[j, 2]
-                        Z_j = nc[j]
-
-                        x = (x_i-x_j)**2
-                        y = (y_i-y_j)**2
-                        z = (z_i-z_j)**2
-
-                        if i == j:
-                            cm[i, j] = (0.5*Z_i**2.4)
-                        else:
-                            cm[i, j] = (conversion_rate*Z_i*Z_j/math.sqrt(x
-                                                                          + y
-                                                                          + z))
-                    else:
-                        break
-
-            # Now the value will be returned.
-            if as_eig:
-                cm_sorted = np.sort(np.linalg.eig(cm)[0])[::-1]
-                # Thanks to SO for the following lines of code.
-                # https://stackoverflow.com/a/43011036
+    if not n == nc.shape[0]:
+        raise ValueError('Compound size is different than the nuclear charge\
+                         size. Arrays are not of the right shape.')
 
-                # Keep zeros at the end.
-                mask = cm_sorted != 0.
-                f_mask = mask.sum(0, keepdims=1) >\
-                    np.arange(cm_sorted.shape[0]-1, -1, -1)
+    if size < n:
+        print(''.join(['Error. Molecule matrix dimension (n) is ',
+                       'greater than size. Using n.']))
+        size = None
 
-                f_mask = f_mask[::-1]
-                cm_sorted[f_mask] = cm_sorted[mask]
-                cm_sorted[~f_mask] = 0.
+    if size:
+        cm = np.zeros((size, size))
+        ml_r = range(size)
 
-                return cm_sorted
-
-            else:
-                return cm
-
-        else:
-            cm_temp = []
-            # Actual calculation of the coulomb matrix.
-            for i in mol_nr:
+        # Actual calculation of the coulomb matrix.
+        for i in ml_r:
+            if i < n:
                 x_i = coords[i, 0]
                 y_i = coords[i, 1]
                 z_i = coords[i, 2]
                 Z_i = nc[i]
+            else:
+                break
 
-                cm_row = []
-                for j in mol_nr:
+            for j in ml_r:
+                if j < n:
                     x_j = coords[j, 0]
                     y_j = coords[j, 1]
                     z_j = coords[j, 2]
@@ -129,17 +80,66 @@ def coulomb_matrix(coords,
                     z = (z_i-z_j)**2
 
                     if i == j:
-                        cm_row.append(0.5*Z_i**2.4)
+                        cm[i, j] = (0.5*Z_i**2.4)
                     else:
-                        cm_row.append(conversion_rate*Z_i*Z_j/math.sqrt(x
-                                                                        + y
-                                                                        + z))
+                        cm[i, j] = (conversion_rate*Z_i*Z_j/math.sqrt(x
+                                                                      + y
+                                                                      + z))
+                else:
+                    break
 
-                cm_temp.append(np.array(cm_row))
+        # Now the value will be returned.
+        if as_eig:
+            cm_sorted = np.sort(np.linalg.eig(cm)[0])[::-1]
+            # Thanks to SO for the following lines of code.
+            # https://stackoverflow.com/a/43011036
 
-            cm = np.array(cm_temp)
-            # Now the value will be returned.
-            if as_eig:
-                return np.sort(np.linalg.eig(cm)[0])[::-1]
-            else:
-                return cm
+            # Keep zeros at the end.
+            mask = cm_sorted != 0.
+            f_mask = mask.sum(0, keepdims=1) >\
+                np.arange(cm_sorted.shape[0]-1, -1, -1)
+
+            f_mask = f_mask[::-1]
+            cm_sorted[f_mask] = cm_sorted[mask]
+            cm_sorted[~f_mask] = 0.
+
+            return cm_sorted
+
+        else:
+            return cm
+
+    else:
+        cm_temp = []
+        # Actual calculation of the coulomb matrix.
+        for i in nr:
+            x_i = coords[i, 0]
+            y_i = coords[i, 1]
+            z_i = coords[i, 2]
+            Z_i = nc[i]
+
+            cm_row = []
+            for j in nr:
+                x_j = coords[j, 0]
+                y_j = coords[j, 1]
+                z_j = coords[j, 2]
+                Z_j = nc[j]
+
+                x = (x_i-x_j)**2
+                y = (y_i-y_j)**2
+                z = (z_i-z_j)**2
+
+                if i == j:
+                    cm_row.append(0.5*Z_i**2.4)
+                else:
+                    cm_row.append(conversion_rate*Z_i*Z_j/math.sqrt(x
+                                                                    + y
+                                                                    + z))
+
+            cm_temp.append(np.array(cm_row))
+
+        cm = np.array(cm_temp)
+        # Now the value will be returned.
+        if as_eig:
+            return np.sort(np.linalg.eig(cm)[0])[::-1]
+        else:
+            return cm