Add flattening options

2 files changed, 13 insertions, 2 deletions

diff --git a/ml_exp/compound.py b/ml_exp/compound.py
index 16b878f7b..fac9015a2 100644
--- a/ml_exp/compound.py
+++ b/ml_exp/compound.py
@@ -78,11 +78,13 @@ class Compound:
 
     def gen_cm(self,
                size=23,
+               flatten=True,
                as_eig=True,
                bohr_ru=False):
         """
         Generate the Coulomb Matrix for the compund.
         size: compound size.
+        flatten: if the representation should be 1D.
         as_eig: if the representation should be as the eigenvalues.
         bohr_ru: if radius units should be in bohr's radius units.
         """
@@ -97,6 +99,7 @@ class Compound:
                 sigma=1.0,
                 epsilon=1.0,
                 size=23,
+                flatten=True,
                 as_eig=True,
                 bohr_ru=False):
         """
@@ -105,6 +108,7 @@ class Compound:
         sigma: sigma value.
         epsilon: epsilon value.
         size: compound size.
+        flatten: if the representation should be 1D.
         as_eig: if the representation should be as the eigenvalues.
         bohr_ru: if radius units should be in bohr's radius units.
         """
@@ -135,11 +139,13 @@ class Compound:
 
     def gen_am(self,
                use_forces=False,
+               flatten=True,
                size=23):
         """
         Generate the Adjacency Matrix for the compund.
         use_forces: if the use of forces instead of k_cx should be used.
         size: compound size.
+        flatten: if the representation should be 1D.
         """
         self.am = adjacency_matrix(self.bonds_i,
                                    self.bonds_k,
diff --git a/ml_exp/representations.py b/ml_exp/representations.py
index b0a5d8553..fcaa0f33a 100644
--- a/ml_exp/representations.py
+++ b/ml_exp/representations.py
@@ -215,7 +215,8 @@ def adjacency_matrix(bonds_i,
                      bonds_k,
                      bonds_f,
                      use_forces=False,
-                     size=23):
+                     size=23,
+                     flatten=True):
     """
     Calculates the adjacency matrix given the bond list.
     bonds: list of bond names.
@@ -224,6 +225,7 @@ def adjacency_matrix(bonds_i,
     bonds_f: list of force values.
     use_forces: if the use of forces instead of k_cx should be used.
     size: compund size.
+    flatten: if the representation should be 1D.
     """
     if bonds_i is None:
         raise ValueError('The helping data hasn\'t been initialized for\
@@ -247,7 +249,10 @@ the current compound.')
                     am[i, j] = bonds_k[i]
                     am[j, i] = am[i, j]
 
-    return np.pad(am, ((0, size - n), (0, size - n)), 'constant')
+    if flatten:
+        return np.pad(am, ((0, size - n), (0, size - n)), 'constant').flatten()
+    else:
+        return np.pad(am, ((0, size - n), (0, size - n)), 'constant')
 
 
 def epsilon_index(am,