1 files changed, 124 insertions, 0 deletions
diff --git a/c_matrix.py b/c_matrix.py
new file mode 100644
index 000000000..b47e15aac
--- /dev/null
+++ b/c_matrix.py
@@ -0,0 +1,124 @@
+import math
+import numpy as np
+from numpy.linalg import eig
+
+
+def c_matrix(mol_data,
+             nc_data,
+             max_len=25,
+             as_eig=False,
+             bohr_radius_units=False):
+    """
+    Creates the coulomb matrix from the molecule data given.
+    mol_data: molecule data, matrix of atom coordinates.
+    nc_data: nuclear charge data, array of atom data.
+    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.
+    """
+    if bohr_radius_units:
+        conversion_rate = 0.52917721067
+    else:
+        conversion_rate = 1
+
+    mol_n = len(mol_data)
+    mol_nr = range(mol_n)
+
+    if not mol_n == len(nc_data):
+        print(''.join(['Error. Molecule matrix dimension is different ',
+                       'than the nuclear charge array dimension.']))
+    else:
+        if max_len < mol_n:
+            print(''.join(['Error. Molecule matrix dimension (mol_n) is ',
+                           'greater than max_len. Using mol_n.']))
+            max_len = None
+
+        if max_len:
+            cm = np.zeros((max_len, max_len))
+            ml_r = range(max_len)
+
+            # Actual calculation of the coulomb matrix.
+            for i in ml_r:
+                if i < mol_n:
+                    x_i = mol_data[i, 0]
+                    y_i = mol_data[i, 1]
+                    z_i = mol_data[i, 2]
+                    Z_i = nc_data[i]
+                else:
+                    break
+
+                for j in ml_r:
+                    if j < mol_n:
+                        x_j = mol_data[j, 0]
+                        y_j = mol_data[j, 1]
+                        z_j = mol_data[j, 2]
+                        Z_j = nc_data[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(eig(cm)[0])[::-1]
+                # Thanks to SO for the following lines of code.
+                # https://stackoverflow.com/a/43011036
+
+                # 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 mol_nr:
+                x_i = mol_data[i, 0]
+                y_i = mol_data[i, 1]
+                z_i = mol_data[i, 2]
+                Z_i = nc_data[i]
+
+                cm_row = []
+                for j in mol_nr:
+                    x_j = mol_data[j, 0]
+                    y_j = mol_data[j, 1]
+                    z_j = mol_data[j, 2]
+                    Z_j = nc_data[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(eig(cm)[0])[::-1]
+            else:
+                return cm