Add epsilon index representation

2 files changed, 47 insertions, 2 deletions

diff --git a/ml_exp/compound.py b/ml_exp/compound.py
index 6c8525195..eae280d38 100644
--- a/ml_exp/compound.py
+++ b/ml_exp/compound.py
@@ -23,7 +23,7 @@ SOFTWARE.
 import numpy as np
 from ml_exp.data import NUCLEAR_CHARGE
 from ml_exp.representations import coulomb_matrix, lennard_jones_matrix,\
-    get_helping_data, adjacency_matrix, bag_of_bonds
+    get_helping_data, adjacency_matrix, epsilon_index, bag_of_bonds
 
 
 class Compound:
@@ -47,6 +47,7 @@ class Compound:
         self.cm = None
         self.ljm = None
         self.am = None
+        self.ei = None
         self.bob = None
         self.bo_atoms = None
         self.bok_cx = None
@@ -133,10 +134,20 @@ class Compound:
                                    use_forces=use_forces,
                                    size=size)
 
+    def gen_ei(self,
+               size=23):
+        """
+        Generates the Epsilon Index for the compound.
+        size: compound size.
+        """
+        self.ei = epsilon_index(self.am,
+                                bonds_i,
+                                size=size)
+
     def gen_bob(self,
                 size=23):
         """
-        Generate the Bag of Stuff for the compound.
+        Generate the Bag of Bonds for the compound.
         size: compound size.
         """
         self.bob = bag_of_bonds(self.cm,
diff --git a/ml_exp/representations.py b/ml_exp/representations.py
index 4a685d388..ff5741816 100644
--- a/ml_exp/representations.py
+++ b/ml_exp/representations.py
@@ -254,6 +254,40 @@ the current compound.')
     return np.pad(am, ((0, size - n), (0, size - n)), 'constant')
 
 
+def epsilon_index(am,
+                  bonds_i,
+                  size=23):
+    """
+    Calculates the Epsilon index of G, presented by Estrada.
+    am: adjacency matrix.
+    bonds_i: list of bond indexes (tuple of indexes).
+    size: compund size.
+    """
+    if am is None:
+        raise ValueError('The adjacency matrix hasn\'t been initialized for\
+the current compound.')
+
+    n = len(bonds_i)
+    if size < n:
+        print('Error. Compound size (n) is greater han (size). Using (n)\
+              instead of (size).')
+        size = n
+
+    deltas = np.zeros(n, dtype=np.float64)
+
+    for i in range(n):
+        deltas[i] = np.sum(am[i, :])
+
+    ei = 0.0
+    for i in range(n - 1):
+        for j in range(i +1, n):
+            if (bonds_i[i][0] in bonds_i[j]) or (bonds_i[i][0] in bonds_i[j]):
+                val = deltas[i]*deltas[j]
+                ei += 1.0/val**0.5
+
+    return ei
+
+
 def check_bond(bags,
                bond):
     """